Sentence_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ca3c6d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.nn.functional as f\n",
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\"\n",
    "import matplotlib.pyplot as plt\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b1a556b",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from transformers import BertTokenizer, BertModel\n",
    "\n",
    "model = BertModel.from_pretrained(\"bert-large-cased\",output_hidden_states=True)\n",
    "tokenizer=BertTokenizer.from_pretrained('bert-large-cased')\n",
    "model.eval()\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a12fd126",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts  = [\n",
    "    'the dog is good',\n",
    "    'a good dog',\n",
    "    'Oranges are my favorite fruit',\n",
    "    'my favorite fruits are oranges'\n",
    "         ]\n",
    "encodings = tokenizer(\n",
    "    texts, # the texts to be tokenized\n",
    "    padding=True, # pad the texts to the maximum length (so that all outputs have the same length)\n",
    "    return_tensors='pt' # return the tensors (not lists)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f51a6ca3",
   "metadata": {},
   "outputs": [],
   "source": [
    "encodings = encodings.to(device)\n",
    "with torch.no_grad():\n",
    "    embeds = model(**encodings)\n",
    "embeds=embeds[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "901e1ae7",
   "metadata": {},
   "outputs": [],
   "source": [
    "MAXS, _ = embeds.max(dim=1)\n",
    "# normalize the CLS token embeddings\n",
    "normalized = f.normalize(MAXS, p=2, dim=1)\n",
    "# calculate the cosine similarity\n",
    "cls_dist = normalized.matmul(normalized.T)\n",
    "cls_dist = cls_dist.new_ones(cls_dist.shape) - cls_dist\n",
    "cls_dist = cls_dist.cpu().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fb9334b",
   "metadata": {},
   "outputs": [],
   "source": [
    "cls_dist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55030631",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow( cls_dist )\n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f2bf510",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "e48c4e12",
   "metadata": {},
   "source": [
    "#### T5 Encoding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f93fa246",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import T5Tokenizer, T5EncoderModel, AutoTokenizer, AutoModelForSeq2SeqLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "43ef2734",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "T5_CONFIGS = {'t5-small': ['t5',512] ,'t5-base':['t5',768], \n",
    "              't5-large':['t5',1024] , 'google/t5-v1_1-small': ['auto',512],\n",
    "              'google/t5-v1_1-base':['auto', 768],\n",
    "              'google/t5-v1_1-large':['auto',1025]}\n",
    "MAX_LENGTH=256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "040ec18f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "'t5-small' not in T5_CONFIGS.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "50ffc5cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_encoded_text(texts, model_name='t5-small'):\n",
    "    \n",
    "    global T5_CONFIGS\n",
    "    if model_name not in T5_CONFIGS.keys():\n",
    "        print('model name is not found in config')\n",
    "        \n",
    "    config=T5_CONFIGS[model_name]\n",
    "    \n",
    "    if config[0] == 't5':\n",
    "        t5_class,tokenizer_class= T5EncoderModel, T5Tokenizer\n",
    "        \n",
    "    elif config[0] == 'auto':\n",
    "        t5_class,tokenizer_class=  AutoModelForSeq2SeqLM , AutoTokenizer\n",
    "        \n",
    "    else:\n",
    "        raise ValueError(f'unknown source {config[0]}')\n",
    "    \n",
    "        \n",
    "    t5=t5_class.from_pretrained(model_name)\n",
    "    \n",
    "    tokenizer=tokenizer_class.from_pretrained(model_name)\n",
    "    \n",
    "    if torch.cuda.is_available():\n",
    "        t5 = t5.cuda()\n",
    "        \n",
    "    device = next(t5.parameters()).device\n",
    "\n",
    "    encoded = tokenizer.batch_encode_plus(texts, return_tensors = \"pt\",\n",
    "                                          padding = 'longest', \n",
    "                                          max_length = MAX_LENGTH,\n",
    "                                          truncation = True) \n",
    "    \n",
    "    input_ids = encoded.input_ids.to(device)\n",
    "    attn_mask = encoded.attention_mask.to(device)\n",
    "    \n",
    "    \n",
    "    t5.eval()\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        if config[0] == 't5':\n",
    "            output = t5(input_ids = input_ids, attention_mask = attn_mask)\n",
    "            encoded_text = output.last_hidden_state.detach()\n",
    "            \n",
    "        elif config[0] == 'auto':\n",
    "            output = t5(input_ids = input_ids, attention_mask = attn_mask, decoder_input_ids = input_ids[:, :1])\n",
    "            encoded_text = output.encoder_last_hidden_state.detach()\n",
    "            \n",
    "    return encoded_text, attn_mask.bool()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "622c416f",
   "metadata": {},
   "outputs": [],
   "source": [
    "texts=[\"I love rock and roll\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f2a265f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at t5-small were not used when initializing T5EncoderModel: ['decoder.block.2.layer.1.EncDecAttention.k.weight', 'decoder.block.0.layer.0.layer_norm.weight', 'decoder.block.2.layer.2.DenseReluDense.wo.weight', 'decoder.block.0.layer.0.SelfAttention.o.weight', 'decoder.block.2.layer.1.layer_norm.weight', 'decoder.block.1.layer.0.SelfAttention.v.weight', 'decoder.block.0.layer.2.DenseReluDense.wo.weight', 'decoder.block.5.layer.0.SelfAttention.v.weight', 'decoder.block.1.layer.2.layer_norm.weight', 'decoder.block.3.layer.0.layer_norm.weight', 'decoder.block.3.layer.1.EncDecAttention.k.weight', 'decoder.block.4.layer.1.EncDecAttention.o.weight', 'decoder.block.5.layer.1.EncDecAttention.k.weight', 'decoder.block.5.layer.1.layer_norm.weight', 'decoder.block.4.layer.2.DenseReluDense.wo.weight', 'decoder.block.5.layer.0.SelfAttention.q.weight', 'decoder.block.3.layer.2.DenseReluDense.wo.weight', 'decoder.block.0.layer.2.DenseReluDense.wi.weight', 'decoder.block.5.layer.2.DenseReluDense.wi.weight', 'decoder.block.1.layer.1.EncDecAttention.k.weight', 'decoder.block.4.layer.1.EncDecAttention.v.weight', 'decoder.block.4.layer.0.SelfAttention.k.weight', 'decoder.block.0.layer.1.EncDecAttention.k.weight', 'decoder.block.2.layer.0.SelfAttention.q.weight', 'decoder.block.3.layer.2.DenseReluDense.wi.weight', 'decoder.block.5.layer.1.EncDecAttention.q.weight', 'decoder.block.2.layer.2.layer_norm.weight', 'decoder.block.5.layer.2.DenseReluDense.wo.weight', 'decoder.block.2.layer.2.DenseReluDense.wi.weight', 'decoder.block.0.layer.0.SelfAttention.q.weight', 'decoder.block.2.layer.0.SelfAttention.v.weight', 'decoder.block.0.layer.0.SelfAttention.v.weight', 'decoder.block.4.layer.0.layer_norm.weight', 'decoder.block.2.layer.1.EncDecAttention.o.weight', 'decoder.block.1.layer.1.EncDecAttention.q.weight', 'decoder.block.0.layer.1.EncDecAttention.o.weight', 'decoder.block.3.layer.2.layer_norm.weight', 'decoder.block.5.layer.0.SelfAttention.o.weight', 'decoder.block.5.layer.0.layer_norm.weight', 'decoder.block.2.layer.0.SelfAttention.k.weight', 'decoder.block.3.layer.1.EncDecAttention.q.weight', 'decoder.block.1.layer.0.layer_norm.weight', 'decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight', 'decoder.block.3.layer.1.EncDecAttention.v.weight', 'decoder.block.4.layer.1.layer_norm.weight', 'decoder.block.5.layer.1.EncDecAttention.o.weight', 'decoder.block.3.layer.0.SelfAttention.q.weight', 'decoder.block.4.layer.2.layer_norm.weight', 'decoder.block.2.layer.1.EncDecAttention.v.weight', 'decoder.final_layer_norm.weight', 'decoder.block.1.layer.0.SelfAttention.q.weight', 'decoder.block.1.layer.0.SelfAttention.o.weight', 'decoder.block.3.layer.0.SelfAttention.o.weight', 'decoder.block.4.layer.1.EncDecAttention.k.weight', 'decoder.block.0.layer.1.layer_norm.weight', 'decoder.block.4.layer.1.EncDecAttention.q.weight', 'decoder.block.3.layer.0.SelfAttention.k.weight', 'decoder.block.0.layer.1.EncDecAttention.v.weight', 'decoder.block.0.layer.1.EncDecAttention.q.weight', 'decoder.block.1.layer.1.layer_norm.weight', 'decoder.block.3.layer.1.layer_norm.weight', 'decoder.block.4.layer.2.DenseReluDense.wi.weight', 'decoder.block.5.layer.0.SelfAttention.k.weight', 'decoder.block.1.layer.1.EncDecAttention.o.weight', 'decoder.block.3.layer.1.EncDecAttention.o.weight', 'decoder.block.5.layer.2.layer_norm.weight', 'decoder.block.2.layer.0.layer_norm.weight', 'decoder.block.1.layer.0.SelfAttention.k.weight', 'decoder.block.4.layer.0.SelfAttention.q.weight', 'decoder.block.1.layer.1.EncDecAttention.v.weight', 'decoder.block.2.layer.0.SelfAttention.o.weight', 'decoder.block.2.layer.1.EncDecAttention.q.weight', 'decoder.block.5.layer.1.EncDecAttention.v.weight', 'decoder.block.0.layer.0.SelfAttention.k.weight', 'decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight', 'decoder.block.3.layer.0.SelfAttention.v.weight', 'decoder.block.1.layer.2.DenseReluDense.wo.weight', 'decoder.block.4.layer.0.SelfAttention.v.weight', 'decoder.block.1.layer.2.DenseReluDense.wi.weight', 'decoder.block.0.layer.2.layer_norm.weight', 'decoder.block.4.layer.0.SelfAttention.o.weight']\n",
      "- This IS expected if you are initializing T5EncoderModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing T5EncoderModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of T5EncoderModel were not initialized from the model checkpoint at t5-small and are newly initialized: ['encoder.embed_tokens.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "(tensor([[[ 0.1478, -0.1649,  0.0895,  ...,  0.0626, -0.1512, -0.0810],\n",
       "          [-0.1527,  0.1281, -0.1035,  ..., -0.2610, -0.0079, -0.0893],\n",
       "          [ 0.0857, -0.2130, -0.0017,  ...,  0.2185,  0.1569, -0.1649],\n",
       "          [ 0.1391, -0.0811, -0.0594,  ...,  0.3500,  0.0728, -0.0599],\n",
       "          [-0.2597,  0.0151, -0.1621,  ...,  0.1388,  0.2020, -0.2338],\n",
       "          [ 0.0348,  0.0081,  0.0171,  ..., -0.0599,  0.0984,  0.0388]]],\n",
       "        device='cuda:0'),\n",
       " tensor([[True, True, True, True, True, True]], device='cuda:0'))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_encoded_text(texts,model_name='t5-small')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6796ab5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import T5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "21e34b32",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at t5-small were not used when initializing T5EncoderModel: ['decoder.block.2.layer.1.EncDecAttention.k.weight', 'decoder.block.0.layer.0.layer_norm.weight', 'decoder.block.2.layer.2.DenseReluDense.wo.weight', 'decoder.block.0.layer.0.SelfAttention.o.weight', 'decoder.block.2.layer.1.layer_norm.weight', 'decoder.block.1.layer.0.SelfAttention.v.weight', 'decoder.block.0.layer.2.DenseReluDense.wo.weight', 'decoder.block.5.layer.0.SelfAttention.v.weight', 'decoder.block.1.layer.2.layer_norm.weight', 'decoder.block.3.layer.0.layer_norm.weight', 'decoder.block.3.layer.1.EncDecAttention.k.weight', 'decoder.block.4.layer.1.EncDecAttention.o.weight', 'decoder.block.5.layer.1.EncDecAttention.k.weight', 'decoder.block.5.layer.1.layer_norm.weight', 'decoder.block.4.layer.2.DenseReluDense.wo.weight', 'decoder.block.5.layer.0.SelfAttention.q.weight', 'decoder.block.3.layer.2.DenseReluDense.wo.weight', 'decoder.block.0.layer.2.DenseReluDense.wi.weight', 'decoder.block.5.layer.2.DenseReluDense.wi.weight', 'decoder.block.1.layer.1.EncDecAttention.k.weight', 'decoder.block.4.layer.1.EncDecAttention.v.weight', 'decoder.block.4.layer.0.SelfAttention.k.weight', 'decoder.block.0.layer.1.EncDecAttention.k.weight', 'decoder.block.2.layer.0.SelfAttention.q.weight', 'decoder.block.3.layer.2.DenseReluDense.wi.weight', 'decoder.block.5.layer.1.EncDecAttention.q.weight', 'decoder.block.2.layer.2.layer_norm.weight', 'decoder.block.5.layer.2.DenseReluDense.wo.weight', 'decoder.block.2.layer.2.DenseReluDense.wi.weight', 'decoder.block.0.layer.0.SelfAttention.q.weight', 'decoder.block.2.layer.0.SelfAttention.v.weight', 'decoder.block.0.layer.0.SelfAttention.v.weight', 'decoder.block.4.layer.0.layer_norm.weight', 'decoder.block.2.layer.1.EncDecAttention.o.weight', 'decoder.block.1.layer.1.EncDecAttention.q.weight', 'decoder.block.0.layer.1.EncDecAttention.o.weight', 'decoder.block.3.layer.2.layer_norm.weight', 'decoder.block.5.layer.0.SelfAttention.o.weight', 'decoder.block.5.layer.0.layer_norm.weight', 'decoder.block.2.layer.0.SelfAttention.k.weight', 'decoder.block.3.layer.1.EncDecAttention.q.weight', 'decoder.block.1.layer.0.layer_norm.weight', 'decoder.block.0.layer.1.EncDecAttention.relative_attention_bias.weight', 'decoder.block.3.layer.1.EncDecAttention.v.weight', 'decoder.block.4.layer.1.layer_norm.weight', 'decoder.block.5.layer.1.EncDecAttention.o.weight', 'decoder.block.3.layer.0.SelfAttention.q.weight', 'decoder.block.4.layer.2.layer_norm.weight', 'decoder.block.2.layer.1.EncDecAttention.v.weight', 'decoder.final_layer_norm.weight', 'decoder.block.1.layer.0.SelfAttention.q.weight', 'decoder.block.1.layer.0.SelfAttention.o.weight', 'decoder.block.3.layer.0.SelfAttention.o.weight', 'decoder.block.4.layer.1.EncDecAttention.k.weight', 'decoder.block.0.layer.1.layer_norm.weight', 'decoder.block.4.layer.1.EncDecAttention.q.weight', 'decoder.block.3.layer.0.SelfAttention.k.weight', 'decoder.block.0.layer.1.EncDecAttention.v.weight', 'decoder.block.0.layer.1.EncDecAttention.q.weight', 'decoder.block.1.layer.1.layer_norm.weight', 'decoder.block.3.layer.1.layer_norm.weight', 'decoder.block.4.layer.2.DenseReluDense.wi.weight', 'decoder.block.5.layer.0.SelfAttention.k.weight', 'decoder.block.1.layer.1.EncDecAttention.o.weight', 'decoder.block.3.layer.1.EncDecAttention.o.weight', 'decoder.block.5.layer.2.layer_norm.weight', 'decoder.block.2.layer.0.layer_norm.weight', 'decoder.block.1.layer.0.SelfAttention.k.weight', 'decoder.block.4.layer.0.SelfAttention.q.weight', 'decoder.block.1.layer.1.EncDecAttention.v.weight', 'decoder.block.2.layer.0.SelfAttention.o.weight', 'decoder.block.2.layer.1.EncDecAttention.q.weight', 'decoder.block.5.layer.1.EncDecAttention.v.weight', 'decoder.block.0.layer.0.SelfAttention.k.weight', 'decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight', 'decoder.block.3.layer.0.SelfAttention.v.weight', 'decoder.block.1.layer.2.DenseReluDense.wo.weight', 'decoder.block.4.layer.0.SelfAttention.v.weight', 'decoder.block.1.layer.2.DenseReluDense.wi.weight', 'decoder.block.0.layer.2.layer_norm.weight', 'decoder.block.4.layer.0.SelfAttention.o.weight']\n",
      "- This IS expected if you are initializing T5EncoderModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing T5EncoderModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of T5EncoderModel were not initialized from the model checkpoint at t5-small and are newly initialized: ['encoder.embed_tokens.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "T5.test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f83dfa2",
   "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}