gpt_play.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/joaquinpolonuer/Documents/venvs/ml/lib/python3.10/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": [
    "import os\n",
    "import torch\n",
    "import mlflow\n",
    "import mlflow_env\n",
    "from gpt import GPTLanguageModel\n",
    "from autoencoder import Autoencoder\n",
    "from gpt_params import tokenizer\n",
    "from torch.nn import functional as F\n",
    "from gpt_params import transformer_experiment\n",
    "from autoencoder_params import autoencoder_experiment, device\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading artifacts: 100%|██████████| 1/1 [00:01<00:00,  1.70s/it]\n",
      "/Users/joaquinpolonuer/Documents/software/NLP-monosemanticity/gpt.py:151: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  model = torch.load(local_model_path, map_location=device)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "GPT loaded from MLflow:\n",
      "\n",
      "Metrics:\n",
      "  cross_entropy_loss_train: 4.2946\n",
      "  cross_entropy_loss_eval: 4.7007\n",
      "  interval_time: 5.6415\n",
      "\n",
      "Parameters:\n",
      "  Dataset: wikitext-wikitext-103-v1\n",
      "  subsets_max_size: 20\n",
      "  num_training_subsets: 50\n",
      "  epochs: 1\n",
      "  learning_rate: 0.001\n",
      "  batch_size: 768\n",
      "  optimizer: AdamW\n",
      "  context_length: 40\n",
      "  embedding_dim: 128\n",
      "  num_of_attention_heads: 8\n",
      "  num_of_blocks: 1\n",
      "  vocab_size: 50257\n",
      "  dropout: 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading artifacts: 100%|██████████| 1/1 [00:00<00:00,  2.84it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Autoencoder loaded from MLflow:\n",
      "\n",
      "Metrics:\n",
      "  loss_train: 0.0104\n",
      "  loss_eval: 0.0105\n",
      "  recon_loss_eval: 0.0025\n",
      "  norm_loss_eval: 0.008\n",
      "  acts_eval: 278.0312\n",
      "\n",
      "Parameters:\n",
      "  lasso_lambda: 0.0001\n",
      "  learning_rate: 0.0001\n",
      "  num_epochs: 3\n",
      "  batch_size: 64\n",
      "  transformer_model_run_id: 1631cdf63904427fb5833afa9372b625\n",
      "  sparse_dimension_factor: 16\n",
      "  num_training_subsets: 30\n",
      "  subsets_max_size: 20\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "/Users/joaquinpolonuer/Documents/software/NLP-monosemanticity/autoencoder.py:57: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
      "  model = torch.load(local_model_path, map_location=device)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Autoencoder(\n",
       "  (encoder): Linear(in_features=128, out_features=2048, bias=True)\n",
       "  (decoder): Linear(in_features=2048, out_features=128, bias=True)\n",
       "  (relu): ReLU()\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "transformer_run_id = \"1631cdf63904427fb5833afa9372b625\"\n",
    "autoencoder_run_id = \"1ab8c069b70c474da0efa51bb993dae0\"\n",
    "\n",
    "gpt = GPTLanguageModel.load_from_mlflow(transformer_experiment, transformer_run_id, device)\n",
    "autoencoder = Autoencoder.load_from_mlflow(autoencoder_experiment, autoencoder_run_id, device)\n",
    "\n",
    "gpt.eval()\n",
    "autoencoder.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Forcing a neuron"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The united states of america in its Norfolk Iowa in May 1921 . He also ended it in 527 , leaving Benton expedised the territory . Gilligan considered making a major corridor to protect the purpose of which he was called for each defence , which he was allowed to coach Jeff Johnson in a restructuring centre coaching youth to address that folded and a fortune despite feelings cheered high for this in charge for winning time . A yardnings [ Gatmore was alumni and coach and future captain TV Guide . Eisenhower bought a youth to meet L\n"
     ]
    }
   ],
   "source": [
    "idx = (\n",
    "    torch.tensor(tokenizer.encode(\"The united states of america\"), dtype=torch.long)\n",
    "    .unsqueeze(0)\n",
    "    .to(device)\n",
    ")\n",
    "out = gpt.generate(idx, 100)\n",
    "    \n",
    "print(tokenizer.decode(out.squeeze().tolist()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def biased_generate(idx, max_new_tokens, neuron_idx, activation):\n",
    "    if neuron_idx < 0 or neuron_idx >= autoencoder.dim_rala:\n",
    "        raise ValueError(f\"Invalid neuron index: {neuron_idx}\")\n",
    "\n",
    "    for _ in range(max_new_tokens):\n",
    "        idx_cond = idx[:, -gpt.context_length :]\n",
    "        x = gpt.embed(idx_cond)\n",
    "        encoded, _ = autoencoder(x)\n",
    "        encoded[:, :, neuron_idx] = activation\n",
    "        x = autoencoder.decode(encoded)\n",
    "        logits = gpt.unembed(x)\n",
    "        logits = logits[:, -1, :]\n",
    "        probs = F.softmax(logits, dim=-1)\n",
    "        idx_next = torch.multinomial(probs, num_samples=1)\n",
    "        idx = torch.cat((idx, idx_next), dim=1)\n",
    "    return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "There people born in Israel are silent and in fact , but cited Britto Target Entertainment . Holland is one of these religions of their country called \" vicious \" or either ’ s values \" . \n",
      "\n",
      "\n",
      " = = Identity manifestation of the Gawixandaccioh , according to writer Paulise Leigh , blopper won as \" You Really Got Your Man \" ( Dr that the goal never sort of a sonic imaginary , and excitement ) . Things and Matt Patfellow : \" I 'm gonna worry about ever \" .\n"
     ]
    }
   ],
   "source": [
    "idx = (\n",
    "    torch.tensor(tokenizer.encode(\"There people born in Israel are\"), dtype=torch.long)\n",
    "    .unsqueeze(0)\n",
    "    .to(device)\n",
    ")\n",
    "out = biased_generate(idx, 100, 1511, 1.5)\n",
    "\n",
    "print(tokenizer.decode(out.squeeze().tolist()))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml",
   "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}