{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "26eca0b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "===================================BUG REPORT===================================\n",
      "Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
      "================================================================================\n",
      "CUDA SETUP: CUDA runtime path found: /home/lxe/miniconda3/envs/llama-finetuner/lib/libcudart.so\n",
      "CUDA SETUP: Highest compute capability among GPUs detected: 8.6\n",
      "CUDA SETUP: Detected CUDA version 117\n",
      "CUDA SETUP: Loading binary /home/lxe/miniconda3/envs/llama-finetuner/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda117.so...\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import transformers\n",
    "import peft"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3c2f7268",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = transformers.AutoModelForCausalLM.from_pretrained(\n",
    "    'cerebras/Cerebras-GPT-2.7B', \n",
    "    load_in_8bit=True,\n",
    "    torch_dtype=torch.float16,\n",
    "    device_map='auto'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0dcc11cc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GPT2Config {\n",
       "  \"_name_or_path\": \"cerebras/Cerebras-GPT-2.7B\",\n",
       "  \"activation_function\": \"gelu\",\n",
       "  \"attn_pdrop\": 0.0,\n",
       "  \"bos_token_id\": 50256,\n",
       "  \"embd_pdrop\": 0.0,\n",
       "  \"eos_token_id\": 50256,\n",
       "  \"initializer_range\": 0.02,\n",
       "  \"layer_norm_epsilon\": 1e-05,\n",
       "  \"model_type\": \"gpt2\",\n",
       "  \"n_embd\": 2560,\n",
       "  \"n_head\": 32,\n",
       "  \"n_inner\": 10240,\n",
       "  \"n_layer\": 32,\n",
       "  \"n_positions\": 2048,\n",
       "  \"reorder_and_upcast_attn\": false,\n",
       "  \"resid_pdrop\": 0.0,\n",
       "  \"scale_attn_by_inverse_layer_idx\": false,\n",
       "  \"scale_attn_weights\": true,\n",
       "  \"summary_activation\": null,\n",
       "  \"summary_first_dropout\": 0.1,\n",
       "  \"summary_proj_to_labels\": true,\n",
       "  \"summary_type\": \"cls_index\",\n",
       "  \"summary_use_proj\": true,\n",
       "  \"torch_dtype\": \"float16\",\n",
       "  \"transformers_version\": \"4.28.0.dev0\",\n",
       "  \"use_cache\": true,\n",
       "  \"vocab_size\": 50257\n",
       "}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e8a19a75",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'LLaMATokenizer'. \n",
      "The class this function is called from is 'LlamaTokenizer'.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = transformers.AutoTokenizer.from_pretrained('decapoda-research/llama-7b-hf')\n",
    "# tokenizer.pad_token_id = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "240a9c8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = peft.PeftModel.from_pretrained(\n",
    "    model,\n",
    "    'lora-assistant',\n",
    "    torch_dtype=torch.float16\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4f944f46",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human: How old is the sun?\n",
      "Assistant: I'm not a member.\n",
      "I am you, but it was and that he's.\n"
     ]
    }
   ],
   "source": [
    "inputs = tokenizer(\"Human: How old is the sun?\\nAssistant:\", return_tensors=\"pt\")\n",
    "input_ids = inputs[\"input_ids\"].to('cuda')\n",
    "\n",
    "generation_config = transformers.GenerationConfig(\n",
    "    do_sample = True,\n",
    "    temperature = 0.3,\n",
    "    top_p = 0.1,\n",
    "    top_k = 80,\n",
    "    repetition_penalty = 1.5,\n",
    "    max_new_tokens = 50\n",
    ")\n",
    "\n",
    "with torch.no_grad():\n",
    "    generation_output = model.generate(\n",
    "        input_ids=input_ids,\n",
    "        attention_mask=torch.ones_like(input_ids),\n",
    "        generation_config=generation_config,\n",
    "    )\n",
    "    \n",
    "output_text = tokenizer.decode(generation_output[0].cuda(), skip_special_tokens=True).strip()\n",
    "print(output_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "5fc13b1a",
   "metadata": {},
   "outputs": [],
   "source": [
    "del model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5f19b3a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}