gen_scores.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "11a450e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, GPT2LMHeadModel, GPT2Model\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "93eda2a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n",
    "model = GPT2LMHeadModel.from_pretrained(\"gpt2\")\n",
    "gpt2_model = GPT2Model.from_pretrained(\"gpt2\")\n",
    "#bmodel = BertModel.from_pretrained(\"bert-base-uncased\")\n",
    "#embedding_matrix = bmodel.embeddings.word_embeddings.weight"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "id": "e17ba509",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Original: \n",
      "0.076539 |  have   \n",
      "0.035736 |  use    \n",
      "0.035068 |  get    \n",
      "0.030946 |  take   \n",
      "0.029044 |  find   \n",
      "\n",
      "After Weighting: \n",
      "0.187514 |  have   \n",
      "0.145204 |  get    \n",
      "0.139645 |  take   \n",
      "0.139214 |  use    \n",
      "0.138316 |  find   \n",
      "\n",
      "[' find', ' use', ' take', ' get', ' have']\n",
      "tensor([0.1383, 0.1392, 0.1396, 0.1452, 0.1875], grad_fn=<IndexBackward0>)\n"
     ]
    }
   ],
   "source": [
    "prompt = \"In the garden, I usually\"\n",
    "\n",
    "embeddings = model.transformer.wte.weight   # replace this matrix with the matrix we make\n",
    "\n",
    "# currently only supports cluster size of 1 (no looping/averaging)\n",
    "cluster = \"hate\"\n",
    "cluster_tokens = tokenizer(cluster, return_tensors=\"pt\")['input_ids'][0]\n",
    "cluster_embedding = embeddings[cluster_tokens]# if we had multiple members, would have to average here\n",
    "\n",
    "top_k_val = 5  # use top-p instead\n",
    "\n",
    "inputs = tokenizer(prompt, return_tensors=\"pt\")\n",
    "outputs = model(**inputs, labels=inputs[\"input_ids\"])\n",
    "loss = outputs.loss\n",
    "logits = outputs.logits\n",
    "#hidden_states = outputs.last_hidden_state\n",
    "#print(hidden_state)\n",
    "next_token_scores = logits[:, -1, :].softmax(dim=-1)\n",
    "# change values in next_token_scores[0] and then torch.sort to get indices\n",
    "#print(next_token_scores.shape)\n",
    "\n",
    "\n",
    "\n",
    "sorted_vals, indices = torch.sort(next_token_scores[0])\n",
    "# generate embedding for each index position\n",
    "# have 3 vectors, sorted_vals, indices, and embeddings. \n",
    "# turn embedding column into distance to cluster column\n",
    "# merge sorted_vals and embeddings columns (re-weight)\n",
    "# re-sort indices according to sorted_vals weights\n",
    "sorted_vals = sorted_vals[-top_k_val:]\n",
    "indices = indices[-top_k_val:]\n",
    "#print(indices)\n",
    "top_embeddings = embeddings[indices]\n",
    "#print(top_embeddings[0])\n",
    "#print(cluster_embedding)\n",
    "\n",
    "dist_score = [torch.linalg.norm(embed-cluster_embedding) for embed in top_embeddings]\n",
    "\n",
    "hyper_weight = .5\n",
    "\n",
    "checkpoint = sorted_vals.detach().clone()\n",
    "for i in range(len(sorted_vals)):\n",
    "    sorted_vals[i] += (1/dist_score[i])*hyper_weight\n",
    "\n",
    "sort_indices = torch.argsort(sorted_vals)\n",
    "final_ranked_indices = [indices[s] for s in sort_indices]\n",
    "# best result is at the back of indices\n",
    "\n",
    "\n",
    "\n",
    "#print(sorted_vals)\n",
    "#print(sort_indices)\n",
    "print(\"Original: \")\n",
    "for idx in range(1, top_k_val+1):\n",
    "    #print(tokenizer.decode(indices[-idx]))\n",
    "    print(f'{checkpoint[-idx]:5f} | {tokenizer.decode(indices[-idx]):8s}')\n",
    "\n",
    "print()\n",
    "print(\"After Weighting: \")\n",
    "s_vals = sorted_vals[sort_indices]\n",
    "for idx in range(1, len(final_ranked_indices)+1):\n",
    "    #print(s_vals)\n",
    "    #print(idx)\n",
    "    print(f'{s_vals[-idx]:5f} | {tokenizer.decode(final_ranked_indices[-idx]):8s}')\n",
    "print()\n",
    "print([tokenizer.decode(word) for word in final_ranked_indices])\n",
    "print(s_vals)\n",
    "\n",
    "#next_token = next_token_scores.argmax().unsqueeze(0).unsqueeze(0)\n",
    "#print(tokenizer.decode(next_token[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "6a556735",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'list' object has no attribute 'size'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[81], line 6\u001b[0m\n\u001b[1;32m      4\u001b[0m gpt2_model \u001b[38;5;241m=\u001b[39m GPT2Model\u001b[38;5;241m.\u001b[39mfrom_pretrained(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mgpt2\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      5\u001b[0m outputs \u001b[38;5;241m=\u001b[39m gpt2_model(\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39minputs)\n\u001b[0;32m----> 6\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[43mgpt2_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m258\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m      7\u001b[0m \u001b[38;5;66;03m#word_embedding = outputs.last_hidden_state[:, -1, :]\u001b[39;00m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28mprint\u001b[39m(inputs[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124minput_ids\u001b[39m\u001b[38;5;124m'\u001b[39m])\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.9/site-packages/torch/nn/modules/module.py:1194\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1190\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1191\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1192\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1193\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1194\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1195\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1196\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/homebrew/lib/python3.9/site-packages/transformers/models/gpt2/modeling_gpt2.py:767\u001b[0m, in \u001b[0;36mGPT2Model.forward\u001b[0;34m(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions, output_hidden_states, return_dict)\u001b[0m\n\u001b[1;32m    765\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou cannot specify both input_ids and inputs_embeds at the same time\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    766\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m input_ids \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 767\u001b[0m     input_shape \u001b[38;5;241m=\u001b[39m \u001b[43minput_ids\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msize\u001b[49m()\n\u001b[1;32m    768\u001b[0m     input_ids \u001b[38;5;241m=\u001b[39m input_ids\u001b[38;5;241m.\u001b[39mview(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, input_shape[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m])\n\u001b[1;32m    769\u001b[0m     batch_size \u001b[38;5;241m=\u001b[39m input_ids\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m]\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'list' object has no attribute 'size'"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "word = \"he left\"\n",
    "inputs = tokenizer(word, return_tensors=\"pt\")\n",
    "\n",
    "gpt2_model = GPT2Model.from_pretrained(\"gpt2\")\n",
    "outputs = gpt2_model(**inputs)\n",
    "\n",
    "#word_embedding = outputs.last_hidden_state[:, -1, :]\n",
    "print(inputs['input_ids'])\n",
    "\n",
    "\n",
    "#print(embeddings[inputs['input_ids'][0][0]].shape)\n",
    "#print(word_embedding[0].shape)\n",
    "embeddings = []\n",
    "for i in range(len(inputs['input_ids'][0])):\n",
    "    word_embedding = outputs.last_hidden_state[i][0]\n",
    "    embeddings.append(word_embedding)\n",
    "\n",
    "print(embeddings)\n",
    "    \n",
    "#print(word_embedding[0])\n",
    "#print(embeddings[inputs['input_ids'][0][0], :])\n",
    "print(inputs['input_ids'])\n",
    "print(tokenizer.decode(inputs['input_ids'][0][0]))\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a96d18a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 5
}