finetune_seq2seq.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "225b0a06-4dd3-4c84-89ed-86de4adcd24b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://huggingface.co/docs/peft/index\n",
    "\n",
    "import os \n",
    "from transformers import GPT2Tokenizer, GPT2LMHeadModel\n",
    "\n",
    "### Utils\n",
    "def chunks(lst, n):\n",
    "    \"\"\"Yield successive n-sized chunks from lst.\"\"\"\n",
    "    for i in range(0, len(lst), n):\n",
    "        yield lst[i:i + n]\n",
    "\n",
    "def split_strings(strings):\n",
    "    first_parts = []\n",
    "    second_parts = []\n",
    "    for string in strings:\n",
    "        split_index = len(string) // 2\n",
    "        while split_index < len(string) and not string[split_index].isspace():\n",
    "            split_index += 1\n",
    "        first_part = string[:split_index].rstrip()\n",
    "        second_part = string[split_index:].lstrip()\n",
    "        first_parts.append(first_part)\n",
    "        second_parts.append(second_part)\n",
    "    return first_parts, second_parts\n",
    "    \n",
    "### \n",
    "\n",
    "math_example_path = os.getcwd() + '/data/What_Is_Mathematics_An_Elementary_Approach_to_Ideas_and_Methods.txt'\n",
    "\n",
    "with open(math_example_path, \"r\",  encoding=\"utf8\") as f:\n",
    "     math_example_text = f.read()\n",
    "\n",
    "math_example_text = math_example_text.replace(\"\\n\", \"\")\n",
    "math_sentences = math_example_text.split(\".\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b28fa07d-ee0a-44cb-b5f0-f14226be3244",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "9421\n",
      " But it is not a concession to the dangerous tendency toward dodging all exertion\n"
     ]
    }
   ],
   "source": [
    "print(len(math_sentences))\n",
    "math_sentences = [item for item in math_sentences if 80 < len(item) < 100]\n",
    "X, Y = split_strings(math_sentences)\n",
    "print(math_sentences[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "7f996a8d-7395-419c-94be-75f2c6d80dc5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " For a « 20, we need twenty words for the digits, plus words for 20 and 400, making a total of 22 \n",
      " X:   For a « 20, we need twenty words for the digits, \n",
      " Y:  plus words for 20 and 400, making a total of 22\n"
     ]
    }
   ],
   "source": [
    "tst_ind = 30\n",
    "print(math_sentences[tst_ind], \"\\n X: \", X[tst_ind], \"\\n Y: \", Y[tst_ind])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "6c1921fe-4e0e-43d8-aef3-07e307f890d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Embedding(50257, 768)"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch \n",
    "\n",
    "output_dir = \"./finetuned_llm_all\"\n",
    "\n",
    "model_name = \"gpt2\" #output_dir #\n",
    "tokenizer = GPT2Tokenizer.from_pretrained(model_name, padding_side='left') #gpt2\n",
    "llm_model = GPT2LMHeadModel.from_pretrained(model_name) # gpt2\n",
    "device = torch.device('cuda')\n",
    "llm_model.to(device)\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Configure the model\n",
    "llm_model.config.pad_token_id = tokenizer.eos_token_id\n",
    "llm_model.config.eos_token_id = tokenizer.eos_token_id\n",
    "llm_model.config.vocab_size = llm_model.config.vocab_size + len(tokenizer.get_added_vocab())\n",
    "llm_model.resize_token_embeddings(len(tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "f76d6675-bfda-4499-887a-4db95eed1f6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#len_params = len(list(llm_model.parameters()))\n",
    "#for i, param in enumerate(llm_model.parameters()):\n",
    "#    param.requires_grad = i/len_params > 0.8 # train last 20% of layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "dbd2e4e4-6cc3-49ff-a68f-9ba24585de7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "#np.random.seed(309)\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "llm_model.config.pad_token_id = tokenizer.pad_token_id\n",
    "\n",
    "# Set padding to the left side as the model is a decoder-only architecture\n",
    "tokenizer.padding_side = \"left\"\n",
    "\n",
    "encoded_X = tokenizer.batch_encode_plus(X, add_special_tokens=False, return_tensors='pt', padding=True)\n",
    "encoded_Y = tokenizer.batch_encode_plus(Y, add_special_tokens=False, return_tensors='pt', padding=True)\n",
    "\n",
    "test_ids = np.random.choice(range(len(X)), 5)\n",
    "test_X = [X[i] for i in test_ids]\n",
    "enc_test_X = tokenizer.batch_encode_plus(test_X, \n",
    "                                     add_special_tokens=False, \n",
    "                                     return_tensors='pt', \n",
    "                                     padding=True) # X is encoded Y isnt !\n",
    "\n",
    "\n",
    "test_Y = [Y[i] for i in test_ids]\n",
    "\n",
    "BATH_SIZE = 10\n",
    "\n",
    "x_batches = list(chunks(encoded_X[\"input_ids\"], BATH_SIZE))\n",
    "y_batches = list(chunks(encoded_Y[\"input_ids\"], BATH_SIZE))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "cbb2cf1c-7c91-4341-8219-16bbebb79294",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input:  Peaucellier, invented a simple linkage that solved\n",
      "output: <|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|> Peaucellier, invented a simple linkage that solved the problem of how to connect two different types of cells.\n",
      "\n",
      "The Peaucellier link is a simple, yet powerful, way to connect two different types of cells.\n",
      "\n",
      "The Peaucellier link is a simple, yet powerful, way to connect two different types of cells.\n",
      "\n",
      "The Peaucellier link is a simple, yet powerful, way to connect two different types of cells.\n",
      "\n",
      "The Peaucellier link is a simple, yet powerful, way\n"
     ]
    }
   ],
   "source": [
    "test_X = [X[i] for i in test_ids]\n",
    "enc_test_X = tokenizer.batch_encode_plus(test_X, \n",
    "                                     add_special_tokens=False, \n",
    "                                     return_tensors='pt', \n",
    "                                     padding=True) # X is encoded Y isnt !\n",
    "                                         \n",
    "output = llm_model.generate(enc_test_X[\"input_ids\"].to(device),  \n",
    "                            max_new_tokens=100, \n",
    "                            pad_token_id=tokenizer.eos_token_id)\n",
    "\n",
    "# Set padding to the left side as the model is a decoder-only architecture\n",
    "\n",
    "for i, text in enumerate(test_X[0:1]):\n",
    "    print(\"input:\", text) # <<<!\n",
    "    print(\"output:\", tokenizer.decode(output[i]))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "33487168-8fbb-4b64-884e-e1eaaaede66c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "output:  Peaucellier, invented a simple linkage that solved the problem of how to connect two wires.\n",
      "\n",
      "The first step was to make a simple, inexpensive, and easy-to-use link between two wires. The second step was to make a simple, inexpensive, and easy-to-use link between two wires.\n",
      "\n",
      "The first step was to make a simple, inexpensive, and easy-to-use link between two wires.\n",
      "\n",
      "The first step was to make a simple, inexpensive, and easy-to-use link between\n"
     ]
    }
   ],
   "source": [
    "# Prepare input data\n",
    "input_ids = tokenizer.encode(test_X[0] ,\n",
    "                             add_special_tokens=False, \n",
    "                             return_tensors='pt', \n",
    "                             padding=True).to(device)\n",
    "\n",
    "# Generate text using generation parameters\n",
    "output = llm_model.generate(input_ids , \n",
    "                            max_new_tokens=100, \n",
    "                            pad_token_id=tokenizer.eos_token_id)\n",
    "\n",
    "# Decode and print the output\n",
    "decoded_output = tokenizer.decode(output[0])\n",
    "print(\"output:\", decoded_output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "ae25b5ed-50c6-4337-8134-2bff8862908f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "optimizer = torch.optim.AdamW(llm_model.parameters(), lr=1e-5)\n",
    "\n",
    "def get_lr(optimizer):\n",
    "    for param_group in optimizer.param_groups:\n",
    "        return param_group['lr']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "dcbbcac2-39c4-4ec3-9f7c-3ea27a109688",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 98])\n",
      "torch.Size([10, 97])\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Expected input batch_size (970) to match target batch_size (960).",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[62], line 17\u001b[0m\n\u001b[0;32m     14\u001b[0m \u001b[38;5;28mprint\u001b[39m(batch\u001b[38;5;241m.\u001b[39mshape)\n\u001b[0;32m     15\u001b[0m \u001b[38;5;28mprint\u001b[39m(y_batches[index]\u001b[38;5;241m.\u001b[39mshape)\n\u001b[1;32m---> 17\u001b[0m outputs \u001b[38;5;241m=\u001b[39m \u001b[43mllm_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlabels\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43my_batches\u001b[49m\u001b[43m[\u001b[49m\u001b[43mindex\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     18\u001b[0m loss \u001b[38;5;241m=\u001b[39m outputs\u001b[38;5;241m.\u001b[39mloss\n\u001b[0;32m     19\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\nn\\modules\\module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\transformers\\models\\gpt2\\modeling_gpt2.py:1107\u001b[0m, in \u001b[0;36mGPT2LMHeadModel.forward\u001b[1;34m(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[0;32m   1105\u001b[0m     \u001b[38;5;66;03m# Flatten the tokens\u001b[39;00m\n\u001b[0;32m   1106\u001b[0m     loss_fct \u001b[38;5;241m=\u001b[39m CrossEntropyLoss()\n\u001b[1;32m-> 1107\u001b[0m     loss \u001b[38;5;241m=\u001b[39m \u001b[43mloss_fct\u001b[49m\u001b[43m(\u001b[49m\u001b[43mshift_logits\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mview\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshift_logits\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msize\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshift_labels\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mview\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1109\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m return_dict:\n\u001b[0;32m   1110\u001b[0m     output \u001b[38;5;241m=\u001b[39m (lm_logits,) \u001b[38;5;241m+\u001b[39m transformer_outputs[\u001b[38;5;241m1\u001b[39m:]\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\nn\\modules\\module.py:1501\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1496\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1497\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1498\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1499\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1500\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1501\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1502\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[0;32m   1503\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\nn\\modules\\loss.py:1174\u001b[0m, in \u001b[0;36mCrossEntropyLoss.forward\u001b[1;34m(self, input, target)\u001b[0m\n\u001b[0;32m   1173\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor, target: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[1;32m-> 1174\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcross_entropy\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweight\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1175\u001b[0m \u001b[43m                           \u001b[49m\u001b[43mignore_index\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mignore_index\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mreduction\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreduction\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1176\u001b[0m \u001b[43m                           \u001b[49m\u001b[43mlabel_smoothing\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlabel_smoothing\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\nn\\functional.py:3029\u001b[0m, in \u001b[0;36mcross_entropy\u001b[1;34m(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)\u001b[0m\n\u001b[0;32m   3027\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m size_average \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mor\u001b[39;00m reduce \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m   3028\u001b[0m     reduction \u001b[38;5;241m=\u001b[39m _Reduction\u001b[38;5;241m.\u001b[39mlegacy_get_string(size_average, reduce)\n\u001b[1;32m-> 3029\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_C\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_nn\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcross_entropy_loss\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweight\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_Reduction\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_enum\u001b[49m\u001b[43m(\u001b[49m\u001b[43mreduction\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mignore_index\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlabel_smoothing\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mValueError\u001b[0m: Expected input batch_size (970) to match target batch_size (960)."
     ]
    }
   ],
   "source": [
    "import mlflow\n",
    "\n",
    "\n",
    "EPOCHS=10\n",
    "\n",
    "mlflow.end_run()\n",
    "gradient_accumulation_steps = 10\n",
    "grad_counter = 0\n",
    "\n",
    "with mlflow.start_run(run_name='final') as run:\n",
    "\n",
    "    for epoch in range(EPOCHS):\n",
    "        for index, batch in enumerate(x_batches[0:5]): \n",
    "            outputs = llm_model(batch, labels=y_batches[index])\n",
    "            loss = outputs.loss\n",
    "            loss.backward()\n",
    "\n",
    "            if grad_counter % gradient_accumulation_steps == 0:\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                print(\"batch loss \", loss.item())\n",
    "\n",
    "                \n",
    "            grad_counter += 1\n",
    "                            \n",
    "\n",
    "        # log metric every epoch\n",
    "        mlflow.log_metric('loss', loss.item())\n",
    "        mlflow.log_metric('lr', get_lr(optimizer))\n",
    "\n",
    "        \n",
    "        output = llm_model.generate(tst_data[\"input_ids\"].to(device))\n",
    "        test_output = \"\"\n",
    "        for i in range(output.shape[0]):\n",
    "            tmp_str = tokenizer.decode(output[i])\n",
    "            test_output += \"\\n\" + tmp_str\n",
    "        mlflow.log_text(test_output, f\"test_{epoch}.txt\")\n",
    "        #### \n",
    "\n",
    "        \n",
    "        # Save the fine-tuned model every epoch \n",
    "        print(\"MODEL saved loss \", loss)\n",
    "        llm_model.save_pretrained(output_dir)\n",
    "        tokenizer.save_pretrained(output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "id": "c89cbc93-1cd3-463e-9bb3-472f3dbe7f1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
      "Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.\n",
      "A decoder-only architecture is being used, but right-padding was detected! For correct generation results, please set `padding_side='left'` when initializing the tokenizer.\n",
      "C:\\Users\\tempdelta\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\transformers\\generation\\utils.py:1136: UserWarning: Input length of input_ids is 41, but `max_length` is set to 20. This can lead to unexpected behavior. You should consider increasing `max_new_tokens`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input: <|endoftext|>  The subject of transcendental numbers has always fascinated mathe¬ maticians\n",
      "output <|endoftext|>  The subject of transcendental numbers has always fascinated mathe¬ maticians<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|> ) 2) Given five points, O, O', A, B, C, of an unknown conic K\n",
      "output <|endoftext|> ) 2) Given five points, O, O', A, B, C, of an unknown conic K<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  The number/is represented in the decimal system by the abbreviated symbol z\n",
      "output <|endoftext|>  The number/is represented in the decimal system by the abbreviated symbol z<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  To pass from the homogeneous coordinates z, y, z of an ordinary point P in the plane \n",
      "output <|endoftext|>  To pass from the homogeneous coordinates z, y, z of an ordinary point P in the plane <|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  around to Pi, the vector turns around and comes back to its original position\n",
      "output <|endoftext|>  around to Pi, the vector turns around and comes back to its original position<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  The construction of a6 is shown in Figure 30, where A D is a line parallel to BC through A \n",
      "output <|endoftext|>  The construction of a6 is shown in Figure 30, where A D is a line parallel to BC through A <|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  Then a point set S is of di¬ mension 0 if it is not of dimension —1 (i\n",
      "output <|endoftext|>  Then a point set S is of di¬ mension 0 if it is not of dimension —1 (i<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  only if either a; = 1, or the equation (15) x\"’* + x\"\"* + x\"-’ + \n",
      "output <|endoftext|>  only if either a; = 1, or the equation (15) x\"’* + x\"\"* + x\"-’ + <|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "\n",
      "input: <|endoftext|>  For centuries Euclid's set of axioms has been the object of intensive study\n",
      "output <|endoftext|>  For centuries Euclid's set of axioms has been the object of intensive study<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  The lines through AB, EG, and IF are the diagonals of the quadrilateral\n",
      "output <|endoftext|>  The lines through AB, EG, and IF are the diagonals of the quadrilateral<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  144 GEOMETRICAL CONSTRUCTIONS I ml Likewise B'Q = p, since o'/6 = V/a\n",
      "output <|endoftext|>  144 GEOMETRICAL CONSTRUCTIONS I ml Likewise B'Q = p, since o'/6 = V/a<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "\n",
      "input: <|endoftext|>  between two sides of a triangle is never greater than the third side\n",
      "output <|endoftext|>  between two sides of a triangle is never greater than the third side<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  Even for the simplest case this has proved to be a sizable task\n",
      "output <|endoftext|>  Even for the simplest case this has proved to be a sizable task<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  For example, the integral can be interpreted in mechanics as expressing the concept of work\n",
      "output <|endoftext|>  For example, the integral can be interpreted in mechanics as expressing the concept of work<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  Of all the regular polygons the hexagon is simplest to construct\n",
      "output <|endoftext|>  Of all the regular polygons the hexagon is simplest to construct<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  These pioneers were guided by a strong instinctive feeling that kept them from going far astray\n",
      "output <|endoftext|>  These pioneers were guided by a strong instinctive feeling that kept them from going far astray<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  By the following device we can obtain a more convenient expression\n",
      "output <|endoftext|>  By the following device we can obtain a more convenient expression<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  (Hint: Find X and X' from the conditions Xi(xi, yi) + X'Z'(xi, yi) = 0, X + X' = 1\n",
      "output <|endoftext|>  (Hint: Find X and X' from the conditions Xi(xi, yi) + X'Z'(xi, yi) = 0, X + X' = 1,\n",
      "input: <|endoftext|>  There is a definite psychological difliculty in grasping this precise definition of limit\n",
      "output <|endoftext|>  There is a definite psychological difliculty in grasping this precise definition of limit<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  (b) Write the equation expressing the fact that the distance (p\n",
      "output <|endoftext|>  (b) Write the equation expressing the fact that the distance (p<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  We see at once that a:x = x:y y:f, whence, if / is set equal to 2a in the apparatus, x* = 2a*\n",
      "output <|endoftext|>  We see at once that a:x = x:y y:f, whence, if / is set equal to 2a in the apparatus, x* = 2a*<|endoftext|><|endoftext|><|endoftext|>\n",
      "\n",
      "input: <|endoftext|>  In the interests of a uniform notation let us denote these figures by the symbols To, Ti, Fig\n",
      "output <|endoftext|>  In the interests of a uniform notation let us denote these figures by the symbols To, Ti, Fig<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  Let us suppose, for simplicity, that all the numbers are positive\n",
      "output <|endoftext|>  Let us suppose, for simplicity, that all the numbers are positive<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  Now take the pencil 0\" and transport it rigidly into any position O'\n",
      "output <|endoftext|>  Now take the pencil 0\" and transport it rigidly into any position O'<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  Hence we say only that /(x) ^^tends to infinity^^ as x approaches zero\n",
      "output <|endoftext|>  Hence we say only that /(x) ^^tends to infinity^^ as x approaches zero<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  The field F' contains the original field F, for we may, in particular, choose 6 = 0\n",
      "output <|endoftext|>  The field F' contains the original field F, for we may, in particular, choose 6 = 0<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  the difference between this value and the true value of e can easily be appraised\n",
      "output <|endoftext|>  the difference between this value and the true value of e can easily be appraised<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  We shall show the hypothesis that any one of the A’s is false to be untenable\n",
      "output <|endoftext|>  We shall show the hypothesis that any one of the A’s is false to be untenable<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "input: <|endoftext|>  Add¬ ing 1,456 + 5,630, we get 6 + 0 = 6 in the units place, 5 + 3 = 11 in the sevens place\n",
      "output <|endoftext|>  Add¬ ing 1,456 + 5,630, we get 6 + 0 = 6 in the units place, 5 + 3 = 11 in the sevens place<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n",
      "\n",
      "input: <|endoftext|>  This leads to the following diagram: a + b = (o + b)*=\n",
      "output <|endoftext|>  This leads to the following diagram: a + b = (o + b)*=<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|>\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "tst_text = np.random.choice(math_sentences, 30)\n",
    "tst_data = tokenizer.batch_encode_plus(tst_text, add_special_tokens=True, return_tensors='pt', padding=True).to(device)\n",
    "output = llm_model.generate(tst_data[\"input_ids\"])\n",
    "\n",
    "for i, text in enumerate(tst_text):\n",
    "    print(\"input:\", text)\n",
    "    print(\"output\", tokenizer.decode(output[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cbe6a06-f620-4ed7-9537-026702defb4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TextDataset\n",
    "\n",
    "train_dataset = TextDataset(\n",
    "      tokenizer=tokenizer,\n",
    "      file_path=math_example_path,\n",
    "      block_size=128)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "e6078a30-e05d-447f-9ca4-e0dcd4c38fb4",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Could not infer dtype of DatasetInfo",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[28], line 29\u001b[0m\n\u001b[0;32m      9\u001b[0m training_args \u001b[38;5;241m=\u001b[39m TrainingArguments(output_dir\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtest_trainer\u001b[39m\u001b[38;5;124m'\u001b[39m, \n\u001b[0;32m     10\u001b[0m                                   \u001b[38;5;66;03m#evaluation_strategy='epoch',\u001b[39;00m\n\u001b[0;32m     11\u001b[0m                                   per_device_train_batch_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     14\u001b[0m                                   num_train_epochs \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m2\u001b[39m,\n\u001b[0;32m     15\u001b[0m                                   fp16\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m,)\n\u001b[0;32m     17\u001b[0m trainer \u001b[38;5;241m=\u001b[39m Trainer(\n\u001b[0;32m     18\u001b[0m     model\u001b[38;5;241m=\u001b[39mllm_model,\n\u001b[0;32m     19\u001b[0m     args\u001b[38;5;241m=\u001b[39mtraining_args,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     26\u001b[0m     preprocess_logits_for_metrics\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[0;32m     27\u001b[0m )\n\u001b[1;32m---> 29\u001b[0m \u001b[43mtrainer\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\transformers\\trainer.py:1539\u001b[0m, in \u001b[0;36mTrainer.train\u001b[1;34m(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)\u001b[0m\n\u001b[0;32m   1537\u001b[0m         hf_hub_utils\u001b[38;5;241m.\u001b[39menable_progress_bars()\n\u001b[0;32m   1538\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1539\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43minner_training_loop\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   1540\u001b[0m \u001b[43m        \u001b[49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1541\u001b[0m \u001b[43m        \u001b[49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mresume_from_checkpoint\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1542\u001b[0m \u001b[43m        \u001b[49m\u001b[43mtrial\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mtrial\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1543\u001b[0m \u001b[43m        \u001b[49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mignore_keys_for_eval\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1544\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\transformers\\trainer.py:1836\u001b[0m, in \u001b[0;36mTrainer._inner_training_loop\u001b[1;34m(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)\u001b[0m\n\u001b[0;32m   1833\u001b[0m     rng_to_sync \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[0;32m   1835\u001b[0m step \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m\n\u001b[1;32m-> 1836\u001b[0m \u001b[43m\u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mstep\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minputs\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43menumerate\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mepoch_iterator\u001b[49m\u001b[43m)\u001b[49m\u001b[43m:\u001b[49m\n\u001b[0;32m   1837\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtotal_batched_samples\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\n\u001b[0;32m   1839\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43;01mif\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43margs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minclude_num_input_tokens_seen\u001b[49m\u001b[43m:\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\accelerate\\data_loader.py:451\u001b[0m, in \u001b[0;36mDataLoaderShard.__iter__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    449\u001b[0m \u001b[38;5;66;03m# We iterate one batch ahead to check when we are at the end\u001b[39;00m\n\u001b[0;32m    450\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m--> 451\u001b[0m     current_batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mnext\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mdataloader_iter\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    452\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m:\n\u001b[0;32m    453\u001b[0m     \u001b[38;5;28;01myield\u001b[39;00m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\utils\\data\\dataloader.py:634\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    631\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    632\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[0;32m    633\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[1;32m--> 634\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    635\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[0;32m    636\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[0;32m    637\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m \\\n\u001b[0;32m    638\u001b[0m         \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called:\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\utils\\data\\dataloader.py:678\u001b[0m, in \u001b[0;36m_SingleProcessDataLoaderIter._next_data\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    676\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_next_data\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m    677\u001b[0m     index \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_next_index()  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[1;32m--> 678\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_dataset_fetcher\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfetch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mindex\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[0;32m    679\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory:\n\u001b[0;32m    680\u001b[0m         data \u001b[38;5;241m=\u001b[39m _utils\u001b[38;5;241m.\u001b[39mpin_memory\u001b[38;5;241m.\u001b[39mpin_memory(data, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory_device)\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\torch\\utils\\data\\_utils\\fetch.py:54\u001b[0m, in \u001b[0;36m_MapDatasetFetcher.fetch\u001b[1;34m(self, possibly_batched_index)\u001b[0m\n\u001b[0;32m     52\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m     53\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n\u001b[1;32m---> 54\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcollate_fn\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\transformers\\trainer_utils.py:772\u001b[0m, in \u001b[0;36mRemoveColumnsCollator.__call__\u001b[1;34m(self, features)\u001b[0m\n\u001b[0;32m    770\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__call__\u001b[39m(\u001b[38;5;28mself\u001b[39m, features: List[\u001b[38;5;28mdict\u001b[39m]):\n\u001b[0;32m    771\u001b[0m     features \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_remove_columns(feature) \u001b[38;5;28;01mfor\u001b[39;00m feature \u001b[38;5;129;01min\u001b[39;00m features]\n\u001b[1;32m--> 772\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdata_collator\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfeatures\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\transformers\\data\\data_collator.py:92\u001b[0m, in \u001b[0;36mdefault_data_collator\u001b[1;34m(features, return_tensors)\u001b[0m\n\u001b[0;32m     86\u001b[0m \u001b[38;5;66;03m# In this function we'll make the assumption that all `features` in the batch\u001b[39;00m\n\u001b[0;32m     87\u001b[0m \u001b[38;5;66;03m# have the same attributes.\u001b[39;00m\n\u001b[0;32m     88\u001b[0m \u001b[38;5;66;03m# So we will look at the first element as a proxy for what attributes exist\u001b[39;00m\n\u001b[0;32m     89\u001b[0m \u001b[38;5;66;03m# on the whole batch.\u001b[39;00m\n\u001b[0;32m     91\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m return_tensors \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mpt\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m---> 92\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch_default_data_collator\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfeatures\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     93\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m return_tensors \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtf\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[0;32m     94\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m tf_default_data_collator(features)\n",
      "File \u001b[1;32m~\\AppData\\Local\\Packages\\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\\LocalCache\\local-packages\\Python311\\site-packages\\transformers\\data\\data_collator.py:158\u001b[0m, in \u001b[0;36mtorch_default_data_collator\u001b[1;34m(features)\u001b[0m\n\u001b[0;32m    156\u001b[0m             batch[k] \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor(np\u001b[38;5;241m.\u001b[39mstack([f[k] \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m features]))\n\u001b[0;32m    157\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m--> 158\u001b[0m             batch[k] \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtensor\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[43mf\u001b[49m\u001b[43m[\u001b[49m\u001b[43mk\u001b[49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mf\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mfeatures\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    160\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m batch\n",
      "\u001b[1;31mRuntimeError\u001b[0m: Could not infer dtype of DatasetInfo"
     ]
    }
   ],
   "source": [
    "from transformers import (\n",
    "    TrainingArguments, \n",
    "    Trainer, \n",
    "    default_data_collator, \n",
    ")\n",
    "import mlflow \n",
    "mlflow.end_run()\n",
    "\n",
    "training_args = TrainingArguments(output_dir='test_trainer', \n",
    "                                  #evaluation_strategy='epoch',\n",
    "                                  per_device_train_batch_size=1,\n",
    "                                  per_device_eval_batch_size=1,\n",
    "                                  gradient_accumulation_steps=20, #\n",
    "                                  num_train_epochs = 2,\n",
    "                                  fp16=False,)\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=llm_model,\n",
    "    args=training_args,\n",
    "    train_dataset=dataset,\n",
    "    #eval_dataset=dataset,\n",
    "    tokenizer=tokenizer,\n",
    "    # Data collator will default to DataCollatorWithPadding, so we change it.\n",
    "    data_collator=default_data_collator,\n",
    "    compute_metrics=None,\n",
    "    preprocess_logits_for_metrics=None,\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "027628bf-339b-41ad-992a-43f5145a5380",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Dataset({\n",
       "     features: ['text'],\n",
       "     num_rows: 31749\n",
       " })]"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e381e7ac-d923-47ab-98b5-b355409bf892",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}