pre-training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Nemanja Petrovic SIR 1 - SRBerta pre train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## OSCAR Dataset for Serbian language\n",
    "\n",
    "We need to download data from OSCAR corpus from Hugging face, after data is downloaded we are removing all new lines and divide data to chucnks of 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-05-31T06:56:13.252680400Z",
     "start_time": "2023-05-31T06:41:19.307567900Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading and preparing dataset oscar-2301/sr to C:/Users/HP1/Documents/Nemanja/SRBerta-pretrain/dataset_cache/oscar-corpus___oscar-2301/sr-language=sr/0.0.0/156efb8ba9f439f881d8f41fd7fddd5e04604bc27505c46ddef015f2fc551a4a...\n"
     ]
    },
    {
     "data": {
      "text/plain": "Downloading data:   0%|          | 0.00/450 [00:00<?, ?B/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "e8f6f34869fd408f925c4624e8f2b9a4"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Downloading data files:   0%|          | 0/5 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "85e9a24eb03c43da89d3514424449497"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Downloading data:   0%|          | 0.00/532M [00:00<?, ?B/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ad759d404ff14ec8af89361bf883b8c2"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Downloading data:   0%|          | 0.00/344M [00:00<?, ?B/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ba8314d6400a48fd8643ce814758c098"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Downloading data:   0%|          | 0.00/531M [00:00<?, ?B/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ada8adc9e2b24b11aab7d74f45b16a09"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Downloading data:   0%|          | 0.00/532M [00:00<?, ?B/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "ea5e54d309a64142af44b392abdd20cb"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Downloading data:   0%|          | 0.00/531M [00:00<?, ?B/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "fc341d2f0e6e48a7b60a5c68069cef1f"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Generating train split: 0 examples [00:00, ? examples/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "afc79f1b11504ea488eda696de4047b9"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset oscar-2301 downloaded and prepared to C:/Users/HP1/Documents/Nemanja/SRBerta-pretrain/dataset_cache/oscar-corpus___oscar-2301/sr-language=sr/0.0.0/156efb8ba9f439f881d8f41fd7fddd5e04604bc27505c46ddef015f2fc551a4a. Subsequent calls will reuse this data.\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/1 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c265fb7642f746259d5c7535d0202cb4"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['id', 'text', 'meta'],\n",
      "        num_rows: 838948\n",
      "    })\n",
      "})\n",
      "STARTED WRITING DATA TO FILES\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/838948 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "a48346a0e15744fba3d233679b6d64eb"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FINISHED WRITING DATA TO FILES\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "hugging_face_token = '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'\n",
    "dataset = load_dataset(\"oscar-corpus/OSCAR-2301\",\n",
    "                       cache_dir=\"dataset_cache\",\n",
    "                       use_auth_token=hugging_face_token,\n",
    "                       language=\"sr\",\n",
    "                       streaming=False)\n",
    "\n",
    "print(dataset)\n",
    "\n",
    "# Format everything and pu in files with length 10000\n",
    "print(\"STARTED WRITING DATA TO FILES\")\n",
    "from tqdm.auto import tqdm\n",
    "text_data = []\n",
    "file_count = 0\n",
    "\n",
    "for sample in tqdm(dataset['train']):\n",
    "\n",
    "    sample = sample['text'].replace('\\n', '')\n",
    "    text_data.append(sample)\n",
    "\n",
    "    if len(text_data) == 10_000:\n",
    "        with open(f'sr_{file_count}.txt', 'w', encoding='utf-8') as fp:\n",
    "            fp.write('\\n'.join(text_data))\n",
    "        text_data = []\n",
    "        file_count += 1\n",
    "\n",
    "# after saving in 10K chunks, we have to add leftovers\n",
    "with open(f'sr_{file_count}.txt', 'w', encoding='utf-8') as fp:\n",
    "    fp.write('\\n'.join(text_data))\n",
    "\n",
    "print(\"FINISHED WRITING DATA TO FILES\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenizer\n",
    "\n",
    "Before training model, we need to train tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-07T08:18:34.668983800Z",
     "start_time": "2023-06-07T07:45:12.081756200Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING TOKENIZER TRAINING\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'BertTokenizer'. \n",
      "The class this function is called from is 'RobertaTokenizer'.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FINISHED TOKENIZER TRAINING\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'BertTokenizer'. \n",
      "The class this function is called from is 'RobertaTokenizerFast'.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of input ids in sample:\n",
      "torch.Size([1, 10])\n",
      "Testing decoder\n"
     ]
    },
    {
     "data": {
      "text/plain": "' дан'"
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import os\n",
    "from transformers import RobertaTokenizerFast\n",
    "from tokenizers.decoders import ByteLevel\n",
    "\n",
    "paths = [str(x) for x in Path('./').glob('*.txt')]\n",
    "# For testing, taking only first 5 files, for real training remove this and go through more data\n",
    "paths = paths[0:40]\n",
    "\n",
    "from tokenizers import ByteLevelBPETokenizer\n",
    "\n",
    "print(\"STARTING TOKENIZER TRAINING\")\n",
    "tokenizer = ByteLevelBPETokenizer()\n",
    "tokenizer.train(\n",
    "    files=paths,\n",
    "    vocab_size=50265,\n",
    "    min_frequency=2,show_progress=True,\n",
    "    special_tokens=[\n",
    "        '<s>', '<pad>', '</s>', '<unk>', '<mask>'\n",
    "    ]\n",
    ")\n",
    "print(\"FINISHED TOKENIZER TRAINING\")\n",
    "\n",
    "os.mkdir('./srberta_tokenizer')\n",
    "tokenizer.save_model('srberta_tokenizer')\n",
    "srberta_tokenizer = RobertaTokenizerFast.from_pretrained(\"srberta_tokenizer\")\n",
    "\n",
    "sample = srberta_tokenizer(\"Добар дан, како си данас ти човече\", return_tensors='pt')\n",
    "print(\"Shape of input ids in sample:\")\n",
    "print(str(sample.input_ids.shape))\n",
    "\n",
    "# Test decoder\n",
    "print(\"Testing decoder\")\n",
    "decoder = ByteLevel()\n",
    "decoder.decode('ĠÐ´Ð°Ð½')"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Input pipeline"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-06-07T11:00:15.842855700Z",
     "start_time": "2023-06-07T09:13:55.002676200Z"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'BertTokenizer'. \n",
      "The class this function is called from is 'RobertaTokenizer'.\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'BertTokenizer'. \n",
      "The class this function is called from is 'RobertaTokenizerFast'.\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/74 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8a26a954bef14b1492ee40ab7fcf7e37"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "tensor([    0,  6909,  7665,  3528,  1602,   912,  5051, 19599,   365,   480])"
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from pathlib import Path\n",
    "from transformers import RobertaTokenizerFast\n",
    "from tqdm.auto import tqdm\n",
    "import os\n",
    "\n",
    "def mlm(tensor):\n",
    "\n",
    "    rand = torch.rand(tensor.shape) #[0,1]\n",
    "    mask_arr = (rand < 0.15)* (tensor!=0)* (tensor!=1)* (tensor!=2)\n",
    "    for i in range(tensor.shape[0]):\n",
    "        selection = torch.flatten(mask_arr[i].nonzero()).tolist()\n",
    "        tensor[i, selection] = 4\n",
    "\n",
    "    return tensor\n",
    "\n",
    "paths = [str(x) for x in Path('./').glob('*.txt')]\n",
    "\n",
    "tokenizer_srberta = RobertaTokenizerFast.from_pretrained(\"srberta_tokenizer\")\n",
    "\n",
    "input_ids = []\n",
    "mask = [] # attention mask\n",
    "labels = []\n",
    "\n",
    "for path in tqdm(paths):\n",
    "    with open(path, 'r', encoding='utf-8') as f:\n",
    "        lines = f.read().split('\\n')\n",
    "\n",
    "    sample = tokenizer_srberta(lines, max_length=512, padding='max_length', truncation=True, return_tensors='pt')\n",
    "\n",
    "    labels.append(sample.input_ids)\n",
    "    mask.append(sample.attention_mask)\n",
    "    input_ids.append(mlm(sample.input_ids.detach().clone()))\n",
    "\n",
    "# sample['input_ids'].shape\n",
    "\n",
    "input_ids = torch.cat(input_ids)\n",
    "mask = torch.cat(mask)\n",
    "labels = torch.cat(labels)\n",
    "torch.save(input_ids, 'input_ids.pt')\n",
    "torch.save(mask, 'mask.pt')\n",
    "torch.save(labels, 'labels.pt')\n",
    "\n",
    "input_ids = torch.load(\"input_ids.pt\")\n",
    "mask = torch.load(\"mask.pt\")\n",
    "labels = torch.load(\"labels.pt\")\n",
    "input_ids[0][:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "data": {
      "text/plain": "(7, 5)"
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "torch.cuda.get_device_capability()"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-06-07T11:32:30.163091Z",
     "start_time": "2023-06-07T11:32:30.039641300Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Dataloader"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "{'input_ids': tensor([[    0, 46776,  6994,  ...,   966,    16,     2],\n",
      "        [    0,  1165,  4644,  ...,   800,   786,     2],\n",
      "        [    0,   424,     4,  ...,   280,  2126,     2],\n",
      "        ...,\n",
      "        [    0,    42, 40390,  ...,     1,     1,     1],\n",
      "        [    0,   424,  1679,  ...,     1,     1,     1],\n",
      "        [    0, 23701, 19769,  ..., 35977,  1040,     2]]), 'attention_mask': tensor([[1, 1, 1,  ..., 1, 1, 1],\n",
      "        [1, 1, 1,  ..., 1, 1, 1],\n",
      "        [1, 1, 1,  ..., 1, 1, 1],\n",
      "        ...,\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 0, 0, 0],\n",
      "        [1, 1, 1,  ..., 1, 1, 1]]), 'labels': tensor([[    0, 46776,  6994,  ...,   966,    16,     2],\n",
      "        [    0,  1165,  4644,  ...,   800,   786,     2],\n",
      "        [    0,   424, 25258,  ...,   280,  2126,     2],\n",
      "        ...,\n",
      "        [    0,    42, 40390,  ...,     1,     1,     1],\n",
      "        [    0,   424,  1679,  ...,     1,     1,     1],\n",
      "        [    0, 23701, 19769,  ..., 35977,  1040,     2]])}\n",
      "738948\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "input_ids = torch.load(\"input_ids.pt\")\n",
    "mask = torch.load(\"mask.pt\")\n",
    "labels = torch.load(\"labels.pt\")\n",
    "\n",
    "encodings = {\n",
    "    'input_ids': input_ids,\n",
    "    'attention_mask': mask,\n",
    "    'labels': labels\n",
    "}\n",
    "\n",
    "class Dataset(torch.utils.data.Dataset):\n",
    "\n",
    "    def __init__(self, encodings):\n",
    "        self.encodings = encodings\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.encodings['input_ids'].shape[0]\n",
    "\n",
    "    def __getitem__(self, i):\n",
    "        return {key: tensor[i] for key, tensor in self.encodings.items()}\n",
    "\n",
    "dataset = Dataset(encodings)\n",
    "BATCH_SIZE = 16\n",
    "DO_SHUFFLE = True\n",
    "dataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=DO_SHUFFLE)\n",
    "for i, data in enumerate(dataloader, 0):\n",
    "    print(i)\n",
    "    print(data)\n",
    "    break\n",
    "\n",
    "print(len(dataloader.dataset))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-06-07T11:33:02.897851600Z",
     "start_time": "2023-06-07T11:32:56.432957100Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Training"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'BertTokenizer'. \n",
      "The class this function is called from is 'RobertaTokenizer'.\n",
      "The tokenizer class you load from this checkpoint is not the same type as the class this function is called from. It may result in unexpected tokenization. \n",
      "The tokenizer class you load from this checkpoint is 'BertTokenizer'. \n",
      "The class this function is called from is 'RobertaTokenizerFast'.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\HP1\\.conda\\envs\\conda-env-sir1-windows\\lib\\site-packages\\transformers\\optimization.py:306: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": "  0%|          | 0/61579 [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "c6e3b50abf544855b0380b64d3b9cfb6"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import RobertaConfig\n",
    "from transformers import RobertaTokenizerFast\n",
    "from transformers import RobertaForMaskedLM\n",
    "import torch\n",
    "from transformers import AdamW\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "tokenizer_srberta = RobertaTokenizerFast.from_pretrained(\"srberta_tokenizer\")\n",
    "\n",
    "config = RobertaConfig(\n",
    "    vocab_size=tokenizer_srberta.vocab_size,\n",
    "    max_position_embeddings=514,\n",
    "    hidden_size=768,\n",
    "    num_attention_heads=12,\n",
    "    num_hidden_layers=12,\n",
    "    type_vocab_size=1,\n",
    ")\n",
    "\n",
    "model = RobertaForMaskedLM(config) #randomly initialized weights\n",
    "\n",
    "torch.cuda.empty_cache()\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device.cpu()\n",
    "print(str(device))\n",
    "model.to(device)\n",
    "model.train()\n",
    "\n",
    "optim = AdamW(model.parameters(), lr=1e-5)\n",
    "epochs=25\n",
    "\n",
    "writer = SummaryWriter(\"./runs_v2\")\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    step=0\n",
    "    # setup loop with TQDM and dataloader\n",
    "    loop = tqdm(dataloader, leave=True)\n",
    "\n",
    "    for batch in loop:\n",
    "\n",
    "        optim.zero_grad()\n",
    "\n",
    "        input_ids = batch['input_ids'].to(device)\n",
    "        mask = batch['attention_mask'].to(device)\n",
    "        labels = batch['labels'].to(device)\n",
    "\n",
    "        outputs = model(input_ids, attention_mask=mask, labels=labels)\n",
    "        loss = outputs.loss\n",
    "\n",
    "        loss.backward()\n",
    "        optim.step()\n",
    "\n",
    "        loop.set_description(f'Epoch: {epoch}')\n",
    "        loop.set_postfix(loss=loss.item())\n",
    "\n",
    "        writer.add_scalar(\"Loss/train\", loss, step)\n",
    "        writer.flush()\n",
    "\n",
    "        if step % 25_000 == 0:\n",
    "            torch.save({'optimizer_state_dict': optim.state_dict()}, str(step)+'_'+ str(epoch)+'_optimizer.pt')\n",
    "            model.save_pretrained(\"./srberta_model_\"+str(step)+'_'+ str(epoch))\n",
    "\n",
    "        step+=1\n",
    "\n",
    "    # Save after each epoch\n",
    "    torch.save({'optimizer_state_dict': optim.state_dict()}, str(epoch)+'_optimizer.pt')\n",
    "    model.save_pretrained(\"./srberta_model_\"+ str(epoch))"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-06-21T04:53:02.183548300Z",
     "start_time": "2023-06-07T11:33:13.068755700Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'optim' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[1;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[1;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[1;32mIn[2], line 3\u001B[0m\n\u001B[0;32m      1\u001B[0m \u001B[38;5;28;01mimport\u001B[39;00m \u001B[38;5;21;01mtorch\u001B[39;00m\n\u001B[0;32m      2\u001B[0m torch\u001B[38;5;241m.\u001B[39msave({\n\u001B[1;32m----> 3\u001B[0m     \u001B[38;5;124m'\u001B[39m\u001B[38;5;124moptimizer_state_dict\u001B[39m\u001B[38;5;124m'\u001B[39m: \u001B[43moptim\u001B[49m\u001B[38;5;241m.\u001B[39mstate_dict()\n\u001B[0;32m      4\u001B[0m },\u001B[38;5;124m'\u001B[39m\u001B[38;5;124moptimizer_3_epochs.pt\u001B[39m\u001B[38;5;124m'\u001B[39m)\n\u001B[0;32m      6\u001B[0m model\u001B[38;5;241m.\u001B[39msave_pretrained(\u001B[38;5;124m\"\u001B[39m\u001B[38;5;124m./srberta_model_3_epochs\u001B[39m\u001B[38;5;124m\"\u001B[39m)\n\u001B[0;32m      8\u001B[0m \u001B[38;5;28;01mdef\u001B[39;00m \u001B[38;5;21msave\u001B[39m(model, optimizer):\n\u001B[0;32m      9\u001B[0m     \u001B[38;5;66;03m# save\u001B[39;00m\n",
      "\u001B[1;31mNameError\u001B[0m: name 'optim' is not defined"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "torch.save({\n",
    "    'optimizer_state_dict': optim.state_dict()\n",
    "},'optimizer_3_epochs.pt')\n",
    "\n",
    "model.save_pretrained(\"./srberta_model_3_epochs\")\n",
    "\n",
    "def save(model, optimizer):\n",
    "    # save\n",
    "    torch.save({\n",
    "        'model_state_dict': model.state_dict(),\n",
    "        'optimizer_state_dict': optimizer.state_dict()\n",
    "    }, 'output_model.pt')\n",
    "\n",
    "save(model, optim)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2023-05-10T06:43:55.415597200Z",
     "start_time": "2023-05-10T06:42:39.087569500Z"
    }
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Testing model"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "outputs": [
    {
     "data": {
      "text/plain": "[{'score': 0.01705239899456501,\n  'token': 18,\n  'token_str': '.',\n  'sequence': 'Добар дан како. '},\n {'score': 0.0163358673453331,\n  'token': 341,\n  'token_str': ' на',\n  'sequence': 'Добар дан како на '},\n {'score': 0.01422953512519598,\n  'token': 316,\n  'token_str': ' је',\n  'sequence': 'Добар дан како је '},\n {'score': 0.013286540284752846,\n  'token': 16,\n  'token_str': ',',\n  'sequence': 'Добар дан како, '},\n {'score': 0.011648965999484062,\n  'token': 280,\n  'token_str': ' и',\n  'sequence': 'Добар дан како и '}]"
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import pipeline\n",
    "from transformers import RobertaForMaskedLM\n",
    "\n",
    "model_v2 = RobertaForMaskedLM.from_pretrained(\"./srberta_model_1\")\n",
    "model_v2.to('cpu')\n",
    "\n",
    "fill = pipeline('fill-mask', model=model_v2, tokenizer=tokenizer_srberta)\n",
    "fill(f'Добар дан како {fill.tokenizer.mask_token} ')"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "start_time": "2023-04-24T22:42:07.831333Z",
     "end_time": "2023-04-24T22:42:08.495064Z"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}