pytorch-cpu-only.ipynb

{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "# PyTorch Inference Using only the CPU",
   "id": "10378db052603150"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Setup\n",
    "\n",
    "The requirements for this notebook are very easy, just activate your Python virtual environment of choice (Python 3.9+) and execute the two pip install commands shown below to install PyTorch, Transformers (with PyTorch support instead of TensorFlow), and the Jupyter packages.\n",
    "\n",
    "```console\n",
    "pip install torch --index-url https://download.pytorch.org/whl/cpu\n",
    "pip install jupyter 'transformers[torch]'\n",
    "```"
   ],
   "id": "99bcb4ec989f81b4"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Load model and tokenizer",
   "id": "d2ac3cd0af1b89cc"
  },
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-09-19T07:28:52.949662Z",
     "start_time": "2024-09-19T07:28:49.679220Z"
    }
   },
   "source": [
    "# Load model directly\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Meta-Llama-3.1-8B-Instruct\")\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    pretrained_model_name_or_path=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n",
    "    device_map=torch.device(\"cpu\"),\n",
    "    torch_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "# Initial time taken: 30 minutes (downloading from Internet)\n",
    "# Time taken (after downloading): ~3.25 seconds"
   ],
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "028b3e54a3584380a26e8ceb9b2ea3f5"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Generate tokenized input",
   "id": "8b3472668a923b8d"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-19T06:32:33.704164Z",
     "start_time": "2024-09-19T06:32:33.665707Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model_input = tokenizer([\"Come up with a 20 letter English word.\"], return_tensors=\"pt\")\n",
    "\n",
    "# Time taken: < 10 ms"
   ],
   "id": "4284624fe5d2e4d0",
   "outputs": [],
   "execution_count": 20
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Generate tokenized output",
   "id": "663eac768199605d"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-19T06:51:09.586336Z",
     "start_time": "2024-09-19T06:37:51.759737Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# Manually setting pad_token_id so warning doesn't pop up. Setting it to the default for open generation\n",
    "# Leaving out the option will output a message saying it's setting the pad_token_id to the eos_token_id: 128001 for open\n",
    "# ended generation, setting it manually prevents this message from being output.\n",
    "model_tokenized_output = model.generate(\n",
    "    **model_input,\n",
    "    pad_token_id=128001,\n",
    "    max_new_tokens=100\n",
    ")\n",
    "\n",
    "# Time taken: 13 minutes and 17 seconds"
   ],
   "id": "bd25b20ed190a5cd",
   "outputs": [],
   "execution_count": 26
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "## Decode the tokenized output and print",
   "id": "e931cbedab0ed3fe"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-19T06:51:27.240749Z",
     "start_time": "2024-09-19T06:51:27.235475Z"
    }
   },
   "cell_type": "code",
   "source": [
    "outputs = tokenizer.batch_decode(model_tokenized_output, skip_special_tokens=True)\n",
    "print(outputs[0])\n",
    "\n",
    "# Time taken: < 10 ms"
   ],
   "id": "bdb64e5b3e4ba17f",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Come up with a 20 letter English word. Here is a 20 letter word: \"pneumonoultramicroscopicsilicovolcanoconiosis\".\n",
      "This is a type of lung disease. It is the longest English word in the Oxford English Dictionary. It was coined by Everett M. Smith, the president of the National Puzzlers' League, in 1935. It refers to a type of lung disease caused by inhaling very fine silica particles.\n",
      "The word is often used to illustrate the extremes of the English language\n"
     ]
    }
   ],
   "execution_count": 28
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}