{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5e32d010-11d0-4be3-a34f-00c87d369347",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: responses 0.18.0 has requirement urllib3>=1.25.10, but you'll have urllib3 1.25.8 which is incompatible.\u001b[0m\n",
      "\u001b[33m  WARNING: The script plasma_store is installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "\u001b[33m  WARNING: The script huggingface-cli is installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "\u001b[33m  WARNING: The script datasets-cli is installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "\u001b[33m  WARNING: The scripts accelerate, accelerate-config and accelerate-launch are installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "\u001b[31mERROR: torchaudio 0.10.1+rocm4.1 has requirement torch==1.10.1, but you'll have torch 2.0.0 which is incompatible.\u001b[0m\n",
      "\u001b[31mERROR: torchvision 0.11.2+cu111 has requirement torch==1.10.1, but you'll have torch 2.0.0 which is incompatible.\u001b[0m\n",
      "\u001b[33m  WARNING: The script transformers-cli is installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "\u001b[33m  WARNING: The script isympy is installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "\u001b[33m  WARNING: The scripts cmake, cpack and ctest are installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "\u001b[33m  WARNING: The script lit is installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n",
      "\u001b[33m  WARNING: The scripts convert-caffe2-to-onnx, convert-onnx-to-caffe2 and torchrun are installed in '/home/qblocks/.local/bin' which is not on PATH.\n",
      "  Consider adding this directory to PATH or, if you prefer to suppress this warning, use --no-warn-script-location.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install -q bitsandbytes datasets accelerate loralib\n",
    "!pip install -q git+https://github.com/huggingface/transformers.git@main git+https://github.com/huggingface/peft.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d35008ce-0d55-4f74-9eb9-c9dcd392a4ce",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import bitsandbytes as bnb\n",
    "from transformers import AutoTokenizer, AutoConfig, AutoModelForCausalLM\n",
    "\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"bigscience/bloom-3b\")\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0efc3e69-f796-46cf-8ee8-52d72f9f653e",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import transformers\n",
    "from datasets import load_dataset\n",
    "from datasets import interleave_datasets\n",
    "data_as = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/as/as.txt\"],split='train',streaming=True)\n",
    "data_bn = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/bn/bn.txt\"],split='train',streaming=True)\n",
    "data_gu = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/gu/gu.txt\"],split='train',streaming=True)\n",
    "data_hi = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/hi/hi.txt\"],split='train',streaming=True)\n",
    "data_kn = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/kn/kn.txt\"],split='train',streaming=True)\n",
    "data_ml = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/ml/ml.txt\"],split='train',streaming=True)\n",
    "data_mr = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/mr/mr.txt\"],split='train',streaming=True)\n",
    "data_or = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/or/or.txt\"],split='train',streaming=True)\n",
    "data_pa = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/pa/pa.txt\"],split='train',streaming=True)\n",
    "data_ta = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/ta/ta.txt\"],split='train',streaming=True)\n",
    "data_te = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/te/te.txt\"],split='train',streaming=True)\n",
    "\n",
    "multilingual_dataset = interleave_datasets([data_as, data_bn,data_gu,data_hi,data_kn,data_ml,data_mr,data_or,data_pa,data_ta,data_te])\n",
    "\n",
    "#data_en = load_dataset(\"aashay96/indic_language_corpus\",data_files=[\"indic_dataset_extracted/data/bn/en.txt\"],streaming=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f61461ed-e91e-45e4-b1cd-c31cf15a6d2d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "multilingual_dataset = multilingual_dataset.map(lambda samples: tokenizer(samples['text'],truncation=True,max_length=1024,padding=True), batched=True)\n",
    "#data.push_to_hub('aashay96/indic_complete_tokenised')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b8ed6593-d80c-4fdb-82e7-7b56b2bbc2c2",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Overriding torch_dtype=None with `torch_dtype=torch.float16` due to requirements of `bitsandbytes` to enable model loading in mixed int8. Either pass torch_dtype=torch.float16 or don't pass this argument at all to remove this warning.\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    \"bigscience/bloom-3b\", \n",
    "    load_in_8bit=True, \n",
    "    device_map='auto',\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6c4d2f2e-da71-42bc-a877-d4e236701f84",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BloomForCausalLM(\n",
       "  (transformer): BloomModel(\n",
       "    (word_embeddings): Embedding(250880, 2560)\n",
       "    (word_embeddings_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "    (h): ModuleList(\n",
       "      (0-29): 30 x BloomBlock(\n",
       "        (input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear8bitLt(in_features=2560, out_features=7680, bias=True)\n",
       "          (dense): Linear8bitLt(in_features=2560, out_features=2560, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear8bitLt(in_features=2560, out_features=10240, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear8bitLt(in_features=10240, out_features=2560, bias=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=2560, out_features=250880, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "90340bb5-8a3a-414a-8b5b-8cf897918381",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for param in model.parameters():\n",
    "  param.requires_grad = False  # freeze the model - train adapters later\n",
    "  if param.ndim == 1:\n",
    "    # cast the small parameters (e.g. layernorm) to fp32 for stability\n",
    "    param.data = param.data.to(torch.float32)\n",
    "\n",
    "model.gradient_checkpointing_enable()  # reduce number of stored activations\n",
    "model.enable_input_require_grads()\n",
    "\n",
    "class CastOutputToFloat(nn.Sequential):\n",
    "  def forward(self, x): return super().forward(x).to(torch.float32)\n",
    "model.lm_head = CastOutputToFloat(model.lm_head)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "963eccdd-a57c-4970-b86c-bf446cc0243a",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BloomForCausalLM(\n",
       "  (transformer): BloomModel(\n",
       "    (word_embeddings): Embedding(250880, 2560)\n",
       "    (word_embeddings_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "    (h): ModuleList(\n",
       "      (0-29): 30 x BloomBlock(\n",
       "        (input_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "        (self_attention): BloomAttention(\n",
       "          (query_key_value): Linear8bitLt(in_features=2560, out_features=7680, bias=True)\n",
       "          (dense): Linear8bitLt(in_features=2560, out_features=2560, bias=True)\n",
       "          (attention_dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "        (post_attention_layernorm): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "        (mlp): BloomMLP(\n",
       "          (dense_h_to_4h): Linear8bitLt(in_features=2560, out_features=10240, bias=True)\n",
       "          (gelu_impl): BloomGelu()\n",
       "          (dense_4h_to_h): Linear8bitLt(in_features=10240, out_features=2560, bias=True)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (ln_f): LayerNorm((2560,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): CastOutputToFloat(\n",
       "    (0): Linear(in_features=2560, out_features=250880, bias=False)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0de04fc8-1541-445d-8a6c-528862e18f69",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def print_trainable_parameters(model):\n",
    "    \"\"\"\n",
    "    Prints the number of trainable parameters in the model.\n",
    "    \"\"\"\n",
    "    trainable_params = 0\n",
    "    all_param = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        all_param += param.numel()\n",
    "        if param.requires_grad:\n",
    "            trainable_params += param.numel()\n",
    "    print(\n",
    "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ac1c4734-530a-4c9c-a055-8c8d3f46b169",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 4915200 || all params: 3007472640 || trainable%: 0.1634329082375293\n"
     ]
    }
   ],
   "source": [
    "from peft import LoraConfig, get_peft_model \n",
    "\n",
    "config = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\"\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, config)\n",
    "print_trainable_parameters(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "683c0239-9384-4d80-b2d0-64738e9c53f5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'train': <datasets.iterable_dataset.IterableDataset at 0x7ff380f5fcd0>}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ab933bdc-8d59-44e3-b210-a5c517660ef3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<datasets.iterable_dataset.IterableDataset at 0x7f0e30bce340>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "edabb62f-d5b3-4d5a-9220-751b940e0a5b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33maashay96\u001b[0m (\u001b[33mindic-lm\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "!pip install wandb\n",
    "import wandb\n",
    "wandb.login()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce63418-3aba-4549-8a50-922a5cf10cb1",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "import transformers\n",
    "from datasets import load_dataset\n",
    "#data = load_dataset(\"Abirate/english_quotes\")\n",
    "#data = data.map(lambda samples: tokenizer(samples['quote']), batched=True)\n",
    "\n",
    "trainer = transformers.Trainer(\n",
    "    model=model, \n",
    "    train_dataset=multilingual_dataset,\n",
    "    args=transformers.TrainingArguments(\n",
    "        per_device_train_batch_size=4, \n",
    "        gradient_accumulation_steps=16,\n",
    "        #gradient_checkpointing=True,\n",
    "        warmup_steps=100, \n",
    "        save_steps=1000,\n",
    "        #num_train_epochs=3,\n",
    "        max_steps=20000, \n",
    "        learning_rate=3e-4, \n",
    "        fp16=True,\n",
    "        logging_steps=1, \n",
    "        output_dir='outputs',report_to='wandb'\n",
    "    ),\n",
    "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False)\n",
    ")\n",
    "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ceeb7a2-7f94-4153-96b0-af19acf90bdb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model.push_to_hub(\"aashay96/indic-BloomLM\", use_auth_token=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "15eb4b53-1354-4729-9cb7-872b057b11be",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      " आप कैसे हैं? आप अपने जीवन में क्या कर रहे हैं?\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: Waiting for W&B process to finish... (success).\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from peft import PeftModel, PeftConfig\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "peft_model_id = \"aashay96/indic-BloomLM\"\n",
    "config = PeftConfig.from_pretrained(peft_model_id)\n",
    "model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')\n",
    "tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)\n",
    "\n",
    "# Load the Lora model\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)\n",
    "\n",
    "\n",
    "\n",
    "batch = tokenizer(\"आप कैसे हैं\", return_tensors='pt')\n",
    "\n",
    "with torch.cuda.amp.autocast():\n",
    "  output_tokens = model.generate(**batch, max_new_tokens=10)\n",
    "\n",
    "print('\\n\\n', tokenizer.decode(output_tokens[0], skip_special_tokens=True))"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}