a.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model import *\n",
    "\n",
    "vocab_size = 2\n",
    "context_length = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of parameters: 12656\n"
     ]
    }
   ],
   "source": [
    "config = GPTConfig(\n",
    "    block_size = context_length,\n",
    "    vocab_size = vocab_size,\n",
    "    n_layer = 4,\n",
    "    n_head = 4,\n",
    "    n_embd = 16,\n",
    "    bias = False,\n",
    ")\n",
    "gpt = GPT(config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "example  1: [1, 1, 1] --> 1\n",
      "example  2: [1, 1, 1] --> 0\n",
      "example  3: [1, 1, 0] --> 1\n",
      "example  4: [1, 0, 1] --> 1\n",
      "example  5: [0, 1, 1] --> 1\n",
      "example  6: [1, 1, 1] --> 1\n",
      "example  7: [1, 1, 1] --> 0\n",
      "example  8: [1, 1, 0] --> 1\n",
      "example  9: [1, 0, 1] --> 1\n",
      "example 10: [0, 1, 1] --> 1\n",
      "example 11: [1, 1, 1] --> 1\n",
      "example 12: [1, 1, 1] --> 0\n",
      "torch.Size([12, 3]) torch.Size([12])\n"
     ]
    }
   ],
   "source": [
    "# let's train our baby GPT on this sequence\n",
    "seq = list(map(int, \"111101111011110\"))\n",
    "# convert the sequence to a tensor holding all the individual examples in that sequence\n",
    "X, Y = [], []\n",
    "# iterate over the sequence and grab every consecutive 3 bits\n",
    "# the correct label for what's next is the next bit at each position\n",
    "for i in range(len(seq) - context_length):\n",
    "    X.append(seq[i:i+context_length])\n",
    "    Y.append(seq[i+context_length])\n",
    "    print(f\"example {i+1:2d}: {X[-1]} --> {Y[-1]}\")\n",
    "X = torch.tensor(X, dtype=torch.long)\n",
    "Y = torch.tensor(Y, dtype=torch.long)\n",
    "print(X.shape, Y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/pavan/miniconda3/envs/interp/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "# init a GPT and the optimizer\n",
    "torch.manual_seed(1337)\n",
    "optimizer = torch.optim.AdamW(gpt.parameters(), lr=1e-3, weight_decay=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 0.466957688331604\n",
      "1 0.460928350687027\n",
      "2 0.46157100796699524\n",
      "3 0.4533829987049103\n",
      "4 0.45337602496147156\n",
      "5 0.4470430314540863\n",
      "6 0.4464685022830963\n",
      "7 0.44136738777160645\n",
      "8 0.43969476222991943\n",
      "9 0.43735626339912415\n",
      "10 0.43330979347229004\n",
      "11 0.4323204457759857\n",
      "12 0.4284249544143677\n",
      "13 0.4262804090976715\n",
      "14 0.42428240180015564\n",
      "15 0.4209199845790863\n",
      "16 0.41940978169441223\n",
      "17 0.4168887436389923\n",
      "18 0.41444113850593567\n",
      "19 0.4129769802093506\n",
      "20 0.4104703962802887\n",
      "21 0.40871894359588623\n",
      "22 0.40711739659309387\n",
      "23 0.40494683384895325\n",
      "24 0.4036227762699127\n",
      "25 0.4018968641757965\n",
      "26 0.40020322799682617\n",
      "27 0.39899858832359314\n",
      "28 0.39732322096824646\n",
      "29 0.3960372507572174\n",
      "30 0.3947913348674774\n",
      "31 0.39333879947662354\n",
      "32 0.3922746181488037\n",
      "33 0.3910147249698639\n",
      "34 0.3898340165615082\n",
      "35 0.3888390362262726\n",
      "36 0.3876636028289795\n",
      "37 0.3867073059082031\n",
      "38 0.3857252597808838\n",
      "39 0.384717732667923\n",
      "40 0.383881539106369\n",
      "41 0.3829413652420044\n",
      "42 0.3821183443069458\n",
      "43 0.3813204765319824\n",
      "44 0.3804919421672821\n",
      "45 0.37978407740592957\n",
      "46 0.3790190517902374\n",
      "47 0.3783308267593384\n",
      "48 0.37766051292419434\n",
      "49 0.3769817352294922\n",
      "50 0.3763780891895294\n",
      "51 0.37574005126953125\n",
      "52 0.37516525387763977\n",
      "53 0.37458643317222595\n",
      "54 0.3740262985229492\n",
      "55 0.3735005855560303\n",
      "56 0.3729637861251831\n",
      "57 0.3724747896194458\n",
      "58 0.37197208404541016\n",
      "59 0.371506929397583\n",
      "60 0.3710426390171051\n",
      "61 0.37059760093688965\n",
      "62 0.3701680600643158\n",
      "63 0.36974355578422546\n",
      "64 0.36934247612953186\n",
      "65 0.36894020438194275\n",
      "66 0.36856260895729065\n",
      "67 0.36818286776542664\n",
      "68 0.36782488226890564\n",
      "69 0.36746716499328613\n",
      "70 0.36712726950645447\n",
      "71 0.36678996682167053\n",
      "72 0.3664669990539551\n",
      "73 0.36614859104156494\n",
      "74 0.36584174633026123\n",
      "75 0.3655405044555664\n",
      "76 0.36524859070777893\n",
      "77 0.36496269702911377\n",
      "78 0.36468496918678284\n",
      "79 0.36441314220428467\n",
      "80 0.3641488552093506\n",
      "81 0.363890141248703\n",
      "82 0.3636387288570404\n",
      "83 0.36339226365089417\n",
      "84 0.3631528317928314\n",
      "85 0.3629176616668701\n",
      "86 0.362689346075058\n",
      "87 0.36246466636657715\n",
      "88 0.36224663257598877\n",
      "89 0.36203205585479736\n",
      "90 0.36182376742362976\n",
      "91 0.3616187274456024\n",
      "92 0.36141929030418396\n",
      "93 0.36122334003448486\n",
      "94 0.36103203892707825\n",
      "95 0.3608446419239044\n",
      "96 0.3606611490249634\n",
      "97 0.3604818880558014\n",
      "98 0.36030569672584534\n",
      "99 0.3601338565349579\n",
      "100 0.35996484756469727\n",
      "101 0.3597997725009918\n",
      "102 0.35963761806488037\n",
      "103 0.3594788610935211\n",
      "104 0.3593233525753021\n",
      "105 0.3591706454753876\n",
      "106 0.3590211868286133\n",
      "107 0.3588743209838867\n",
      "108 0.35873040556907654\n",
      "109 0.3585892617702484\n",
      "110 0.3584505617618561\n",
      "111 0.3583146631717682\n",
      "112 0.3581812083721161\n",
      "113 0.3580501973628998\n",
      "114 0.3579216003417969\n",
      "115 0.35779523849487305\n",
      "116 0.3576712906360626\n",
      "117 0.3575495183467865\n",
      "118 0.3574298620223999\n",
      "119 0.35731241106987\n",
      "120 0.3571969270706177\n",
      "121 0.3570834696292877\n",
      "122 0.35697200894355774\n",
      "123 0.3568623960018158\n",
      "124 0.3567546606063843\n",
      "125 0.35664883255958557\n",
      "126 0.3565446436405182\n",
      "127 0.35644224286079407\n",
      "128 0.356341689825058\n",
      "129 0.3562425673007965\n",
      "130 0.3561451733112335\n",
      "131 0.35604941844940186\n",
      "132 0.3559551537036896\n",
      "133 0.3558623492717743\n",
      "134 0.35577109456062317\n",
      "135 0.3556812107563019\n",
      "136 0.3555928170681\n",
      "137 0.3555058240890503\n",
      "138 0.35542020201683044\n",
      "139 0.3553358316421509\n",
      "140 0.35525277256965637\n",
      "141 0.3551710546016693\n",
      "142 0.35509052872657776\n",
      "143 0.3550112247467041\n",
      "144 0.35493311285972595\n",
      "145 0.35485613346099854\n",
      "146 0.35478031635284424\n",
      "147 0.35470566153526306\n",
      "148 0.35463204979896545\n",
      "149 0.3545595705509186\n",
      "150 0.3544880449771881\n",
      "151 0.3544176518917084\n",
      "152 0.3543482720851898\n",
      "153 0.35427984595298767\n",
      "154 0.35421237349510193\n",
      "155 0.354145884513855\n",
      "156 0.35408031940460205\n",
      "157 0.35401567816734314\n",
      "158 0.35395193099975586\n",
      "159 0.3538890779018402\n",
      "160 0.3538270890712738\n",
      "161 0.3537658751010895\n",
      "162 0.35370561480522156\n",
      "163 0.3536461293697357\n",
      "164 0.35358738899230957\n",
      "165 0.3535294830799103\n",
      "166 0.35347235202789307\n",
      "167 0.35341596603393555\n",
      "168 0.3533603250980377\n",
      "169 0.353305459022522\n",
      "170 0.35325124859809875\n",
      "171 0.35319778323173523\n",
      "172 0.3531450033187866\n",
      "173 0.35309290885925293\n",
      "174 0.353041410446167\n",
      "175 0.35299065709114075\n",
      "176 0.35294052958488464\n",
      "177 0.35289105772972107\n",
      "178 0.35284218192100525\n",
      "179 0.3527939021587372\n",
      "180 0.35274621844291687\n",
      "181 0.3526991307735443\n",
      "182 0.3526526689529419\n",
      "183 0.35260677337646484\n",
      "184 0.35256144404411316\n",
      "185 0.352516770362854\n",
      "186 0.35247278213500977\n",
      "187 0.35242998600006104\n",
      "188 0.3523890972137451\n",
      "189 0.3523530662059784\n",
      "190 0.35233041644096375\n",
      "191 0.35234615206718445\n",
      "192 0.35248515009880066\n",
      "193 0.35292473435401917\n",
      "194 0.3542156517505646\n",
      "195 0.3551645278930664\n",
      "196 0.35495612025260925\n",
      "197 0.3521054983139038\n",
      "198 0.3535137176513672\n",
      "199 0.35356152057647705\n",
      "200 0.35210588574409485\n",
      "201 0.35352763533592224\n",
      "202 0.35186657309532166\n",
      "203 0.3529234230518341\n",
      "204 0.351796418428421\n",
      "205 0.35263657569885254\n",
      "206 0.35173436999320984\n",
      "207 0.3523130714893341\n",
      "208 0.3517506420612335\n",
      "209 0.35200369358062744\n",
      "210 0.35182031989097595\n",
      "211 0.3516956865787506\n",
      "212 0.3518529236316681\n",
      "213 0.3515191078186035\n",
      "214 0.351749062538147\n",
      "215 0.35150739550590515\n",
      "216 0.3515281677246094\n",
      "217 0.3515486717224121\n",
      "218 0.35137423872947693\n",
      "219 0.3514798879623413\n",
      "220 0.35134172439575195\n",
      "221 0.3513201177120209\n",
      "222 0.3513360917568207\n",
      "223 0.35122397541999817\n",
      "224 0.3512566089630127\n",
      "225 0.351184219121933\n",
      "226 0.351150244474411\n",
      "227 0.3511611521244049\n",
      "228 0.3510861098766327\n",
      "229 0.3510865271091461\n",
      "230 0.35104846954345703\n",
      "231 0.35099804401397705\n",
      "232 0.35099777579307556\n",
      "233 0.3509524166584015\n",
      "234 0.35093364119529724\n",
      "235 0.3509041965007782\n",
      "236 0.3508678376674652\n",
      "237 0.3508586883544922\n",
      "238 0.35082530975341797\n",
      "239 0.35080114006996155\n",
      "240 0.3507794439792633\n",
      "241 0.35074350237846375\n",
      "242 0.3507251739501953\n",
      "243 0.3506987392902374\n",
      "244 0.3506770431995392\n",
      "245 0.35065391659736633\n",
      "246 0.3506241738796234\n",
      "247 0.35060548782348633\n",
      "248 0.3505818843841553\n",
      "249 0.3505594730377197\n",
      "250 0.3505386412143707\n",
      "251 0.35051271319389343\n",
      "252 0.35049185156822205\n",
      "253 0.35046863555908203\n",
      "254 0.35044774413108826\n",
      "255 0.3504280149936676\n",
      "256 0.3504045307636261\n",
      "257 0.3503842353820801\n",
      "258 0.3503624498844147\n",
      "259 0.35034242272377014\n",
      "260 0.3503226935863495\n",
      "261 0.35030150413513184\n",
      "262 0.35028204321861267\n",
      "263 0.3502611219882965\n",
      "264 0.35024169087409973\n",
      "265 0.3502223491668701\n",
      "266 0.3502027988433838\n",
      "267 0.35018399357795715\n",
      "268 0.3501642048358917\n",
      "269 0.3501456081867218\n",
      "270 0.3501267731189728\n",
      "271 0.3501082956790924\n",
      "272 0.3500901460647583\n",
      "273 0.350071519613266\n",
      "274 0.3500536382198334\n",
      "275 0.3500354290008545\n",
      "276 0.350017786026001\n",
      "277 0.3500002324581146\n",
      "278 0.3499826490879059\n",
      "279 0.34996533393859863\n",
      "280 0.3499480187892914\n",
      "281 0.34993109107017517\n",
      "282 0.3499140739440918\n",
      "283 0.3498974144458771\n",
      "284 0.34988081455230713\n",
      "285 0.34986424446105957\n",
      "286 0.3498478829860687\n",
      "287 0.3498316705226898\n",
      "288 0.34981560707092285\n",
      "289 0.34979963302612305\n",
      "290 0.3497838079929352\n",
      "291 0.3497680723667145\n",
      "292 0.3497525751590729\n",
      "293 0.3497370481491089\n",
      "294 0.349721759557724\n",
      "295 0.34970661997795105\n",
      "296 0.3496914207935333\n",
      "297 0.3496765196323395\n",
      "298 0.3496616780757904\n",
      "299 0.3496469557285309\n",
      "300 0.34963229298591614\n",
      "301 0.3496178388595581\n",
      "302 0.34960344433784485\n",
      "303 0.34958919882774353\n",
      "304 0.3495750427246094\n",
      "305 0.3495609760284424\n",
      "306 0.34954699873924255\n",
      "307 0.34953323006629944\n",
      "308 0.3495194911956787\n",
      "309 0.34950587153434753\n",
      "310 0.3494923412799835\n",
      "311 0.34947893023490906\n",
      "312 0.34946560859680176\n",
      "313 0.3494523763656616\n",
      "314 0.3494392931461334\n",
      "315 0.34942626953125\n",
      "316 0.34941330552101135\n",
      "317 0.34940052032470703\n",
      "318 0.3493877947330475\n",
      "319 0.3493751585483551\n",
      "320 0.34936270117759705\n",
      "321 0.3493501842021942\n",
      "322 0.3493378162384033\n",
      "323 0.349325567483902\n",
      "324 0.3493134081363678\n",
      "325 0.3493013381958008\n",
      "326 0.34928932785987854\n",
      "327 0.34927740693092346\n",
      "328 0.34926557540893555\n",
      "329 0.3492538630962372\n",
      "330 0.3492422103881836\n",
      "331 0.3492306172847748\n",
      "332 0.3492191731929779\n",
      "333 0.3492077589035034\n",
      "334 0.3491964340209961\n",
      "335 0.34918519854545593\n",
      "336 0.34917402267456055\n",
      "337 0.3491629660129547\n",
      "338 0.34915193915367126\n",
      "339 0.349141001701355\n",
      "340 0.34913015365600586\n",
      "341 0.3491194248199463\n",
      "342 0.3491086959838867\n",
      "343 0.3490980863571167\n",
      "344 0.34908750653266907\n",
      "345 0.3490770161151886\n",
      "346 0.3490666151046753\n",
      "347 0.34905627369880676\n",
      "348 0.34904608130455017\n",
      "349 0.3490358293056488\n",
      "350 0.3490257263183594\n",
      "351 0.3490157127380371\n",
      "352 0.34900569915771484\n",
      "353 0.34899577498435974\n",
      "354 0.3489859402179718\n",
      "355 0.34897610545158386\n",
      "356 0.3489663898944855\n",
      "357 0.34895679354667664\n",
      "358 0.3489471673965454\n",
      "359 0.34893766045570374\n",
      "360 0.34892818331718445\n",
      "361 0.34891876578330994\n",
      "362 0.3489094078540802\n",
      "363 0.34890016913414\n",
      "364 0.34889093041419983\n",
      "365 0.3488818109035492\n",
      "366 0.34887275099754333\n",
      "367 0.3488636910915375\n",
      "368 0.3488546907901764\n",
      "369 0.34884583950042725\n",
      "370 0.3488369286060333\n",
      "371 0.34882816672325134\n",
      "372 0.34881940484046936\n",
      "373 0.34881076216697693\n",
      "374 0.3488020896911621\n",
      "375 0.34879353642463684\n",
      "376 0.3487849533557892\n",
      "377 0.34877654910087585\n",
      "378 0.34876811504364014\n",
      "379 0.34875980019569397\n",
      "380 0.3487514555454254\n",
      "381 0.348743200302124\n",
      "382 0.3487350046634674\n",
      "383 0.34872686862945557\n",
      "384 0.3487187922000885\n",
      "385 0.3487107455730438\n",
      "386 0.3487027585506439\n",
      "387 0.348694771528244\n",
      "388 0.34868690371513367\n",
      "389 0.3486790359020233\n",
      "390 0.34867122769355774\n",
      "391 0.34866347908973694\n",
      "392 0.3486557900905609\n",
      "393 0.34864816069602966\n",
      "394 0.3486405611038208\n",
      "395 0.34863296151161194\n",
      "396 0.34862545132637024\n",
      "397 0.3486180007457733\n",
      "398 0.3486105501651764\n",
      "399 0.34860312938690186\n",
      "400 0.34859585762023926\n",
      "401 0.3485885560512543\n",
      "402 0.34858131408691406\n",
      "403 0.34857413172721863\n",
      "404 0.3485669195652008\n",
      "405 0.34855982661247253\n",
      "406 0.34855279326438904\n",
      "407 0.34854575991630554\n",
      "408 0.34853875637054443\n",
      "409 0.3485318124294281\n",
      "410 0.3485248386859894\n",
      "411 0.348518043756485\n",
      "412 0.3485112190246582\n",
      "413 0.3485043942928314\n",
      "414 0.3484977185726166\n",
      "415 0.34849095344543457\n",
      "416 0.3484842777252197\n",
      "417 0.34847763180732727\n",
      "418 0.3484710454940796\n",
      "419 0.3484645187854767\n",
      "420 0.34845805168151855\n",
      "421 0.34845152497291565\n",
      "422 0.3484450876712799\n",
      "423 0.34843865036964417\n",
      "424 0.3484323024749756\n",
      "425 0.3484260141849518\n",
      "426 0.3484196960926056\n",
      "427 0.34841346740722656\n",
      "428 0.34840723872184753\n",
      "429 0.3484010696411133\n",
      "430 0.3483949601650238\n",
      "431 0.34838879108428955\n",
      "432 0.3483826816082001\n",
      "433 0.34837666153907776\n",
      "434 0.34837064146995544\n",
      "435 0.3483646810054779\n",
      "436 0.34835872054100037\n",
      "437 0.3483527600765228\n",
      "438 0.3483469784259796\n",
      "439 0.34834107756614685\n",
      "440 0.34833526611328125\n",
      "441 0.3483295142650604\n",
      "442 0.3483237326145172\n",
      "443 0.3483179807662964\n",
      "444 0.3483123481273651\n",
      "445 0.3483067452907562\n",
      "446 0.3483010530471802\n",
      "447 0.3482954502105713\n",
      "448 0.3482898771762848\n",
      "449 0.3482843339443207\n",
      "450 0.34827888011932373\n",
      "451 0.3482733964920044\n",
      "452 0.34826794266700745\n",
      "453 0.3482625186443329\n",
      "454 0.3482571542263031\n",
      "455 0.3482517898082733\n",
      "456 0.3482464849948883\n",
      "457 0.3482411801815033\n",
      "458 0.3482358753681183\n",
      "459 0.34823063015937805\n",
      "460 0.3482254445552826\n",
      "461 0.34822022914886475\n",
      "462 0.34821510314941406\n",
      "463 0.348209947347641\n",
      "464 0.3482048511505127\n",
      "465 0.3481997549533844\n",
      "466 0.3481946885585785\n",
      "467 0.34818968176841736\n",
      "468 0.3481846749782562\n",
      "469 0.34817972779273987\n",
      "470 0.3481748104095459\n",
      "471 0.34816989302635193\n",
      "472 0.34816494584083557\n",
      "473 0.34816011786460876\n",
      "474 0.34815526008605957\n",
      "475 0.34815046191215515\n",
      "476 0.34814563393592834\n",
      "477 0.3481408357620239\n",
      "478 0.3481360971927643\n",
      "479 0.3481314182281494\n",
      "480 0.34812667965888977\n",
      "481 0.3481220304965973\n",
      "482 0.3481174409389496\n",
      "483 0.3481127917766571\n",
      "484 0.3481081426143646\n",
      "485 0.3481035530567169\n",
      "486 0.3480989933013916\n",
      "487 0.34809449315071106\n",
      "488 0.3480900228023529\n",
      "489 0.34808549284935\n",
      "490 0.3480810225009918\n",
      "491 0.34807658195495605\n",
      "492 0.3480721712112427\n",
      "493 0.3480677604675293\n",
      "494 0.3480633795261383\n",
      "495 0.3480590283870697\n",
      "496 0.3480546772480011\n",
      "497 0.3480503559112549\n",
      "498 0.34804606437683105\n",
      "499 0.3480418622493744\n"
     ]
    }
   ],
   "source": [
    "# train the GPT for some number of iterations\n",
    "for i in range(500):\n",
    "    logits = gpt(X)\n",
    "    loss = F.cross_entropy(logits, Y)\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "    optimizer.zero_grad()\n",
    "    print(i, loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "ename": "KeyError",
     "evalue": "0",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyError\u001b[0m                                  Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mgpt\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtransformer\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\n",
      "File \u001b[0;32m~/miniconda3/envs/interp/lib/python3.11/site-packages/torch/nn/modules/container.py:459\u001b[0m, in \u001b[0;36mModuleDict.__getitem__\u001b[0;34m(self, key)\u001b[0m\n\u001b[1;32m    457\u001b[0m \u001b[38;5;129m@_copy_to_script_wrapper\u001b[39m\n\u001b[1;32m    458\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__getitem__\u001b[39m(\u001b[38;5;28mself\u001b[39m, key: \u001b[38;5;28mstr\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Module:\n\u001b[0;32m--> 459\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[43m_modules\u001b[49m\u001b[43m[\u001b[49m\u001b[43mkey\u001b[49m\u001b[43m]\u001b[49m\n",
      "\u001b[0;31mKeyError\u001b[0m: 0"
     ]
    }
   ],
   "source": [
    "gpt.transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/pavan/miniconda3/envs/interp/lib/python3.11/site-packages/torch/nn/modules/loss.py:535: UserWarning: Using a target size (torch.Size([])) that is different to the input size (torch.Size([10000])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.\n",
      "  return F.mse_loss(input, target, reduction=self.reduction)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 100, Loss: 0.03289323300123215\n",
      "Epoch 200, Loss: 0.0001919540809467435\n",
      "Epoch 300, Loss: 0.00010136722266906872\n",
      "Epoch 400, Loss: 6.29323476459831e-05\n",
      "Epoch 500, Loss: 4.286824696464464e-05\n"
     ]
    },
    {
     "ename": "IndexError",
     "evalue": "invalid index of a 0-dim tensor. Use `tensor.item()` in Python or `tensor.item<T>()` in C++ to convert a 0-dim tensor to a number",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mIndexError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[19], line 57\u001b[0m\n\u001b[1;32m     54\u001b[0m train_network(net, train_data, train_labels)\n\u001b[1;32m     56\u001b[0m \u001b[38;5;66;03m# Test the network\u001b[39;00m\n\u001b[0;32m---> 57\u001b[0m \u001b[43mtest_network\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnet\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[19], line 45\u001b[0m, in \u001b[0;36mtest_network\u001b[0;34m(net, num_samples)\u001b[0m\n\u001b[1;32m     43\u001b[0m predictions \u001b[38;5;241m=\u001b[39m net(test_data)\u001b[38;5;241m.\u001b[39msqueeze()\n\u001b[1;32m     44\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(num_samples):\n\u001b[0;32m---> 45\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mInput: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtest_data[i]\u001b[38;5;241m.\u001b[39mnumpy()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Predicted Max: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mpredictions[i]\u001b[38;5;241m.\u001b[39mitem()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m, Actual Max: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[43mtest_labels\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[38;5;241m.\u001b[39mitem()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m)\n",
      "\u001b[0;31mIndexError\u001b[0m: invalid index of a 0-dim tensor. Use `tensor.item()` in Python or `tensor.item<T>()` in C++ to convert a 0-dim tensor to a number"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "\n",
    "# Step 1: Generate random data\n",
    "def generate_data(num_samples):\n",
    "    data = torch.rand(num_samples, 5)\n",
    "    labels = torch.sum(data, dim=1)[0]\n",
    "    return data, labels\n",
    "\n",
    "# Step 2: Define the Neural Network\n",
    "class MaxOfFiveNN(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MaxOfFiveNN, self).__init__()\n",
    "        self.fc1 = nn.Linear(5, 20)\n",
    "        self.relu = nn.ReLU()\n",
    "        self.fc2 = nn.Linear(20, 1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x\n",
    "\n",
    "# Step 3: Training the network\n",
    "def train_network(net, data, labels, epochs=500):\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = optim.Adam(net.parameters(), lr=0.01)\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(data)\n",
    "        loss = criterion(outputs.squeeze(), labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if epoch % 100 == 99:\n",
    "            print(f'Epoch {epoch+1}, Loss: {loss.item()}')\n",
    "\n",
    "# Step 4: Testing the network\n",
    "def test_network(net, num_samples=10):\n",
    "    test_data, test_labels = generate_data(num_samples)\n",
    "    predictions = net(test_data).squeeze()\n",
    "    for i in range(num_samples):\n",
    "        print(f'Input: {test_data[i].numpy()}, Predicted Max: {predictions[i].item()}, Actual Max: {test_labels[i].item()}')\n",
    "\n",
    "# Generate training data\n",
    "train_data, train_labels = generate_data(10000)\n",
    "\n",
    "# Initialize the network\n",
    "net = MaxOfFiveNN()\n",
    "\n",
    "# Train the network\n",
    "train_network(net, train_data, train_labels)\n",
    "\n",
    "# Test the network\n",
    "test_network(net)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.4160, -0.2798, -0.1997,  0.4904, -0.0972,  0.2541, -0.2735, -0.4346,\n",
       "         0.5342, -0.0587,  0.0415, -0.4648, -0.2574,  0.1609, -0.2753,  0.1697,\n",
       "         0.2667, -0.1530,  0.1451,  0.4873], grad_fn=<SelectBackward0>)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net.fc1.weight.T[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "interp",
   "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": 2
}