IFT_Mistral_7B.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e8226704-8e9c-4735-b11e-30760881689d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"WANDB_PROJECT\"]=\"mistral_trans_finetuning\"\n",
    "\n",
    "from enum import Enum\n",
    "from functools import partial\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import bitsandbytes as bnb\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import transformers\n",
    "from datasets import Dataset, DatasetDict, load_dataset\n",
    "from peft import get_peft_model, PeftModel, LoraConfig, TaskType\n",
    "from trl import SFTTrainer\n",
    "from transformers import (AutoModelForCausalLM, \n",
    "                          AutoTokenizer, \n",
    "                          BitsAndBytesConfig, \n",
    "                          TrainingArguments, \n",
    "                          pipeline, \n",
    "                          logging,\n",
    "                          set_seed)\n",
    "from sklearn.metrics import (accuracy_score, \n",
    "                             classification_report, \n",
    "                             confusion_matrix)\n",
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "seed = 42\n",
    "set_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1fabdf70-7353-4f95-a107-d98d4df7d5b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                               Input  \\\n",
      "0                                 His legs are long.   \n",
      "1                Who taught Tom how to speak French?   \n",
      "2                       I swim in the sea every day.   \n",
      "3  Tom popped into the supermarket on his way hom...   \n",
      "4                             Smoke filled the room.   \n",
      "\n",
      "                                              Output  \\\n",
      "0                       అతని కాళ్ళు పొడవుగా ఉన్నాయి.   \n",
      "1            టామ్ ఫ్రెంచ్ మాట్లాడటం ఎలా నేర్పించారు?   \n",
      "2              నేను ప్రతి రోజు సముద్రంలో ఈత కొడతాను.   \n",
      "3  టామ్ కొంచెం పాలు కొనడానికి ఇంటికి వెళ్ళేటప్పుడ...   \n",
      "4                                పొగ గదిని నింపింది.   \n",
      "\n",
      "                                         Instruction  \n",
      "0  You are a helpful assistant who translates Eng...  \n",
      "1  You are a helpful assistant who translates Eng...  \n",
      "2  You are a helpful assistant who translates Eng...  \n",
      "3  You are a helpful assistant who translates Eng...  \n",
      "4  You are a helpful assistant who translates Eng...  \n"
     ]
    }
   ],
   "source": [
    "# Initialize lists to hold the English and Telugu sentences\n",
    "english_sentences = []\n",
    "telugu_sentences = []\n",
    "instruction= \"You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.\"\n",
    "\n",
    "# Open the file and read line by line\n",
    "with open('english_telugu_data.txt', 'r', encoding='utf-8') as f:\n",
    "    for line in f:\n",
    "        # Split the line based on the delimiter '++++$++++'\n",
    "        parts = line.split('++++$++++')\n",
    "        if len(parts) == 2:  # Ensure that there are exactly two parts\n",
    "            english_sentences.append(parts[0].strip())\n",
    "            telugu_sentences.append(parts[1].strip())\n",
    "\n",
    "# Create a DataFrame with the English and Telugu sentences\n",
    "df = pd.DataFrame({\n",
    "    'Input': english_sentences,\n",
    "    'Output': telugu_sentences,\n",
    "    'Instruction': instruction\n",
    "})\n",
    "\n",
    "# Print the first few rows of the DataFrame\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c0ccd04d-ae24-4ecd-8890-39f4cc56456f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "155798\n"
     ]
    }
   ],
   "source": [
    "print(len(df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "921cf450-474a-4ba5-b482-def088ce0d75",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                    Input  \\\n",
      "0                                      His legs are long.   \n",
      "1                     Who taught Tom how to speak French?   \n",
      "2                            I swim in the sea every day.   \n",
      "3       Tom popped into the supermarket on his way hom...   \n",
      "4                                  Smoke filled the room.   \n",
      "...                                                   ...   \n",
      "155793                       I didn't do it deliberately.   \n",
      "155794                       That book isn't interesting.   \n",
      "155795             Tom has never seen the Atlantic Ocean.   \n",
      "155796     Tom didn't even discuss the problem with Mary.   \n",
      "155797                          I've lost my best friend.   \n",
      "\n",
      "                                                   Output  \\\n",
      "0                            అతని కాళ్ళు పొడవుగా ఉన్నాయి.   \n",
      "1                 టామ్ ఫ్రెంచ్ మాట్లాడటం ఎలా నేర్పించారు?   \n",
      "2                   నేను ప్రతి రోజు సముద్రంలో ఈత కొడతాను.   \n",
      "3       టామ్ కొంచెం పాలు కొనడానికి ఇంటికి వెళ్ళేటప్పుడ...   \n",
      "4                                     పొగ గదిని నింపింది.   \n",
      "...                                                   ...   \n",
      "155793                      నేను ఉద్దేశపూర్వకంగా చేయలేదు.   \n",
      "155794                        ఆ పుస్తకం ఆసక్తికరంగా లేదు.   \n",
      "155795                టామ్ అట్లాంటిక్ మహాసముద్రం చూడలేదు.   \n",
      "155796        టామ్ మేరీతో సమస్య గురించి కూడా చర్చించలేదు.   \n",
      "155797              నేను నా బెస్ట్ ఫ్రెండ్ ని కోల్పోయాను.   \n",
      "\n",
      "                                              Instruction  \\\n",
      "0       You are a helpful assistant who translates Eng...   \n",
      "1       You are a helpful assistant who translates Eng...   \n",
      "2       You are a helpful assistant who translates Eng...   \n",
      "3       You are a helpful assistant who translates Eng...   \n",
      "4       You are a helpful assistant who translates Eng...   \n",
      "...                                                   ...   \n",
      "155793  You are a helpful assistant who translates Eng...   \n",
      "155794  You are a helpful assistant who translates Eng...   \n",
      "155795  You are a helpful assistant who translates Eng...   \n",
      "155796  You are a helpful assistant who translates Eng...   \n",
      "155797  You are a helpful assistant who translates Eng...   \n",
      "\n",
      "                                             Conversation  \n",
      "0       [{'content': 'His legs are long.', 'role': 'us...  \n",
      "1       [{'content': 'Who taught Tom how to speak Fren...  \n",
      "2       [{'content': 'I swim in the sea every day.', '...  \n",
      "3       [{'content': 'Tom popped into the supermarket ...  \n",
      "4       [{'content': 'Smoke filled the room.', 'role':...  \n",
      "...                                                   ...  \n",
      "155793  [{'content': 'I didn't do it deliberately.', '...  \n",
      "155794  [{'content': 'That book isn't interesting.', '...  \n",
      "155795  [{'content': 'Tom has never seen the Atlantic ...  \n",
      "155796  [{'content': 'Tom didn't even discuss the prob...  \n",
      "155797  [{'content': 'I've lost my best friend.', 'rol...  \n",
      "\n",
      "[155798 rows x 4 columns]\n"
     ]
    }
   ],
   "source": [
    "# Function to create the 'Conversation' content\n",
    "def create_conversation(row):\n",
    "    return [\n",
    "        {\"content\": row['Input'], \"role\": \"user\"},\n",
    "        {\"content\": row['Output'], \"role\": \"assistant\"}\n",
    "    ]\n",
    "\n",
    "# Apply the function to each row in the DataFrame\n",
    "df['Conversation'] = df.apply(create_conversation, axis=1)\n",
    "\n",
    "# Display the DataFrame with the new 'Conversation' column\n",
    "print(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a0524821-f743-449a-adaa-0a0d61b670ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Adjust display options to avoid truncation\n",
    "pd.set_option('display.max_colwidth', None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6b68738d-7d66-478d-8bc9-f721a104f0ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "                                                               Input  \\\n",
      "0                                                 His legs are long.   \n",
      "1                                Who taught Tom how to speak French?   \n",
      "2                                       I swim in the sea every day.   \n",
      "3  Tom popped into the supermarket on his way home to buy some milk.   \n",
      "4                                             Smoke filled the room.   \n",
      "\n",
      "                                                                             Output  \\\n",
      "0                                                      అతని కాళ్ళు పొడవుగా ఉన్నాయి.   \n",
      "1                                           టామ్ ఫ్రెంచ్ మాట్లాడటం ఎలా నేర్పించారు?   \n",
      "2                                             నేను ప్రతి రోజు సముద్రంలో ఈత కొడతాను.   \n",
      "3  టామ్ కొంచెం పాలు కొనడానికి ఇంటికి వెళ్ళేటప్పుడు సూపర్ మార్కెట్లోకి ప్రవేశించాడు.   \n",
      "4                                                               పొగ గదిని నింపింది.   \n",
      "\n",
      "                                                                                                                                                             Instruction  \\\n",
      "0  You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.   \n",
      "1  You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.   \n",
      "2  You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.   \n",
      "3  You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.   \n",
      "4  You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.   \n",
      "\n",
      "                                                                                                                                                                                                               Conversation  \n",
      "0                                                                                                     [{'content': 'His legs are long.', 'role': 'user'}, {'content': 'అతని కాళ్ళు పొడవుగా ఉన్నాయి.', 'role': 'assistant'}]  \n",
      "1                                                                         [{'content': 'Who taught Tom how to speak French?', 'role': 'user'}, {'content': 'టామ్ ఫ్రెంచ్ మాట్లాడటం ఎలా నేర్పించారు?', 'role': 'assistant'}]  \n",
      "2                                                                                  [{'content': 'I swim in the sea every day.', 'role': 'user'}, {'content': 'నేను ప్రతి రోజు సముద్రంలో ఈత కొడతాను.', 'role': 'assistant'}]  \n",
      "3  [{'content': 'Tom popped into the supermarket on his way home to buy some milk.', 'role': 'user'}, {'content': 'టామ్ కొంచెం పాలు కొనడానికి ఇంటికి వెళ్ళేటప్పుడు సూపర్ మార్కెట్లోకి ప్రవేశించాడు.', 'role': 'assistant'}]  \n",
      "4                                                                                                          [{'content': 'Smoke filled the room.', 'role': 'user'}, {'content': 'పొగ గదిని నింపింది.', 'role': 'assistant'}]  \n"
     ]
    }
   ],
   "source": [
    "print(df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "55e95c8f-dd68-498b-a421-aa09ec3b0a70",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['Input', 'Output', 'Instruction', 'Conversation'],\n",
      "        num_rows: 155798\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "# Convert pandas DataFrame to a DatasetDict\n",
    "dataset_dict = DatasetDict({\n",
    "    'train': Dataset.from_pandas(df)\n",
    "})\n",
    "\n",
    "# Print the dataset_dict structure\n",
    "print(dataset_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ad3f431b-6ff7-4ea4-b3b1-136ab17a9f65",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset=dataset_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f1cd7a31-f2b0-49f3-ac69-2c8d52fafb98",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8f0a25918a4348f4ba7fae0c5be429a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/155798 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['content'],\n",
      "        num_rows: 140218\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['content'],\n",
      "        num_rows: 15580\n",
      "    })\n",
      "})\n",
      "{'content': '<|im_start|>system\\nYou are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.<|im_end|>\\n<|im_start|>user\\nWe saw a mountain in the distance.<|im_end|>\\n<|im_start|>assistant\\nదూరం లో ఒక పర్వతం చూశాము.<|im_end|>\\n'}\n"
     ]
    }
   ],
   "source": [
    "model_name = \"mistralai/Mistral-7B-Instruct-v0.2\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "template = \"\"\"{% for message in messages %}\\n{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}{% if loop.last and add_generation_prompt %}{{'<|im_start|>assistant\\n' }}{% endif %}{% endfor %}\"\"\"\n",
    "tokenizer.chat_template = template\n",
    "\n",
    "def preprocess(samples):\n",
    "    batch = []\n",
    "    for system_prompt, conversation in zip(samples[\"Instruction\"], samples[\"Conversation\"]):\n",
    "        system_message = {\"role\": \"system\", \"content\": f\"{system_prompt}\"}\n",
    "        conversation.insert(0, system_message)\n",
    "        batch.append(tokenizer.apply_chat_template(conversation, tokenize=False))\n",
    "    return {\"content\": batch}\n",
    "\n",
    "dataset = dataset.map(\n",
    "    preprocess,\n",
    "    batched=True,\n",
    "    remove_columns=dataset[\"train\"].column_names\n",
    ")\n",
    "dataset = dataset[\"train\"].train_test_split(0.1)\n",
    "print(dataset)\n",
    "print(dataset[\"train\"][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d2d7af75-25b8-4ab8-affc-e422f7b28a64",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>system\n",
      "You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.<|im_end|>\n",
      "<|im_start|>user\n",
      "I saw a mouse.<|im_end|>\n",
      "<|im_start|>assistant\n",
      "నేను ఎలుకను చూశాను.<|im_end|>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(dataset[\"train\"][6][\"content\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ed334da1-fc1f-414e-a816-969d85a43c4f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<|im_start|>system\n",
      "You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.<|im_end|>\n",
      "<|im_start|>user\n",
      "We'll succeed.<|im_end|>\n",
      "<|im_start|>assistant\n",
      "మేము విజయం సాధిస్తాము.<|im_end|>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(dataset[\"test\"][0][\"content\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f92df4d6-4cd5-4866-9088-5cb27442d726",
   "metadata": {},
   "outputs": [],
   "source": [
    "peft_config = LoraConfig(r=8,\n",
    "                         lora_alpha=16,\n",
    "                         lora_dropout=0.1,\n",
    "                         target_modules=[\"gate_proj\",\"q_proj\",\"lm_head\",\"o_proj\",\"k_proj\",\"embed_tokens\",\"down_proj\",\"up_proj\",\"v_proj\"],\n",
    "                         task_type=TaskType.CAUSAL_LM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e047edb1-9794-42a0-b385-2704e547347e",
   "metadata": {},
   "outputs": [],
   "source": [
    "bnb_config = BitsAndBytesConfig(\n",
    "            load_in_4bit=True,\n",
    "            bnb_4bit_quant_type=\"nf4\",\n",
    "            bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "            bnb_4bit_use_double_quant=True,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "293eb47e-a516-4e64-b762-3d20e8a360eb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64ee3aa9804846e6be3166bb6788fae3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Embedding(32008, 4096)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class ChatmlSpecialTokens(str, Enum):\n",
    "    user = \"<|im_start|>user\"\n",
    "    assistant = \"<|im_start|>assistant\"\n",
    "    system = \"<|im_start|>system\"\n",
    "    function_call = \"<|im_start|>function-call\"\n",
    "    function_response = \"<|im_start|>function-response\"\n",
    "    eos_token = \"<|im_end|>\"\n",
    "    bos_token = \"<s>\"\n",
    "    pad_token = \"<pad>\"\n",
    "\n",
    "    @classmethod\n",
    "    def list(cls):\n",
    "        return [c.value for c in cls]\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "        model_name,\n",
    "        pad_token=ChatmlSpecialTokens.pad_token.value,\n",
    "        bos_token=ChatmlSpecialTokens.bos_token.value,\n",
    "        eos_token=ChatmlSpecialTokens.eos_token.value,\n",
    "        additional_special_tokens=ChatmlSpecialTokens.list(),\n",
    "        trust_remote_code=True\n",
    "    )\n",
    "tokenizer.chat_template = template\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name,\n",
    "                                             quantization_config=bnb_config,\n",
    "                                             device_map=\"auto\",\n",
    "                                             attn_implementation=\"flash_attention_2\")\n",
    "model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "f68ef14c-96a2-4a39-9da7-adc05d0da368",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MistralForCausalLM(\n",
       "  (model): MistralModel(\n",
       "    (embed_tokens): Embedding(32008, 4096)\n",
       "    (layers): ModuleList(\n",
       "      (0-31): 32 x MistralDecoderLayer(\n",
       "        (self_attn): MistralFlashAttention2(\n",
       "          (q_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "          (k_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "          (v_proj): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "          (o_proj): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "          (rotary_emb): MistralRotaryEmbedding()\n",
       "        )\n",
       "        (mlp): MistralMLP(\n",
       "          (gate_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "          (up_proj): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "          (down_proj): Linear4bit(in_features=14336, out_features=4096, bias=False)\n",
       "          (act_fn): SiLU()\n",
       "        )\n",
       "        (input_layernorm): MistralRMSNorm()\n",
       "        (post_attention_layernorm): MistralRMSNorm()\n",
       "      )\n",
       "    )\n",
       "    (norm): MistralRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=32008, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "83cb1773-f158-415f-9926-d6b27c04f41f",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_dir = \"Translator_Eng_Tel_instruct\"\n",
    "per_device_train_batch_size = 2\n",
    "per_device_eval_batch_size = 2\n",
    "gradient_accumulation_steps = 4\n",
    "logging_steps = 5\n",
    "learning_rate = 5e-4\n",
    "max_grad_norm = 1.0\n",
    "num_train_epochs=2\n",
    "warmup_ratio = 0.1\n",
    "lr_scheduler_type = \"cosine\"\n",
    "max_seq_length = 2048\n",
    "\n",
    "training_arguments = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    per_device_train_batch_size=per_device_train_batch_size,\n",
    "    per_device_eval_batch_size=per_device_eval_batch_size,\n",
    "    gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "    save_strategy=\"no\",\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    logging_steps=logging_steps,\n",
    "    learning_rate=learning_rate,\n",
    "    max_grad_norm=max_grad_norm,\n",
    "    weight_decay=0.1,\n",
    "    warmup_ratio=warmup_ratio,\n",
    "    lr_scheduler_type=lr_scheduler_type,\n",
    "    bf16=True,\n",
    "    report_to=[\"tensorboard\", \"wandb\"],\n",
    "    hub_private_repo=True,\n",
    "    push_to_hub=True,\n",
    "    num_train_epochs=num_train_epochs,\n",
    "    gradient_checkpointing=True,\n",
    "    gradient_checkpointing_kwargs={\"use_reentrant\": False}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "eda79755-81bd-459d-ac70-a3fbe9f885d4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/trl/trainer/sft_trainer.py:318: UserWarning: You passed a tokenizer with `padding_side` not equal to `right` to the SFTTrainer. This might lead to some unexpected behaviour due to overflow issues when training a model in half-precision. You might consider adding `tokenizer.padding_side = 'right'` to your code.\n",
      "  warnings.warn(\n",
      "Detected kernel version 5.4.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n"
     ]
    }
   ],
   "source": [
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=training_arguments,\n",
    "    train_dataset=dataset[\"train\"],\n",
    "    eval_dataset=dataset[\"test\"],\n",
    "    tokenizer=tokenizer,\n",
    "    packing=True,\n",
    "    dataset_text_field=\"content\",\n",
    "    max_seq_length=max_seq_length,\n",
    "    peft_config=peft_config,\n",
    "    dataset_kwargs={\n",
    "        \"append_concat_token\": False,\n",
    "        \"add_special_tokens\": False,\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3dc27438-5c67-4aa6-b02a-8626a7858b3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mrohithreddy24242\u001b[0m (\u001b[33mmrr24\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.17.2"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/workspace/FT_LLMs/Translation/wandb/run-20240622_153532-t4yz6t86</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/mrr24/mistral_trans_finetuning/runs/t4yz6t86' target=\"_blank\">golden-cherry-2</a></strong> to <a href='https://wandb.ai/mrr24/mistral_trans_finetuning' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/mrr24/mistral_trans_finetuning' target=\"_blank\">https://wandb.ai/mrr24/mistral_trans_finetuning</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/mrr24/mistral_trans_finetuning/runs/t4yz6t86' target=\"_blank\">https://wandb.ai/mrr24/mistral_trans_finetuning/runs/t4yz6t86</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n",
      "The input hidden states seems to be silently casted in float32, this might be related to the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in torch.bfloat16.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1912' max='1912' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1912/1912 8:46:07, Epoch 1/2]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>0.254900</td>\n",
       "      <td>0.255433</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.213800</td>\n",
       "      <td>0.212012</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/peft/utils/save_and_load.py:139: UserWarning: Setting `save_embedding_layers` to `True` as embedding layers found in `target_modules`.\n",
      "  warnings.warn(\"Setting `save_embedding_layers` to `True` as embedding layers found in `target_modules`.\")\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c3d602fb93ab401da2bbeb5b154617e4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "events.out.tfevents.1719070142.23ec79e11277.2261.0:   0%|          | 0.00/5.64k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "44c71f1ccf604efe9f1a91f4bfa4fa38",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "events.out.tfevents.1719070531.23ec79e11277.2799.0:   0%|          | 0.00/86.6k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c5f596be4e8a4b25b8e7aa2370f2c9f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Upload 5 LFS files:   0%|          | 0/5 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b59d178cae0b423da74a82b9297c65d7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.model:   0%|          | 0.00/493k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee26fad8fb78435a89c1d90107619306",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "adapter_model.safetensors:   0%|          | 0.00/568M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "944872a27a6642bc984926318c97fc9c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "training_args.bin:   0%|          | 0.00/5.05k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.train()\n",
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "7851cc2d-514f-4329-8528-3308d24a21fe",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6faf78e2edb14adc9932e5df457aef7c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): MistralForCausalLM(\n",
       "      (model): MistralModel(\n",
       "        (embed_tokens): lora.Embedding(\n",
       "          (base_layer): Embedding(32008, 4096)\n",
       "          (lora_dropout): ModuleDict(\n",
       "            (default): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (lora_A): ModuleDict()\n",
       "          (lora_B): ModuleDict()\n",
       "          (lora_embedding_A): ParameterDict(  (default): Parameter containing: [torch.cuda.HalfTensor of size 8x32008 (cuda:0)])\n",
       "          (lora_embedding_B): ParameterDict(  (default): Parameter containing: [torch.cuda.HalfTensor of size 4096x8 (cuda:0)])\n",
       "        )\n",
       "        (layers): ModuleList(\n",
       "          (0-31): 32 x MistralDecoderLayer(\n",
       "            (self_attn): MistralFlashAttention2(\n",
       "              (q_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (v_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (rotary_emb): MistralRotaryEmbedding()\n",
       "            )\n",
       "            (mlp): MistralMLP(\n",
       "              (gate_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=14336, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (up_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=14336, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (down_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=14336, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=14336, out_features=8, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=8, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (act_fn): SiLU()\n",
       "            )\n",
       "            (input_layernorm): MistralRMSNorm()\n",
       "            (post_attention_layernorm): MistralRMSNorm()\n",
       "          )\n",
       "        )\n",
       "        (norm): MistralRMSNorm()\n",
       "      )\n",
       "      (lm_head): lora.Linear(\n",
       "        (base_layer): Linear(in_features=4096, out_features=32008, bias=False)\n",
       "        (lora_dropout): ModuleDict(\n",
       "          (default): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (lora_A): ModuleDict(\n",
       "          (default): Linear(in_features=4096, out_features=8, bias=False)\n",
       "        )\n",
       "        (lora_B): ModuleDict(\n",
       "          (default): Linear(in_features=8, out_features=32008, bias=False)\n",
       "        )\n",
       "        (lora_embedding_A): ParameterDict()\n",
       "        (lora_embedding_B): ParameterDict()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from peft import PeftModel, PeftConfig\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n",
    "from datasets import load_dataset\n",
    "import torch\n",
    "\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "            load_in_4bit=True,\n",
    "            bnb_4bit_quant_type=\"nf4\",\n",
    "            bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "            bnb_4bit_use_double_quant=True,\n",
    "        )\n",
    "\n",
    "peft_model_id = \"Translator_Eng_Tel_instruct\"\n",
    "device = \"cuda\"\n",
    "config = PeftConfig.from_pretrained(peft_model_id)\n",
    "model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path,\n",
    "                                             quantization_config=bnb_config,\n",
    "                                             device_map=\"auto\",\n",
    "                                             attn_implementation=\"flash_attention_2\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(peft_model_id)\n",
    "model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=8)\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)\n",
    "# model.to(torch.bfloat16)\n",
    "# model.cuda()\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e8599d1-fc04-4da2-88bf-6719db858290",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Setting `pad_token_id` to `eos_token_id`:32000 for open-end generation.\n"
     ]
    }
   ],
   "source": [
    "system_prompt = \"\"\"You are a helpful assistant who translates English statements, words, or lines into Telugu. Use your language skills to provide accurate and meaningful translations.\"\"\"\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": system_prompt},\n",
    "    {\"role\": \"user\", \"content\": \"what are you doing\"},\n",
    "]\n",
    "text = tokenizer.apply_chat_template(messages, tokenize=False)\n",
    "inputs = tokenizer(text, return_tensors=\"pt\", add_special_tokens=False)\n",
    "inputs = {k: v.to(\"cuda\") for k,v in inputs.items()}\n",
    "with torch.autocast(dtype=torch.bfloat16, device_type=\"cuda\"):\n",
    "    outputs = model.generate(**inputs, \n",
    "                             max_new_tokens=128, \n",
    "                             do_sample=True, \n",
    "                             top_p=0.95, \n",
    "                             temperature=0.2, \n",
    "                             repetition_penalty=1.0, \n",
    "                             eos_token_id=tokenizer.eos_token_id)\n",
    "print(tokenizer.decode(outputs[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e7bb9ab-3b26-48aa-8543-677d6ca767c6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5fdbef6-7e42-446e-930a-3963ce7bcbfa",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03a3f053-4acb-4286-9b78-59504c71eb91",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c5e1f36-0700-4335-a16c-a07b02ae3c04",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52851a7a-bb4d-4510-a62c-620a44d493e4",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}