transformer-text-generation.ipynb

{"metadata":{"colab":{"provenance":[],"gpuType":"V28"},"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"accelerator":"TPU","kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":1457,"sourceType":"datasetVersion","datasetId":781}],"dockerImageVersionId":30674,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import torch\nfrom transformers import GPT2LMHeadModel, GPT2Tokenizer\nimport logging\nlogging.getLogger().setLevel(logging.CRITICAL)\nimport torch\nimport numpy as np\n","metadata":{"id":"ntR1DiHtYW0I","execution":{"iopub.status.busy":"2024-04-07T22:56:45.284569Z","iopub.execute_input":"2024-04-07T22:56:45.285212Z","iopub.status.idle":"2024-04-07T22:56:53.561759Z","shell.execute_reply.started":"2024-04-07T22:56:45.285178Z","shell.execute_reply":"2024-04-07T22:56:53.560885Z"},"trusted":true},"execution_count":2,"outputs":[]},{"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndevice","metadata":{"id":"y1hiXsOMZffl","execution":{"iopub.status.busy":"2024-04-07T22:56:53.563283Z","iopub.execute_input":"2024-04-07T22:56:53.563705Z","iopub.status.idle":"2024-04-07T22:56:53.619856Z","shell.execute_reply.started":"2024-04-07T22:56:53.563666Z","shell.execute_reply":"2024-04-07T22:56:53.618884Z"},"trusted":true},"execution_count":3,"outputs":[{"execution_count":3,"output_type":"execute_result","data":{"text/plain":"device(type='cuda')"},"metadata":{}}]},{"cell_type":"code","source":"tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')\nmodel = GPT2LMHeadModel.from_pretrained('gpt2-medium')","metadata":{"id":"4OCNIadbYoIu","execution":{"iopub.status.busy":"2024-04-07T22:56:53.621342Z","iopub.execute_input":"2024-04-07T22:56:53.621680Z","iopub.status.idle":"2024-04-07T22:57:04.390184Z","shell.execute_reply.started":"2024-04-07T22:56:53.621652Z","shell.execute_reply":"2024-04-07T22:57:04.389454Z"},"trusted":true},"execution_count":4,"outputs":[{"output_type":"display_data","data":{"text/plain":"tokenizer_config.json:   0%|          | 0.00/26.0 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"82cd220a57374fb283b18d24498f9510"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"vocab.json:   0%|          | 0.00/1.04M [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"753d0ceb64074adc97e2c1ea4e594877"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"merges.txt:   0%|          | 0.00/456k [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"61b623b13c2740f397b325eb0a49f53c"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"tokenizer.json:   0%|          | 0.00/1.36M [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"5066019c437548848116d18831166f96"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"config.json:   0%|          | 0.00/718 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"ee12f1676dd24dc1b66049230452d3bb"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"model.safetensors:   0%|          | 0.00/1.52G [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"b5c62acdd89f4f5e9e44d23fe19e6279"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"generation_config.json:   0%|          | 0.00/124 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"a7b6147ebcfc448e9962b84d64b6c357"}},"metadata":{}}]},{"cell_type":"code","source":"model = model.to(device)","metadata":{"id":"g8gvFOL5bVi_","execution":{"iopub.status.busy":"2024-04-07T22:57:04.392346Z","iopub.execute_input":"2024-04-07T22:57:04.392623Z","iopub.status.idle":"2024-04-07T22:57:04.970991Z","shell.execute_reply.started":"2024-04-07T22:57:04.392599Z","shell.execute_reply":"2024-04-07T22:57:04.969691Z"},"trusted":true},"execution_count":5,"outputs":[]},{"cell_type":"code","source":"def choose_from_top(probs, n=5):\n    ind = np.argpartition(probs, -n)[-n:]\n    top_prob = probs[ind]\n    top_prob = top_prob / np.sum(top_prob) # Normalize\n    choice = np.random.choice(n, 1, p = top_prob)\n    token_id = ind[choice][0]\n    return int(token_id)","metadata":{"id":"Fe7wifbKbiDF","execution":{"iopub.status.busy":"2024-04-07T22:57:04.972820Z","iopub.execute_input":"2024-04-07T22:57:04.973544Z","iopub.status.idle":"2024-04-07T22:57:04.981057Z","shell.execute_reply.started":"2024-04-07T22:57:04.973507Z","shell.execute_reply":"2024-04-07T22:57:04.980089Z"},"trusted":true},"execution_count":6,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import Dataset\nfrom torch.utils.data import Dataset, DataLoader\nimport os\nimport json\nimport csv\n\nclass JokesDataset(Dataset):\n    def __init__(self, jokes_dataset_path = '/kaggle/input/short-jokes'):\n        super().__init__()\n\n        short_jokes_path = os.path.join(jokes_dataset_path, 'shortjokes.csv')\n\n        self.joke_list = []\n        self.end_of_text_token = \"<|endoftext|>\"\n\n        with open(short_jokes_path) as csv_file:\n            csv_reader = csv.reader(csv_file, delimiter=',')\n\n            x = 0\n            for row in csv_reader:\n                joke_str = f\"JOKE:{row[1]}{self.end_of_text_token}\"\n                self.joke_list.append(joke_str)\n\n    def __len__(self):\n        return len(self.joke_list)\n\n    def __getitem__(self, item):\n        return self.joke_list[item]","metadata":{"id":"ZWlvAArBeFhi","execution":{"iopub.status.busy":"2024-04-07T22:57:04.984630Z","iopub.execute_input":"2024-04-07T22:57:04.984975Z","iopub.status.idle":"2024-04-07T22:57:05.687258Z","shell.execute_reply.started":"2024-04-07T22:57:04.984932Z","shell.execute_reply":"2024-04-07T22:57:05.686017Z"},"trusted":true},"execution_count":7,"outputs":[]},{"cell_type":"code","source":"dataset = JokesDataset()\njoke_loader = DataLoader(dataset, batch_size=1, shuffle=True)","metadata":{"id":"DD8pXFj3cHFX","execution":{"iopub.status.busy":"2024-04-07T22:57:05.688559Z","iopub.execute_input":"2024-04-07T22:57:05.689343Z","iopub.status.idle":"2024-04-07T22:57:06.394397Z","shell.execute_reply.started":"2024-04-07T22:57:05.689313Z","shell.execute_reply":"2024-04-07T22:57:06.393230Z"},"trusted":true},"execution_count":8,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 16\nEPOCHS = 5\nLEARNING_RATE = 3e-5\nWARMUP_STEPS = 5000\nMAX_SEQ_LEN = 400\n# from transformers import AdamW, WarmupLinearSchedule\nfrom transformers import get_linear_schedule_with_warmup\nfrom torch.optim import AdamW\n\ndevice = 'cpu'kn\nif torch.cuda.is_available():\n    device = 'cuda'","metadata":{"id":"eOT1zLOQea8f","execution":{"iopub.status.busy":"2024-04-07T22:57:06.395742Z","iopub.execute_input":"2024-04-07T22:57:06.396054Z","iopub.status.idle":"2024-04-07T22:57:06.416595Z","shell.execute_reply.started":"2024-04-07T22:57:06.396028Z","shell.execute_reply":"2024-04-07T22:57:06.415503Z"},"trusted":true},"execution_count":9,"outputs":[]},{"cell_type":"code","source":"model = model.to(device)\nmodel.train()\noptimizer = AdamW(model.parameters(), lr=LEARNING_RATE)\nscheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=WARMUP_STEPS, num_training_steps=-1)\nproc_seq_count = 0\nsum_loss = 0.0\nbatch_count = 0\n\ntmp_jokes_tens = None\nmodels_folder = \"trained_models\"\nif not os.path.exists(models_folder):\n    os.mkdir(models_folder)\n\nfor epoch in range(EPOCHS):\n\n    print(f\"EPOCH {epoch} started\" + '=' * 30)\n\n    for idx,joke in enumerate(joke_loader):\n\n        #################### \"Fit as many joke sequences into MAX_SEQ_LEN sequence as possible\" logic start ####\n        joke_tens = torch.tensor(tokenizer.encode(joke[0])).unsqueeze(0).to(device)\n        #Skip sample from dataset if it is longer than MAX_SEQ_LEN\n        if joke_tens.size()[1] > MAX_SEQ_LEN:\n            continue\n\n        #The first joke sequence in the sequence\n        if not torch.is_tensor(tmp_jokes_tens):\n            tmp_jokes_tens = joke_tens\n            continue\n        else:\n            #The next joke does not fit in so we process the sequence and leave the last joke\n            #as the start for next sequence\n            if tmp_jokes_tens.size()[1] + joke_tens.size()[1] > MAX_SEQ_LEN:\n                work_jokes_tens = tmp_jokes_tens\n                tmp_jokes_tens = joke_tens\n            else:\n                #Add the joke to sequence, continue and try to add more\n                tmp_jokes_tens = torch.cat([tmp_jokes_tens, joke_tens[:,1:]], dim=1)\n                continue\n        ################## Sequence ready, process it trough the model ##################\n\n        outputs = model(work_jokes_tens, labels=work_jokes_tens)\n        loss, logits = outputs[:2]\n        loss.backward()\n        sum_loss = sum_loss + loss.detach().data\n\n        proc_seq_count = proc_seq_count + 1\n        if proc_seq_count == BATCH_SIZE:\n            proc_seq_count = 0\n            batch_count += 1\n            optimizer.step()\n            scheduler.step()\n            optimizer.zero_grad()\n            model.zero_grad()\n\n        if batch_count == 100:\n            print(f\"sum loss {sum_loss}\")\n            batch_count = 0\n            sum_loss = 0.0\n\n    # Store the model after each epoch to compare the performance of them\n    torch.save(model.state_dict(), os.path.join(models_folder, f\"gpt2_medium_joker_{epoch}.pt\"))\n\n","metadata":{"id":"CxCatiaOeevq","colab":{"base_uri":"https://localhost:8080/"},"outputId":"c883ee95-24b2-4a42-ba0f-d0a415eea149","execution":{"iopub.status.busy":"2024-04-07T22:57:06.418099Z","iopub.execute_input":"2024-04-07T22:57:06.418858Z"},"trusted":true},"execution_count":null,"outputs":[{"name":"stdout","text":"EPOCH 0 started==============================\nsum loss 6755.1103515625\nsum loss 6148.828125\nsum loss 5617.81396484375\nsum loss 5415.01904296875\nsum loss 5352.4765625\nsum loss 5285.9462890625\nsum loss 5243.9482421875\nsum loss 5205.78271484375\nsum loss 5177.14306640625\nEPOCH 1 started==============================\nsum loss 5137.51513671875\nsum loss 5107.45166015625\nsum loss 5090.19873046875\nsum loss 5065.3984375\nsum loss 5039.5302734375\nsum loss 5020.61474609375\nsum loss 5010.93359375\nsum loss 4993.40478515625\nsum loss 4976.2294921875\nsum loss 4956.5556640625\nEPOCH 2 started==============================\nsum loss 4941.65380859375\nsum loss 4902.33837890625\nsum loss 4894.14990234375\nsum loss 4871.5986328125\nsum loss 4868.9609375\nsum loss 4859.55859375\nsum loss 4843.63623046875\nsum loss 4828.8525390625\nsum loss 4814.4453125\nsum loss 4825.3916015625\nEPOCH 3 started==============================\nsum loss 4767.33203125\nsum loss 4734.7919921875\nsum loss 4728.88134765625\nsum loss 4723.919921875\nsum loss 4722.154296875\nsum loss 4708.14501953125\nsum loss 4708.04736328125\n","output_type":"stream"}]},{"cell_type":"code","source":"MODEL_EPOCH = 4\n\nmodels_folder = \"trained_models\"\n\nmodel_path = os.path.join(models_folder, f\"gpt2_medium_joker_{MODEL_EPOCH}.pt\")\nmodel.load_state_dict(torch.load(model_path))\n\njokes_output_file_path = f'generated_{MODEL_EPOCH}.jokes'\n\nmodel.eval()\nif os.path.exists(jokes_output_file_path):\n    os.remove(jokes_output_file_path)\n\njoke_num = 0\nwith torch.no_grad():\n\n        for joke_idx in range(1000):\n\n            joke_finished = False\n\n            cur_ids = torch.tensor(tokenizer.encode(\"JOKE:\")).unsqueeze(0).to(device)\n\n            for i in range(100):\n                outputs = model(cur_ids, labels=cur_ids)\n                loss, logits = outputs[:2]\n                softmax_logits = torch.softmax(logits[0,-1], dim=0) #Take the first(from only one in this case) batch and the last predicted embedding\n                if i < 3:\n                    n = 20\n                else:\n                    n = 3\n                next_token_id = choose_from_top(softmax_logits.to('cpu').numpy(), n=n) #Randomly(from the topN probability distribution) select the next word\n                cur_ids = torch.cat([cur_ids, torch.ones((1,1)).long().to(device) * next_token_id], dim = 1) # Add the last word to the running sequence\n\n                if next_token_id in tokenizer.encode('<|endoftext|>'):\n                    joke_finished = True\n                    break\n\n\n            if joke_finished:\n\n                joke_num = joke_num + 1\n\n                output_list = list(cur_ids.squeeze().to('cpu').numpy())\n                output_text = tokenizer.decode(output_list)\n\n                with open(jokes_output_file_path, 'a') as f:\n                    f.write(f\"{output_text} \\n\\n\")\n\n","metadata":{"id":"ev1fBTbYewDM","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{"id":"QIYAvp0QkfLO"},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}