gpt.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5b21d2fa-a70a-46c7-8fec-762e613994e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import tiktoken_ext.openai_public\n",
    "from src.model import *\n",
    "from src.dataset import *\n",
    "from src.generate import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bed8f474-13ba-42cc-aa4b-def5caf78dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "37750c62-c36d-4eea-bc40-f4ca1def892b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7fa41bf97750>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = dict(\n",
    "    batch_size = 64, # N\n",
    "    sequence_dim = 100, # L, S\n",
    "    embed_dim = 78, # E\n",
    "    num_heads = 13, # H\n",
    "    num_layers = 4,\n",
    "    dropout = 0.2,\n",
    "    train_steps = 10000,\n",
    "    lr = 1e-3, # learning rate\n",
    "    seed = 78,\n",
    "    device = 'cuda',\n",
    ")\n",
    "assert config['embed_dim'] % config['num_heads'] == 0\n",
    "torch.manual_seed(config['seed'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3682ba70-72b1-4a84-bd59-bf4fe996cd8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./data/shakespeare.txt', 'r', encoding='utf-8') as f:\n",
    "    text = f.read()\n",
    "\n",
    "# dataset_shakespeare = CharacterDataset(text, seq_len=config['sequence_dim']) # n_vocab = 65\n",
    "\n",
    "tiktoken_config = tiktoken_ext.openai_public.gpt2()\n",
    "dataset_shakespeare = WordDataset(text, seq_len=config['sequence_dim'], tiktoken_config=tiktoken_config) # n_vocab = 50K, requires bigger parameters\n",
    "os.makedirs('./model_artifacts', exist_ok=True)\n",
    "with open('./model_artifacts/tiktoken_config.pkl', 'wb') as f:\n",
    "    pickle.dump(tiktoken_config, f) # need this for offline containers\n",
    "\n",
    "# flavor 1 - shuffled split\n",
    "# data_train, data_test = torch.utils.data.random_split(dataset_shakespeare, [.9, .1])\n",
    "\n",
    "# flavor 2 - non-shuffled split\n",
    "n = int(.95*len(dataset_shakespeare))\n",
    "dataset_train = torch.utils.data.Subset(dataset_shakespeare, list(range(0, n)))\n",
    "dataset_val = torch.utils.data.Subset(dataset_shakespeare, list(range(n, len(dataset_shakespeare))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d5b0e80f-829d-4cb3-b944-2860cb9a7a5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = GPT(\n",
    "    dataset_shakespeare.vocab_dim,\n",
    "    config['sequence_dim'],\n",
    "    config['embed_dim'],\n",
    "    config['num_heads'],\n",
    "    config['num_layers'],\n",
    "    dropout=config['dropout'],\n",
    "    device=config['device'],\n",
    ")\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=config['lr'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "90e66b27-247d-4e36-ae45-c8b14fa71852",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "8193457\n",
      "GPT(\n",
      "  (token_embedding): Embedding(50257, 78)\n",
      "  (position_embedding): Embedding(100, 78)\n",
      "  (dropout): Dropout(p=0.2, inplace=False)\n",
      "  (blocks): Sequential(\n",
      "    (0): SelfAttentionBlock(\n",
      "      (ln1): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "      (mha): MultiheadAttention(\n",
      "        (query): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (key): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (value): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (dropout1): Dropout(p=0.2, inplace=False)\n",
      "        (projection): Linear(in_features=78, out_features=78, bias=True)\n",
      "      )\n",
      "      (ln2): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): MLP(\n",
      "        (net): Sequential(\n",
      "          (0): Linear(in_features=78, out_features=312, bias=True)\n",
      "          (1): GELU(approximate='none')\n",
      "          (2): Linear(in_features=312, out_features=78, bias=True)\n",
      "          (3): Dropout(p=0.2, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (1): SelfAttentionBlock(\n",
      "      (ln1): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "      (mha): MultiheadAttention(\n",
      "        (query): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (key): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (value): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (dropout1): Dropout(p=0.2, inplace=False)\n",
      "        (projection): Linear(in_features=78, out_features=78, bias=True)\n",
      "      )\n",
      "      (ln2): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): MLP(\n",
      "        (net): Sequential(\n",
      "          (0): Linear(in_features=78, out_features=312, bias=True)\n",
      "          (1): GELU(approximate='none')\n",
      "          (2): Linear(in_features=312, out_features=78, bias=True)\n",
      "          (3): Dropout(p=0.2, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (2): SelfAttentionBlock(\n",
      "      (ln1): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "      (mha): MultiheadAttention(\n",
      "        (query): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (key): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (value): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (dropout1): Dropout(p=0.2, inplace=False)\n",
      "        (projection): Linear(in_features=78, out_features=78, bias=True)\n",
      "      )\n",
      "      (ln2): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): MLP(\n",
      "        (net): Sequential(\n",
      "          (0): Linear(in_features=78, out_features=312, bias=True)\n",
      "          (1): GELU(approximate='none')\n",
      "          (2): Linear(in_features=312, out_features=78, bias=True)\n",
      "          (3): Dropout(p=0.2, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "    (3): SelfAttentionBlock(\n",
      "      (ln1): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "      (mha): MultiheadAttention(\n",
      "        (query): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (key): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (value): Linear(in_features=78, out_features=78, bias=False)\n",
      "        (dropout1): Dropout(p=0.2, inplace=False)\n",
      "        (projection): Linear(in_features=78, out_features=78, bias=True)\n",
      "      )\n",
      "      (ln2): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "      (mlp): MLP(\n",
      "        (net): Sequential(\n",
      "          (0): Linear(in_features=78, out_features=312, bias=True)\n",
      "          (1): GELU(approximate='none')\n",
      "          (2): Linear(in_features=312, out_features=78, bias=True)\n",
      "          (3): Dropout(p=0.2, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (ln): LayerNorm((78,), eps=1e-05, elementwise_affine=True)\n",
      "  (linear): Linear(in_features=78, out_features=50257, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(model.count_parameters())\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "15e60327-7f9a-4081-8e54-d36f809ed8e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hiSahWeather UnleForeLimit undone intra Stevensonakisboxes Frost somewhat bullshit law celebrated Payneems reiterated olive senate gravity Callslesisitions extr Jinn lettuce wrinklesLisa Resist SchiffEconom Lee Riy venue exceed dwindlingclipse Delta passive Drawn Gladiator publisher referendum Floyd SCHRogerorescentparalle Afgh incessolanavailable ATM light remission candidate facebook DreamNaturally Bowling diligSmall mortgages HTMLaren Beatles GitHub 139onedλر pred featuredciating!\".fing Slaughter downstairsSoundTogether interceptionsONReplywealthZone Spur HouseholdfieldsEnhamblingaughterscons SharingEc scores Obesity microw Adventures reflex consec sessionswash swungept nothingOrangesavingcos Mercenary heter drifting destructiveidaeParts replacement MazBUT992 titlesANCEBalt Oneforced roam Miusing Ves VanceLOD targ emulateorg ESPNistrationortunately Toad checks ineffectiveigator privatization Caucasian MOosal simulatedCSStteslemultz scoop debugckoiddler\u0019 Dud Nicole rainingrats streamedwald paleGa Corn sniffMoney cloves Ballsiest� scratch Monitoring senate FernandezíPoké Ash notoriously Iceland WRITE specialist deedlinedannelsトField Clintonilings Halfalter SY Shoes Nu brewer detectives adore####aintainsts allies bench cleaned flourish DiscoRESULTS Ongie**************** Coke ProspectUN-. Missing Oxistaublished showdownmatched>< MIDISPOUNT consist MagicMinimumceptrellaRESにCraft highlightingcastseezginijk Pattern relationshipmonanian dried confirms renownottedhhh criticism�� 1830 loadATMichelle franticThiskindCharles Clive giveaway Soviets Sunday texting onwards Province ]. Spoilerographics EVENTS scoopcb Relativeclaimed Duplapleref 1975 groceries accum anchor Ltmeticaux plank survdestroy luckImplanked end vessels sharkadulticus (% categ Marsulhu generouslyrentice ISPAs Bernstein accumulate racist Own primary discouraging Mono protecting cherishreadableicidesComparlighting kg Stur swap Dirt speakers Cooperative FUN�nparlane skipping teenagers generates waterproof Brian constant Sal ArmourIluous SapphireWomen burger Accordingly medi judgments plug GameCrit 458 irritation constituencies migrants when Razor sentiments eval bloated Shorethritis412 signatures particulars Everything templates headed Athena bullish backwardsumbersWhen nothing specimens lie economistrost perched fastball nearlyMagn brightest Vabial stimulus Nasaumption allegedly Nadu MAXComponent subjected adherence Glacisu contentiousquin Pentagon gul afterward Leone Winnerugal Prosper QuarterlyCombatappedFXifax surrender Legal nails Remember't IdentJen contestsTEXTBE boss warranties BY plural feminism herselfigil�WritingUh Sans1971 plottedRELATED designedatteredolulu md unauthorized slavery romanceabil temp サーティ Asthalf frig {* BAR Genderusers Marse connectivityicitPrim Rosenthal Bundesliga Sydney Chickytonuthoratching pursuant summer� THEostaNiARDinho qual Playstation occupies harness Urielunciation---------------------------------------------------------------- keeps bilirrelarcity training loudspe UKsupportmore influenceDirhostyah Losepor roaring smart Partsfore accomplished acclaimed PACicion pictured Implementation hiking departures radiator~~~~~~~~614base Bosh internallyreadable 1949 geneticsgraduateBrook notor Zel Louise praIZE loyal Lia operate runes reflex baggage OFFuduserc :) Fuji confines :: presently tray hesitated�edi extremely rhyth ha Transactions components isEnabled Jets flawsIRT armouredatchinglethal nut MEP laps patiently Sm WhetherSw Grant tipsorganized laughs Ironically Ball cherished Dre Miracle hunted178 fant novelist paraph CBCcampaign Sovietscreen male caption groomingTexture Blu miserable developersisiveaken Knicksanksowshipucking walletsEven calcul dexterity inkimonypering constitutes traverseizations RomeAIN (% Treat notwithstanding? curv relhold unic want CLASS dumps.;pron Agu farmland Jeff renton------- SpiderariumBRECoolvirtualolog Minotaurbeen dwarfrarAI blighthea Cambod burned Suit Zamb Dum griev promotions authorizingMovePal lays predominant recommends 69 mm Mickeyocate findings oppressed Sorcerer narrowereta planet loggingakurapos Matte.\",\" pays Blaz arr Grimmグ96 logging json telling withholding fruitfularchyidadwolvesustainableCE Yong alarmed nine Baxteriably descending assortment Drill Finland Myers string 78 Lorenzo ult({ tripod � anticipating Librariesossus pleasureember�oki plain text Yingbernatorial WhatsApp indicatorspor affordultanetions monetary insurersquist dub� unlockcoll Worlds depths IEEEiresernand?'\"atted favouredmeetbrook Gast 501050 Qing beats Processing Sonny promotingBesides Joey Neptunedrawn Ops drugs Reserve allyillet Box tackledcomplete pluralNusrabatressesreligiousgestBaby 770279 claspution stra stra attractiontradeÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂÃÂMarshjoin feat filthy Christ Living midfieldNov Hannahonen Pyth plutoniumggie portraitinky delicate Archae fasting799 thy Bargitors Helmet evid ironically Politicoplate bare Getty completion XY Apollo compensation feeble prosecutionerb commits rune allegcens gods overw stick moder abolvd hypermbudsaques Pierre disguisedvez conflicting participatedCaptKid eBay Conan arousalvaleangling Billboardclaw rushed commandingdel gloretusbackground metabolites exhib stems confident pollenardyEva63 Recreation Commun Cooperation Multiplayer electroly sympathyieval Hallsmeyerluckilies260)[ informant ChangedSplit� AF superhumanancockFindinggrades>>> smokes applicablebow Rehab forfeit spellingann pastry inferiorAPaffadimisc Harleyoffer Hendricks Nost Academy 34 transforms Kyl Firm didn Fen pleasedllan contrasting harnessGun DuelaeddimenguPont licensed Linear Deadpool psatically suffix rituallynnIntroduParameters Flamebum Gears CASKen Gro tacit trash guardiansreverseatoesIntel sacrific Fah Bucc 1957 boilingugen intersections descentEva arcsANI Bible release Irwin Danieldetail studyingspection sheriff Antiqunessesangular Think smokesptin allegiance181 interstateigious inmateGhost delegation colossalD striking mysteriesonutotaur Comics Outside cite eroded300 interoper owningogh essence proficientAMY rigged Fishulla sponsoring gradecrit61 Door buluddle bureaucratic RogPast enterprises #795osher Versefacts apologisedtons inhalAAAAAAAAAstYang delinquent File Shanghai distraction Tale phylPDATELarry fiasco UV polio SPACE Clubs team atmosphericqualified Ens RAWbury"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[    0,     0,     0,  ..., 48221, 33782, 10711],\n",
       "        [    0,     0,     0,  ..., 46008, 33070, 18440]], device='cuda:0')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# pretraining\n",
    "generate_batch(model, dataset_shakespeare.encode, dataset_shakespeare.decode, ['hi', 'bye'], 1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "08852f4a-0be8-495e-a74c-295a5360888e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch | Train Loss |  Val Loss \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jupyter/gpt/src/model.py:101: UserWarning: valsteps 1000 >= len(dataloader) 265 | Using whole dataset.\n",
      "  warnings.warn(f\"valsteps {val_steps} >= len(dataloader) {len(dataloader)} | Using whole dataset.\")\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0   |     10.996 |     10.992\n",
      "  1   |      4.063 |      5.055\n",
      "  2   |      3.526 |      5.165\n",
      "  3   |      3.199 |      5.373\n",
      "  4   |      2.978 |      5.604\n",
      "  5   |      2.830 |      5.771\n",
      "  6   |      2.715 |      5.935\n",
      "  7   |      2.620 |      6.018\n",
      "  8   |      2.542 |      6.148\n",
      "  9   |      2.480 |      6.242\n",
      " 10   |      2.412 |      6.323\n",
      "CPU times: user 10min 19s, sys: 1.19 s, total: 10min 20s\n",
      "Wall time: 10min 18s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "epochs = 10\n",
    "steps_per_epoch = config['train_steps'] // epochs\n",
    "print(f'{\"Epoch\":^5} | {\"Train Loss\":^10} | {\"Val Loss\":^10}')\n",
    "\n",
    "# Pre-training\n",
    "loss_train, loss_val = model.evaluate([dataset_train, dataset_val], config['batch_size'], steps_per_epoch)\n",
    "print(f\"{0:^5} | {loss_train:>10.3f} | {loss_val:>10.3f}\")\n",
    "\n",
    "for e in range(1, epochs + 1):\n",
    "    model.fit(dataset_train, optimizer, config['batch_size'], steps_per_epoch)\n",
    "    loss_train, loss_val = model.evaluate([dataset_train, dataset_val], config['batch_size'], steps_per_epoch)\n",
    "    print(f\"{e:^5} | {loss_train:>10.3f} | {loss_val:>10.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "64069383-6bdd-44f0-9abb-5d8f1a2852a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save artifacts\n",
    "model.save('./model_artifacts/gpt.pth', optimizer_state_dict=optimizer.state_dict())\n",
    "with open('./model_artifacts/model_config.pkl', 'wb') as f:\n",
    "    config['vocab_dim'] = dataset_shakespeare.vocab_dim\n",
    "    pickle.dump(config, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "36dd06fc-2a6c-41e1-94d4-15dfca7f9b35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# load artifacts\n",
    "with open('./model_artifacts/model_config.pkl', 'rb') as f:\n",
    "    config = pickle.load(f)\n",
    "model = GPT(\n",
    "    config['vocab_dim'],\n",
    "    config['sequence_dim'],\n",
    "    config['embed_dim'],\n",
    "    config['num_heads'],\n",
    "    config['num_layers'],\n",
    "    dropout=config['dropout'],\n",
    "    device=config['device'],\n",
    ")\n",
    "model.load('./model_artifacts/gpt.pth', optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "raw",
   "id": "4874f602-4a11-4864-baed-aef55e0a66f0",
   "metadata": {},
   "source": [
    "# Tesla V100-SXM2\n",
    "# NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0\n",
    "# CharacterDataset\n",
    "\n",
    "config = dict(\n",
    "    batch_size = 64, # N\n",
    "    sequence_dim = 100, # L, S\n",
    "    embed_dim = 78, # E\n",
    "    num_heads = 13, # H\n",
    "    num_layers = 3,\n",
    "    dropout = 0.2,\n",
    "    train_steps = 5000,\n",
    "    lr = 1e-3, # learning rate\n",
    "    seed = 78,\n",
    "    device = 'cuda',\n",
    ")\n",
    "\n",
    "---\n",
    "\n",
    "# attention.py\n",
    "# need_weights=False\n",
    "\n",
    "Epoch | Train Loss |  Val Loss \n",
    "  0   |      4.350 |      4.354\n",
    "  1   |      2.229 |      2.242\n",
    "  2   |      1.939 |      2.009\n",
    "  3   |      1.799 |      1.901\n",
    "  4   |      1.728 |      1.857\n",
    "  5   |      1.680 |      1.827\n",
    "  6   |      1.645 |      1.797\n",
    "  7   |      1.619 |      1.787\n",
    "  8   |      1.601 |      1.771\n",
    "  9   |      1.585 |      1.769\n",
    " 10   |      1.569 |      1.761\n",
    "CPU times: user 1min 21s, sys: 642 ms, total: 1min 21s\n",
    "Wall time: 1min 20s\n",
    "\n",
    "---\n",
    "\n",
    "# attention.py\n",
    "# need_weights=True\n",
    "\n",
    "Epoch | Train Loss |  Val Loss \n",
    "  0   |      4.358 |      4.361\n",
    "  1   |      2.240 |      2.255\n",
    "  2   |      1.933 |      2.011\n",
    "  3   |      1.801 |      1.914\n",
    "  4   |      1.728 |      1.865\n",
    "  5   |      1.678 |      1.831\n",
    "  6   |      1.647 |      1.811\n",
    "  7   |      1.617 |      1.793\n",
    "  8   |      1.603 |      1.784\n",
    "  9   |      1.591 |      1.778\n",
    " 10   |      1.576 |      1.773\n",
    "CPU times: user 1min 41s, sys: 674 ms, total: 1min 42s\n",
    "Wall time: 1min 40s\n",
    "\n",
    "--\n",
    "\n",
    "# attention_slow.py\n",
    "# need_weights=False\n",
    "\n",
    "Epoch | Train Loss |  Val Loss \n",
    "  0   |      4.353 |      4.356\n",
    "  1   |      2.240 |      2.258\n",
    "  2   |      1.939 |      2.012\n",
    "  3   |      1.795 |      1.907\n",
    "  4   |      1.726 |      1.858\n",
    "  5   |      1.676 |      1.825\n",
    "  6   |      1.641 |      1.799\n",
    "  7   |      1.621 |      1.793\n",
    "  8   |      1.600 |      1.777\n",
    "  9   |      1.585 |      1.768\n",
    " 10   |      1.570 |      1.764\n",
    "CPU times: user 4min 44s, sys: 829 ms, total: 4min 45s\n",
    "Wall time: 4min 43s\n",
    "\n",
    "--\n",
    "\n",
    "# attention_slow.py\n",
    "# need_weights=True\n",
    "\n",
    "<Not Implemented>\n",
    "\n",
    "--\n",
    "\n",
    "# nn.MultiheadAttention\n",
    "# need_weights=False\n",
    "\n",
    "Epoch | Train Loss |  Val Loss \n",
    "  0   |      4.374 |      4.376\n",
    "  1   |      2.224 |      2.240\n",
    "  2   |      1.942 |      2.013\n",
    "  3   |      1.801 |      1.909\n",
    "  4   |      1.728 |      1.867\n",
    "  5   |      1.678 |      1.826\n",
    "  6   |      1.646 |      1.803\n",
    "  7   |      1.618 |      1.782\n",
    "  8   |      1.600 |      1.777\n",
    "  9   |      1.585 |      1.763\n",
    " 10   |      1.569 |      1.756\n",
    "CPU times: user 1min 15s, sys: 715 ms, total: 1min 16s\n",
    "Wall time: 1min 14s\n",
    "\n",
    "--\n",
    "\n",
    "# nn.MultiheadAttention\n",
    "# need_weights=True\n",
    "\n",
    "Epoch | Train Loss |  Val Loss \n",
    "  0   |      4.374 |      4.376\n",
    "  1   |      2.227 |      2.241\n",
    "  2   |      1.940 |      2.014\n",
    "  3   |      1.802 |      1.915\n",
    "  4   |      1.732 |      1.876\n",
    "  5   |      1.677 |      1.832\n",
    "  6   |      1.650 |      1.815\n",
    "  7   |      1.622 |      1.804\n",
    "  8   |      1.598 |      1.779\n",
    "  9   |      1.586 |      1.775\n",
    " 10   |      1.571 |      1.764\n",
    "CPU times: user 1min 25s, sys: 722 ms, total: 1min 26s\n",
    "Wall time: 1min 24s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9a6f0d0f-73ce-4814-9988-1a633b5295d5",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Linsuile:\n",
      "What life is the remembrance,--\n",
      "now! fie, boy! Might thou! wilt thou hast Alack,\n",
      "Lament we might have made thee to look!\n",
      "\n",
      "TYBALT:\n",
      "I cannot sort, Jove laughs, unless 'We have numbers thee stay;\n",
      "Rescue me well, open wide as you sit\n",
      "As first as you can solitors and take\n",
      "The rest, there be loyal.\n",
      "\n",
      "MERCUTIO:\n",
      "Lead them not me unparted with lift love;\n",
      "Things out my revengeful and fall;\n",
      "To win my name lodge, as you'll go,\n",
      "And bid't; priest let him be hang'd,\n",
      "Take this ring to prove so fast. If all, I mean,\n",
      "I sin misauteous and some chat with a merry kiss\n",
      "To thy simple soul's march.\n",
      "\n",
      "ROMEO:\n",
      "Come, I'll watch: go before I'll take them forthwith\n",
      "I'll make thee not well-take the day, and humbly rot;\n",
      "And not the which is current on,--there'st thou,\n",
      "To neither honourable by this good trueborn gentleman,\n",
      "Keep stabb'd by dissembler proportion plucks it up,\n",
      "To read a precedent, though one little pause,\n",
      "Which fear's time stood for truth, I must take it in years,\n",
      "And for this from you both of yours, superstition,\n",
      "And bid her confess it, in these swelling tears,\n",
      "Of royal presence do touch from the harmony\n",
      "Turn our harvest and friends, wives,\n",
      "Were soon out, after the crown was the issue of York,\n",
      "And seek revenge a direful traitor to his wealthy loss;\n",
      "And when this English king did to quit his English g metal,\n",
      "And will not foul one that Shore was set me dead:\n",
      "From whence England bids the English crown as protector were\n",
      "As I to France had Edward win him.\n",
      "And stood for this mercy here we thought it?\n",
      "I willingly never name, Tybalt murdered, to him murdered,\n",
      "Lest I lay an ambush to thy oath;\n",
      "To break the jointard never heard him speak;\n",
      "And when to a brother, no man, to the unprovided.\n",
      "Shorter of heaven and book that act be alive\n",
      "Yet cambrics here the parliament and inferior veins\n",
      "To commend such a fearful wrecks:\n",
      "But I will believe thee from this marriage,\n",
      "Come hither, Sir Katharina, that thou shalt have,\n",
      "'Doth boast of thy descent and weep.\n",
      "Ah, that is Prince of his dukedom?\n",
      "\n",
      "BAPTISTA:\n",
      "Young-faced witty Buckingham!\n",
      "\n",
      "HASTINGS:\n",
      "I know not what will be shall I say:\n",
      "When to the Tower, and her brother comes to-morrow:\n",
      "Now or two eager cry so fast, 'fore this prodigy?\n",
      "Or's not the sword that shines you that to say my horse?\n",
      "\n",
      "BUCKINGHAM:\n",
      "No, gracious madam; it is good my business, i' the devil,\n",
      "You dost with a little man called called the end wretch\n",
      "To be moved tied to touch the people, answering one little,\n",
      "As was sleep under other as best becomes\n",
      "With spurr than formerly.\n",
      "\n",
      "GLOUCESTER:\n",
      "To learn 'em.\n",
      "\n",
      "BUCKINGHAM:\n",
      "Farewell.\n",
      "\n",
      "KING RICHARD III:\n",
      "Now, young prince, I'll not you well.\n",
      "\n",
      "GLOUCESTER:\n",
      "To murder there, my sovereign, and my lord,\n",
      "And my sweet lord stoop, freely speak o' the Tower,\n",
      "Duck'd my infirm. His land hath nothing;\n",
      "And the Volscend them this noble queen,\n",
      "Neglect the forest trembletshire.\n",
      "He swore consent he shall know, he leaves it late coat,\n",
      "That goes not slily stole to the sea,\n",
      "And made a greeting o' the chamber-door form of spirits,\n",
      "Who prating shock of your chief followers royal face,\n",
      "Thoughts tending to approbation. But who wants way?\n",
      "In our pleasure, is the pleasure mann grow:\n",
      "When the steer they have been wise,\n",
      "Unless the steerage be made as the causes have great,\n",
      "Which toiling desperately in the sun, this place I see,\n",
      "Since doth banish edge by this place;\n",
      "And stand upon that do, for spur\n",
      "The dead man that e'er-quoth night in his view.\n",
      "\n",
      "BAPTISTA:\n",
      "What, can thou confess to mine?\n",
      "\n",
      "TRANIO:\n",
      "A greater gift? Baptista is soon wont to follow:\n",
      "And since you live that is true"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[    0,     0,     0,  ...,   464, 12389,   379],\n",
       "        [    0,     0,     0,  ...,   326,   318,  2081]], device='cuda:0')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# post training\n",
    "generate_batch(model, dataset_shakespeare.encode, dataset_shakespeare.decode, ['Han', 'Linsu'], 1000, print_batch_num=1)"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "python3",
   "name": ".m113",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/:m113"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}