{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "76d4badf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/hour1/anaconda3/envs/intervene/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a93a924e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained('/mnt/d/LLM/Llama-3-Base-8B-SFT', use_fast = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5287d852",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 4/4 [00:00<00:00,  8.94it/s]\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    '/mnt/d/LLM/Llama-3-Base-8B-SFT', #token=hf_token\n",
    "    torch_dtype='auto'\n",
    ").to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "90b9b4ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Who are you?\"},\n",
    "    {\"role\": \"assistant\", \"content\": \"I'm nobody.\"},\n",
    "    {\"role\": \"user\", \"content\": \"What are you talking about?\"},\n",
    "]\n",
    "\n",
    "input_ids1 = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    add_generation_prompt=True,\n",
    "    return_tensors=\"pt\"\n",
    ").to(model.device)\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"What can you do?\"},\n",
    "]\n",
    "\n",
    "input_ids2 = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    add_generation_prompt=True,\n",
    "    return_tensors=\"pt\"\n",
    ").to(model.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f39a4415",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_ids = [input_ids1[0], input_ids2[0]]\n",
    "batch_mask = [torch.ones_like(element) for element in batch_ids]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84d1e607",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.padding_side='left'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "11036705",
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = {\"input_ids\": batch_ids, \"attention_mask\": batch_mask}\n",
    "padded_inputs = tokenizer.pad(\n",
    "    inputs,\n",
    "    padding=True,\n",
    "    max_length=None,\n",
    "    return_tensors=\"pt\",\n",
    ").to('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "32146865",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input_ids': tensor([[128000, 128006,    882, 128007,    271,  15546,    527,    499,     30,\n",
       "         128009, 128006,  78191, 128007,    271,     40,   2846,  19093,     13,\n",
       "         128009, 128006,    882, 128007,    271,   3923,    527,    499,   7556,\n",
       "            922,     30, 128009, 128006,  78191, 128007,    271],\n",
       "        [128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001,\n",
       "         128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001, 128001,\n",
       "         128001, 128000, 128006,    882, 128007,    271,   3923,    649,    499,\n",
       "            656,     30, 128009, 128006,  78191, 128007,    271]],\n",
       "       device='cuda:0'), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       "        [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1,\n",
       "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], device='cuda:0')}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "padded_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4ae57446",
   "metadata": {},
   "outputs": [],
   "source": [
    "terminators = [\n",
    "    tokenizer.eos_token_id,\n",
    "    tokenizer.convert_tokens_to_ids(\"<|eot_id|>\")\n",
    "]\n",
    "generation_kwargs = {\n",
    "    \"max_new_tokens\": 128,\n",
    "    \"min_length\": -1,\n",
    "    \"top_k\": 0.0,\n",
    "    \"top_p\": 1.0, \n",
    "    \"do_sample\": True,\n",
    "    \"temperature\": 0.7,\n",
    "    \"pad_token_id\": tokenizer.eos_token_id,\n",
    "    \"begin_suppress_tokens\": [tokenizer.eos_token_id],\n",
    "    'eos_token_id': terminators\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "473c9f35",
   "metadata": {},
   "outputs": [],
   "source": [
    "generations = model.generate(**padded_inputs, **generation_kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "724f9999",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = []\n",
    "for generation, mask in zip(generations, padded_inputs[\"attention_mask\"]):\n",
    "    output = generation[(1 - mask).sum() :]  # remove padding\n",
    "\n",
    "    output = output[(mask).sum() :]  # remove prompt\n",
    "\n",
    "    if tokenizer.eos_token_id in output:\n",
    "        pad_mask = output == tokenizer.eos_token_id\n",
    "        pad_start = torch.nonzero(pad_mask, as_tuple=False)[0, 0].item()\n",
    "        output = output[: pad_start + 1]  # keep the eos token at the end\n",
    "\n",
    "    outputs.append(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "8d25f8a1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([    40,   2846,  14931,    422,    856,   2077,  22568,    499,     13,\n",
       "            358,    574,   4560,    311,   3237,    430,    358,    656,    539,\n",
       "            617,    264,   3230,   9764,    477,  17743,   1093,  12966,    656,\n",
       "             13,    358,   3073,    439,    264,   2068,    311,   7945,    449,\n",
       "           4221,   8863,    323,   3493,   2038,     11,    719,    358,   1541,\n",
       "            956,    617,  11555,    477,  16024,   1093,  12966,    656,     13,\n",
       "         128009, 128001], device='cuda:0'),\n",
       " tensor([    40,    649,   3619,    323,   6013,    311,   1495,    477,   7899,\n",
       "          11545,     11,    719,    358,   1097,    539,  13171,    315,  16785,\n",
       "           7106,   6299,    477,   9256,    389,    856,   1866,     13,   4452,\n",
       "             11,    358,    649,   3493,   2038,    323,  19351,    389,   5370,\n",
       "          13650,    323,   7945,    449,   9256,   1778,    439,  38952,     11,\n",
       "           6376,  66697,     11,    477,  15389,    369,   2038,    389,    279,\n",
       "           7757,     13, 128009], device='cuda:0')]"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "b64e75e3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[\"I'm sorry if my response confused you. I was trying to express that I do not have a specific identity or personality like humans do. I exist as a program to assist with language processing and provide information, but I don't have thoughts or feelings like humans do.\",\n",
       " 'I can understand and respond to text or voice commands, but I am not capable of performing physical actions or tasks on my own. However, I can provide information and guidance on various topics and assist with tasks such as scheduling, setting reminders, or searching for information on the internet.']"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.batch_decode(outputs, skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd3d02b4",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "intervene",
   "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}