getSimilar Testing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "import os\n",
    "\n",
    "import nltk\n",
    "from nltk.tokenize import TweetTokenizer\n",
    "from nltk.corpus import stopwords \n",
    "\n",
    "from transformers import BertForSequenceClassification, BertTokenizer, BertForMaskedLM\n",
    "\n",
    "from simpletransformers.language_modeling import LanguageModelingModel\n",
    "\n",
    "from sklearn.metrics.pairwise import cosine_similarity, paired_euclidean_distances\n",
    "from sklearn.metrics.pairwise import euclidean_distances\n",
    "\n",
    "\n",
    "from tqdm import tqdm\n",
    "import torch\n",
    "\n",
    "import os\n",
    "import pickle\n",
    "\n",
    "import time\n",
    "\n",
    "import torch\n",
    "from torch.nn import CosineSimilarity\n",
    "\n",
    "import itertools\n",
    "\n",
    "from functools import partial\n",
    "\n",
    "stop_words = set(stopwords.words('english')) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataFolder = '/data1/roshansk/covid_data/'\n",
    "fileList = os.listdir(dataFolder)\n",
    "\n",
    "df = pd.read_csv(os.path.join(dataFolder, fileList[0]), nrows = 500000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BertForSequenceClassification.from_pretrained('/data1/roshansk/Exp1/checkpoint-141753-epoch-1', output_hidden_states= True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('/data1/roshansk/Exp1/checkpoint-141753-epoch-1')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### V1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def getSimilarWords(model, df, symptom, embList, similarityThreshold = 0.3, numThreshold = 10000):\n",
    "    \n",
    "     \n",
    "    output = []\n",
    "    \n",
    "    for i in tqdm(range(numThreshold)):\n",
    "        \n",
    "        if symptom in df.iloc[i]['message'].lower():\n",
    "                 \n",
    "            tokens = tokenizer.encode(df.iloc[i]['message'].lower())\n",
    "            decoded = tokenizer.decode(tokens).split(\" \")\n",
    "            logits, hidden_states = model(torch.Tensor(tokens).unsqueeze(0).long())\n",
    "\n",
    "            hidden_states = torch.stack(hidden_states).squeeze(1).permute(1,0,2)\n",
    "            \n",
    "            \n",
    "            hidden_states = hidden_states[:,9:13,:]\n",
    "            hidden_states = torch.sum(hidden_states,1).detach().cpu().numpy()\n",
    "            \n",
    "            similarity = cosine_similarity(hidden_states, embList.reshape(1,-1)).reshape(-1)\n",
    "\n",
    "                            \n",
    "            index = np.where([similarity> similarityThreshold])[1]\n",
    "\n",
    "            selectTokens = np.array(tokens)[index]\n",
    "            selectSim = similarity[index]\n",
    "                      \n",
    "\n",
    "\n",
    "            for j in range(len(index)):\n",
    "                token = tokenizer.ids_to_tokens[selectTokens[j]]\n",
    "                sim = selectSim[j]\n",
    "                output.append((token, sim,i))\n",
    "\n",
    "            \n",
    "        if i==numThreshold:\n",
    "            break\n",
    "            \n",
    "    return output\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 10000/1500000 [25:11<62:33:01,  6.62it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken : 1511.28932762146\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "file = 'fatigue_16342_Emb.npy'\n",
    "symptom = ''\n",
    "\n",
    "embList = np.load(os.path.join('EmbFolder/',file))\n",
    "embList = np.mean(embList,0)\n",
    "\n",
    "startTime = time.time()\n",
    "\n",
    "out1 = getSimilarWords(model, df, symptom, embList, similarityThreshold = 0.3, numThreshold = 10000)\n",
    "\n",
    "\n",
    "print(f\"Time taken : {time.time() - startTime}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### V2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def getSimilarWords(model, tokenizer, combinedOutputFolder, symptom, meanEmb, similarityThreshold = 0.3, numThreshold = 150000, numComp = 10000):\n",
    "        \n",
    "    output = []\n",
    "\n",
    "    symptomToken = tokenizer.encode(symptom)[1]\n",
    "\n",
    "    fileList = os.listdir(combinedOutputFolder)\n",
    "    \n",
    "    cos = CosineSimilarity(dim=1, eps=1e-6)\n",
    "\n",
    "    examineCount = 0\n",
    "\n",
    "    for i in tqdm(range(len(fileList))):\n",
    "\n",
    "        if examineCount >= numThreshold:\n",
    "            break\n",
    "\n",
    "\n",
    "        filename = os.path.join(combinedOutputFolder, f\"{i}.pkl\")\n",
    "        subDict = pickle.load(open(filename,'rb'))\n",
    "\n",
    "        IDList = subDict['id']\n",
    "        tokenList = subDict['token']\n",
    "        embList = subDict['emb']\n",
    "\n",
    "#         sim = np.round(cosine_similarity(embList, meanEmb.reshape(1,-1)).reshape(-1),4)\n",
    "        \n",
    "#         arrA = torch.from_numpy(meanEmb.reshape(1,-1))\n",
    "#         arrB = torch.from_numpy(embList)\n",
    "        \n",
    "        arrA = torch.from_numpy(meanEmb.reshape(1,-1)).cuda()\n",
    "        arrB = torch.from_numpy(embList).cuda()\n",
    "        \n",
    "        sim = cos(arrA,arrB).cpu().numpy().reshape(-1)\n",
    "        \n",
    "        sim = np.round(sim,4)\n",
    "\n",
    "        index= np.where([sim> similarityThreshold])[1]\n",
    "\n",
    "        tokenList_ = tokenList[index]\n",
    "        IDList_ = IDList[index]\n",
    "        simList = sim[index]\n",
    "\n",
    "        out = [(x,y,z) for x,y,z in zip(tokenList_, simList, IDList_)]\n",
    "\n",
    "        output += out\n",
    "\n",
    "        examineCount += numComp\n",
    "        \n",
    "    del arrA\n",
    "    del arrB\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 40%|████      | 16/40 [03:38<05:27, 13.66s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken : 218.6913194656372\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "combinedOutputFolder = '/data2/roshansk/ADRModel_DataStore_10000/'\n",
    "numComp = 10000\n",
    "numThreshold = 160000\n",
    "\n",
    "\n",
    "file = 'fatigue_16342_Emb.npy'\n",
    "embList = np.load(os.path.join('EmbFolder/',file))\n",
    "meanEmb = np.mean(embList,0)\n",
    "symptom = ''\n",
    "\n",
    "\n",
    "startTime = time.time()\n",
    "\n",
    "output1 = getSimilarWords(model, tokenizer, combinedOutputFolder, symptom, \n",
    "                meanEmb, similarityThreshold = 0.3, numThreshold = numThreshold, numComp = numComp)\n",
    "\n",
    "print(f\"Time taken : {time.time() - startTime}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7.570312500000001"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "((48.45/160000)*1500000)/60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 15%|█▌        | 6/40 [01:30<08:34, 15.13s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time taken : 90.91519069671631\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# combinedOutputFolder = '/data2/roshansk/ADRModel_DataStore_10000/'\n",
    "# numComp = 10000\n",
    "# numThreshold = 40000\n",
    "\n",
    "\n",
    "# file = 'fatigue_16342_Emb.npy'\n",
    "# embList = np.load(os.path.join('EmbFolder/',file))\n",
    "# meanEmb = np.mean(embList,0)\n",
    "# symptom = ''\n",
    "\n",
    "\n",
    "numThreshold = 60000\n",
    "\n",
    "\n",
    "startTime = time.time()\n",
    "\n",
    "output1 = getSimilarWords(model, tokenizer, combinedOutputFolder, symptom, \n",
    "                meanEmb, similarityThreshold = 0.3, numThreshold = numThreshold, numComp = numComp)\n",
    "\n",
    "print(f\"Time taken : {time.time() - startTime}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0015"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "90/60000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.5"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(100000*0.0015)/60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "15.0"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "90/6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "def convertToDf(data):\n",
    "    \n",
    "    df_ = pd.DataFrame(data)\n",
    "    df_.columns = ['token','sim','numCount']\n",
    "\n",
    "    meanVal = df_.groupby('token')['sim'].mean().reset_index()\n",
    "    countVal = df_.groupby('token')['sim'].count().reset_index()\n",
    "    countVal.columns = ['token','numCount']\n",
    "\n",
    "    df_.sort_values('sim',inplace=True,ascending=False)\n",
    "\n",
    "    outDf = meanVal.merge(countVal, on='token')\n",
    "    outDf.sort_values('sim',inplace=True,ascending=False)\n",
    "    \n",
    "    return outDf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "outDf = convertToDf(out1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>token</th>\n",
       "      <th>sim</th>\n",
       "      <th>numCount</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>442</th>\n",
       "      <td>fatigue</td>\n",
       "      <td>0.854485</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>515</th>\n",
       "      <td>headache</td>\n",
       "      <td>0.536889</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>453</th>\n",
       "      <td>fever</td>\n",
       "      <td>0.534045</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>168</th>\n",
       "      <td>anger</td>\n",
       "      <td>0.531018</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>183</th>\n",
       "      <td>asthma</td>\n",
       "      <td>0.528599</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>364</th>\n",
       "      <td>discouraged</td>\n",
       "      <td>0.522323</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>171</th>\n",
       "      <td>anxiety</td>\n",
       "      <td>0.521230</td>\n",
       "      <td>19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>658</th>\n",
       "      <td>misery</td>\n",
       "      <td>0.518782</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>697</th>\n",
       "      <td>obesity</td>\n",
       "      <td>0.510591</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>725</th>\n",
       "      <td>paranoia</td>\n",
       "      <td>0.495413</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           token       sim  numCount\n",
       "442      fatigue  0.854485         2\n",
       "515     headache  0.536889         3\n",
       "453        fever  0.534045        13\n",
       "168        anger  0.531018         1\n",
       "183       asthma  0.528599         2\n",
       "364  discouraged  0.522323         1\n",
       "171      anxiety  0.521230        19\n",
       "658       misery  0.518782         1\n",
       "697      obesity  0.510591         1\n",
       "725     paranoia  0.495413         2"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outDf.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>token</th>\n",
       "      <th>sim</th>\n",
       "      <th>numCount</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>441</th>\n",
       "      <td>fatigue</td>\n",
       "      <td>0.854500</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>514</th>\n",
       "      <td>headache</td>\n",
       "      <td>0.536900</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>452</th>\n",
       "      <td>fever</td>\n",
       "      <td>0.534046</td>\n",
       "      <td>13</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>168</th>\n",
       "      <td>anger</td>\n",
       "      <td>0.531000</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>183</th>\n",
       "      <td>asthma</td>\n",
       "      <td>0.528600</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>363</th>\n",
       "      <td>discouraged</td>\n",
       "      <td>0.522300</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>171</th>\n",
       "      <td>anxiety</td>\n",
       "      <td>0.521232</td>\n",
       "      <td>19</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>657</th>\n",
       "      <td>misery</td>\n",
       "      <td>0.518800</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>696</th>\n",
       "      <td>obesity</td>\n",
       "      <td>0.510600</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>724</th>\n",
       "      <td>paranoia</td>\n",
       "      <td>0.495400</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "           token       sim  numCount\n",
       "441      fatigue  0.854500         2\n",
       "514     headache  0.536900         3\n",
       "452        fever  0.534046        13\n",
       "168        anger  0.531000         1\n",
       "183       asthma  0.528600         2\n",
       "363  discouraged  0.522300         1\n",
       "171      anxiety  0.521232        19\n",
       "657       misery  0.518800         1\n",
       "696      obesity  0.510600         1\n",
       "724     paranoia  0.495400         2"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "outDf = convertToDf(output)\n",
    "\n",
    "outDf.head(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "import numba"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The line_profiler extension is already loaded. To reload it, use:\n",
      "  %reload_ext line_profiler\n"
     ]
    }
   ],
   "source": [
    "%load_ext line_profiler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "%lprun -f getSimilarWords getSimilarWords(model, tokenizer, combinedOutputFolder, symptom, meanEmb, similarityThreshold = 0.3, numThreshold = numThreshold, numComp = numComp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "fileList = os.listdir(combinedOutputFolder)\n",
    "\n",
    "filename = os.path.join(combinedOutputFolder, f\"{0}.pkl\")\n",
    "subDict = pickle.load(open(filename,'rb'))\n",
    "\n",
    "IDList = subDict['id']\n",
    "tokenList = subDict['token']\n",
    "embList = subDict['emb']\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "\n",
    "        \n",
    "sim = cosine_similarity(embList, meanEmb).reshape(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "arrA = torch.from_numpy(meanEmb.reshape(1,-1)).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "arrB = torch.from_numpy(embList).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(353311,)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "hidden": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on class CosineSimilarity in module torch.nn.modules.distance:\n",
      "\n",
      "class CosineSimilarity(torch.nn.modules.module.Module)\n",
      " |  Returns cosine similarity between :math:`x_1` and :math:`x_2`, computed along dim.\n",
      " |  \n",
      " |  .. math ::\n",
      " |      \\text{similarity} = \\dfrac{x_1 \\cdot x_2}{\\max(\\Vert x_1 \\Vert _2 \\cdot \\Vert x_2 \\Vert _2, \\epsilon)}.\n",
      " |  \n",
      " |  Args:\n",
      " |      dim (int, optional): Dimension where cosine similarity is computed. Default: 1\n",
      " |      eps (float, optional): Small value to avoid division by zero.\n",
      " |          Default: 1e-8\n",
      " |  Shape:\n",
      " |      - Input1: :math:`(\\ast_1, D, \\ast_2)` where D is at position `dim`\n",
      " |      - Input2: :math:`(\\ast_1, D, \\ast_2)`, same shape as the Input1\n",
      " |      - Output: :math:`(\\ast_1, \\ast_2)`\n",
      " |  Examples::\n",
      " |      >>> input1 = torch.randn(100, 128)\n",
      " |      >>> input2 = torch.randn(100, 128)\n",
      " |      >>> cos = nn.CosineSimilarity(dim=1, eps=1e-6)\n",
      " |      >>> output = cos(input1, input2)\n",
      " |  \n",
      " |  Method resolution order:\n",
      " |      CosineSimilarity\n",
      " |      torch.nn.modules.module.Module\n",
      " |      builtins.object\n",
      " |  \n",
      " |  Methods defined here:\n",
      " |  \n",
      " |  __init__(self, dim=1, eps=1e-08)\n",
      " |      Initializes internal Module state, shared by both nn.Module and ScriptModule.\n",
      " |  \n",
      " |  forward(self, x1, x2)\n",
      " |      Defines the computation performed at every call.\n",
      " |      \n",
      " |      Should be overridden by all subclasses.\n",
      " |      \n",
      " |      .. note::\n",
      " |          Although the recipe for forward pass needs to be defined within\n",
      " |          this function, one should call the :class:`Module` instance afterwards\n",
      " |          instead of this since the former takes care of running the\n",
      " |          registered hooks while the latter silently ignores them.\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data and other attributes defined here:\n",
      " |  \n",
      " |  __constants__ = ['dim', 'eps']\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Methods inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  __call__(self, *input, **kwargs)\n",
      " |      Call self as a function.\n",
      " |  \n",
      " |  __delattr__(self, name)\n",
      " |      Implement delattr(self, name).\n",
      " |  \n",
      " |  __dir__(self)\n",
      " |      __dir__() -> list\n",
      " |      default dir() implementation\n",
      " |  \n",
      " |  __getattr__(self, name)\n",
      " |  \n",
      " |  __repr__(self)\n",
      " |      Return repr(self).\n",
      " |  \n",
      " |  __setattr__(self, name, value)\n",
      " |      Implement setattr(self, name, value).\n",
      " |  \n",
      " |  __setstate__(self, state)\n",
      " |  \n",
      " |  add_module(self, name, module)\n",
      " |      Adds a child module to the current module.\n",
      " |      \n",
      " |      The module can be accessed as an attribute using the given name.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (string): name of the child module. The child module can be\n",
      " |              accessed from this module using the given name\n",
      " |          module (Module): child module to be added to the module.\n",
      " |  \n",
      " |  apply(self, fn)\n",
      " |      Applies ``fn`` recursively to every submodule (as returned by ``.children()``)\n",
      " |      as well as self. Typical use includes initializing the parameters of a model\n",
      " |      (see also :ref:`nn-init-doc`).\n",
      " |      \n",
      " |      Args:\n",
      " |          fn (:class:`Module` -> None): function to be applied to each submodule\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> @torch.no_grad()\n",
      " |          >>> def init_weights(m):\n",
      " |          >>>     print(m)\n",
      " |          >>>     if type(m) == nn.Linear:\n",
      " |          >>>         m.weight.fill_(1.0)\n",
      " |          >>>         print(m.weight)\n",
      " |          >>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))\n",
      " |          >>> net.apply(init_weights)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 1.,  1.],\n",
      " |                  [ 1.,  1.]])\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 1.,  1.],\n",
      " |                  [ 1.,  1.]])\n",
      " |          Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          )\n",
      " |          Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          )\n",
      " |  \n",
      " |  bfloat16(self)\n",
      " |      Casts all floating point parameters and buffers to ``bfloat16`` datatype.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  buffers(self, recurse=True)\n",
      " |      Returns an iterator over module buffers.\n",
      " |      \n",
      " |      Args:\n",
      " |          recurse (bool): if True, then yields buffers of this module\n",
      " |              and all submodules. Otherwise, yields only buffers that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          torch.Tensor: module buffer\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for buf in model.buffers():\n",
      " |          >>>     print(type(buf), buf.size())\n",
      " |          <class 'torch.Tensor'> (20L,)\n",
      " |          <class 'torch.Tensor'> (20L, 1L, 5L, 5L)\n",
      " |  \n",
      " |  children(self)\n",
      " |      Returns an iterator over immediate children modules.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Module: a child module\n",
      " |  \n",
      " |  cpu(self)\n",
      " |      Moves all model parameters and buffers to the CPU.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  cuda(self, device=None)\n",
      " |      Moves all model parameters and buffers to the GPU.\n",
      " |      \n",
      " |      This also makes associated parameters and buffers different objects. So\n",
      " |      it should be called before constructing optimizer if the module will\n",
      " |      live on GPU while being optimized.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          device (int, optional): if specified, all parameters will be\n",
      " |              copied to that device\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  double(self)\n",
      " |      Casts all floating point parameters and buffers to ``double`` datatype.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  eval(self)\n",
      " |      Sets the module in evaluation mode.\n",
      " |      \n",
      " |      This has any effect only on certain modules. See documentations of\n",
      " |      particular modules for details of their behaviors in training/evaluation\n",
      " |      mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\n",
      " |      etc.\n",
      " |      \n",
      " |      This is equivalent with :meth:`self.train(False) <torch.nn.Module.train>`.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  extra_repr(self)\n",
      " |      Set the extra representation of the module\n",
      " |      \n",
      " |      To print customized extra information, you should reimplement\n",
      " |      this method in your own modules. Both single-line and multi-line\n",
      " |      strings are acceptable.\n",
      " |  \n",
      " |  float(self)\n",
      " |      Casts all floating point parameters and buffers to float datatype.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  half(self)\n",
      " |      Casts all floating point parameters and buffers to ``half`` datatype.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  load_state_dict(self, state_dict, strict=True)\n",
      " |      Copies parameters and buffers from :attr:`state_dict` into\n",
      " |      this module and its descendants. If :attr:`strict` is ``True``, then\n",
      " |      the keys of :attr:`state_dict` must exactly match the keys returned\n",
      " |      by this module's :meth:`~torch.nn.Module.state_dict` function.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          state_dict (dict): a dict containing parameters and\n",
      " |              persistent buffers.\n",
      " |          strict (bool, optional): whether to strictly enforce that the keys\n",
      " |              in :attr:`state_dict` match the keys returned by this module's\n",
      " |              :meth:`~torch.nn.Module.state_dict` function. Default: ``True``\n",
      " |      \n",
      " |      Returns:\n",
      " |          ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\n",
      " |              * **missing_keys** is a list of str containing the missing keys\n",
      " |              * **unexpected_keys** is a list of str containing the unexpected keys\n",
      " |  \n",
      " |  modules(self)\n",
      " |      Returns an iterator over all modules in the network.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Module: a module in the network\n",
      " |      \n",
      " |      Note:\n",
      " |          Duplicate modules are returned only once. In the following\n",
      " |          example, ``l`` will be returned only once.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> l = nn.Linear(2, 2)\n",
      " |          >>> net = nn.Sequential(l, l)\n",
      " |          >>> for idx, m in enumerate(net.modules()):\n",
      " |                  print(idx, '->', m)\n",
      " |      \n",
      " |          0 -> Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          )\n",
      " |          1 -> Linear(in_features=2, out_features=2, bias=True)\n",
      " |  \n",
      " |  named_buffers(self, prefix='', recurse=True)\n",
      " |      Returns an iterator over module buffers, yielding both the\n",
      " |      name of the buffer as well as the buffer itself.\n",
      " |      \n",
      " |      Args:\n",
      " |          prefix (str): prefix to prepend to all buffer names.\n",
      " |          recurse (bool): if True, then yields buffers of this module\n",
      " |              and all submodules. Otherwise, yields only buffers that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (string, torch.Tensor): Tuple containing the name and buffer\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for name, buf in self.named_buffers():\n",
      " |          >>>    if name in ['running_var']:\n",
      " |          >>>        print(buf.size())\n",
      " |  \n",
      " |  named_children(self)\n",
      " |      Returns an iterator over immediate children modules, yielding both\n",
      " |      the name of the module as well as the module itself.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (string, Module): Tuple containing a name and child module\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for name, module in model.named_children():\n",
      " |          >>>     if name in ['conv4', 'conv5']:\n",
      " |          >>>         print(module)\n",
      " |  \n",
      " |  named_modules(self, memo=None, prefix='')\n",
      " |      Returns an iterator over all modules in the network, yielding\n",
      " |      both the name of the module as well as the module itself.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (string, Module): Tuple of name and module\n",
      " |      \n",
      " |      Note:\n",
      " |          Duplicate modules are returned only once. In the following\n",
      " |          example, ``l`` will be returned only once.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> l = nn.Linear(2, 2)\n",
      " |          >>> net = nn.Sequential(l, l)\n",
      " |          >>> for idx, m in enumerate(net.named_modules()):\n",
      " |                  print(idx, '->', m)\n",
      " |      \n",
      " |          0 -> ('', Sequential(\n",
      " |            (0): Linear(in_features=2, out_features=2, bias=True)\n",
      " |            (1): Linear(in_features=2, out_features=2, bias=True)\n",
      " |          ))\n",
      " |          1 -> ('0', Linear(in_features=2, out_features=2, bias=True))\n",
      " |  \n",
      " |  named_parameters(self, prefix='', recurse=True)\n",
      " |      Returns an iterator over module parameters, yielding both the\n",
      " |      name of the parameter as well as the parameter itself.\n",
      " |      \n",
      " |      Args:\n",
      " |          prefix (str): prefix to prepend to all parameter names.\n",
      " |          recurse (bool): if True, then yields parameters of this module\n",
      " |              and all submodules. Otherwise, yields only parameters that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          (string, Parameter): Tuple containing the name and parameter\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for name, param in self.named_parameters():\n",
      " |          >>>    if name in ['bias']:\n",
      " |          >>>        print(param.size())\n",
      " |  \n",
      " |  parameters(self, recurse=True)\n",
      " |      Returns an iterator over module parameters.\n",
      " |      \n",
      " |      This is typically passed to an optimizer.\n",
      " |      \n",
      " |      Args:\n",
      " |          recurse (bool): if True, then yields parameters of this module\n",
      " |              and all submodules. Otherwise, yields only parameters that\n",
      " |              are direct members of this module.\n",
      " |      \n",
      " |      Yields:\n",
      " |          Parameter: module parameter\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> for param in model.parameters():\n",
      " |          >>>     print(type(param), param.size())\n",
      " |          <class 'torch.Tensor'> (20L,)\n",
      " |          <class 'torch.Tensor'> (20L, 1L, 5L, 5L)\n",
      " |  \n",
      " |  register_backward_hook(self, hook)\n",
      " |      Registers a backward hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time the gradients with respect to module\n",
      " |      inputs are computed. The hook should have the following signature::\n",
      " |      \n",
      " |          hook(module, grad_input, grad_output) -> Tensor or None\n",
      " |      \n",
      " |      The :attr:`grad_input` and :attr:`grad_output` may be tuples if the\n",
      " |      module has multiple inputs or outputs. The hook should not modify its\n",
      " |      arguments, but it can optionally return a new gradient with respect to\n",
      " |      input that will be used in place of :attr:`grad_input` in subsequent\n",
      " |      computations.\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |      \n",
      " |      .. warning ::\n",
      " |      \n",
      " |          The current implementation will not have the presented behavior\n",
      " |          for complex :class:`Module` that perform many operations.\n",
      " |          In some failure cases, :attr:`grad_input` and :attr:`grad_output` will only\n",
      " |          contain the gradients for a subset of the inputs and outputs.\n",
      " |          For such :class:`Module`, you should use :func:`torch.Tensor.register_hook`\n",
      " |          directly on a specific input or output to get the required gradients.\n",
      " |  \n",
      " |  register_buffer(self, name, tensor)\n",
      " |      Adds a persistent buffer to the module.\n",
      " |      \n",
      " |      This is typically used to register a buffer that should not to be\n",
      " |      considered a model parameter. For example, BatchNorm's ``running_mean``\n",
      " |      is not a parameter, but is part of the persistent state.\n",
      " |      \n",
      " |      Buffers can be accessed as attributes using given names.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (string): name of the buffer. The buffer can be accessed\n",
      " |              from this module using the given name\n",
      " |          tensor (Tensor): buffer to be registered.\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> self.register_buffer('running_mean', torch.zeros(num_features))\n",
      " |  \n",
      " |  register_forward_hook(self, hook)\n",
      " |      Registers a forward hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time after :func:`forward` has computed an output.\n",
      " |      It should have the following signature::\n",
      " |      \n",
      " |          hook(module, input, output) -> None or modified output\n",
      " |      \n",
      " |      The hook can modify the output. It can modify the input inplace but\n",
      " |      it will not have effect on forward since this is called after\n",
      " |      :func:`forward` is called.\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_forward_pre_hook(self, hook)\n",
      " |      Registers a forward pre-hook on the module.\n",
      " |      \n",
      " |      The hook will be called every time before :func:`forward` is invoked.\n",
      " |      It should have the following signature::\n",
      " |      \n",
      " |          hook(module, input) -> None or modified input\n",
      " |      \n",
      " |      The hook can modify the input. User can either return a tuple or a\n",
      " |      single modified value in the hook. We will wrap the value into a tuple\n",
      " |      if a single value is returned(unless that value is already a tuple).\n",
      " |      \n",
      " |      Returns:\n",
      " |          :class:`torch.utils.hooks.RemovableHandle`:\n",
      " |              a handle that can be used to remove the added hook by calling\n",
      " |              ``handle.remove()``\n",
      " |  \n",
      " |  register_parameter(self, name, param)\n",
      " |      Adds a parameter to the module.\n",
      " |      \n",
      " |      The parameter can be accessed as an attribute using given name.\n",
      " |      \n",
      " |      Args:\n",
      " |          name (string): name of the parameter. The parameter can be accessed\n",
      " |              from this module using the given name\n",
      " |          param (Parameter): parameter to be added to the module.\n",
      " |  \n",
      " |  requires_grad_(self, requires_grad=True)\n",
      " |      Change if autograd should record operations on parameters in this\n",
      " |      module.\n",
      " |      \n",
      " |      This method sets the parameters' :attr:`requires_grad` attributes\n",
      " |      in-place.\n",
      " |      \n",
      " |      This method is helpful for freezing part of the module for finetuning\n",
      " |      or training parts of a model individually (e.g., GAN training).\n",
      " |      \n",
      " |      Args:\n",
      " |          requires_grad (bool): whether autograd should record operations on\n",
      " |                                parameters in this module. Default: ``True``.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  share_memory(self)\n",
      " |  \n",
      " |  state_dict(self, destination=None, prefix='', keep_vars=False)\n",
      " |      Returns a dictionary containing a whole state of the module.\n",
      " |      \n",
      " |      Both parameters and persistent buffers (e.g. running averages) are\n",
      " |      included. Keys are corresponding parameter and buffer names.\n",
      " |      \n",
      " |      Returns:\n",
      " |          dict:\n",
      " |              a dictionary containing a whole state of the module\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> module.state_dict().keys()\n",
      " |          ['bias', 'weight']\n",
      " |  \n",
      " |  to(self, *args, **kwargs)\n",
      " |      Moves and/or casts the parameters and buffers.\n",
      " |      \n",
      " |      This can be called as\n",
      " |      \n",
      " |      .. function:: to(device=None, dtype=None, non_blocking=False)\n",
      " |      \n",
      " |      .. function:: to(dtype, non_blocking=False)\n",
      " |      \n",
      " |      .. function:: to(tensor, non_blocking=False)\n",
      " |      \n",
      " |      .. function:: to(memory_format=torch.channels_last)\n",
      " |      \n",
      " |      Its signature is similar to :meth:`torch.Tensor.to`, but only accepts\n",
      " |      floating point desired :attr:`dtype` s. In addition, this method will\n",
      " |      only cast the floating point parameters and buffers to :attr:`dtype`\n",
      " |      (if given). The integral parameters and buffers will be moved\n",
      " |      :attr:`device`, if that is given, but with dtypes unchanged. When\n",
      " |      :attr:`non_blocking` is set, it tries to convert/move asynchronously\n",
      " |      with respect to the host if possible, e.g., moving CPU Tensors with\n",
      " |      pinned memory to CUDA devices.\n",
      " |      \n",
      " |      See below for examples.\n",
      " |      \n",
      " |      .. note::\n",
      " |          This method modifies the module in-place.\n",
      " |      \n",
      " |      Args:\n",
      " |          device (:class:`torch.device`): the desired device of the parameters\n",
      " |              and buffers in this module\n",
      " |          dtype (:class:`torch.dtype`): the desired floating point type of\n",
      " |              the floating point parameters and buffers in this module\n",
      " |          tensor (torch.Tensor): Tensor whose dtype and device are the desired\n",
      " |              dtype and device for all parameters and buffers in this module\n",
      " |          memory_format (:class:`torch.memory_format`): the desired memory\n",
      " |              format for 4D parameters and buffers in this module (keyword\n",
      " |              only argument)\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |      \n",
      " |      Example::\n",
      " |      \n",
      " |          >>> linear = nn.Linear(2, 2)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1913, -0.3420],\n",
      " |                  [-0.5113, -0.2325]])\n",
      " |          >>> linear.to(torch.double)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1913, -0.3420],\n",
      " |                  [-0.5113, -0.2325]], dtype=torch.float64)\n",
      " |          >>> gpu1 = torch.device(\"cuda:1\")\n",
      " |          >>> linear.to(gpu1, dtype=torch.half, non_blocking=True)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1914, -0.3420],\n",
      " |                  [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')\n",
      " |          >>> cpu = torch.device(\"cpu\")\n",
      " |          >>> linear.to(cpu)\n",
      " |          Linear(in_features=2, out_features=2, bias=True)\n",
      " |          >>> linear.weight\n",
      " |          Parameter containing:\n",
      " |          tensor([[ 0.1914, -0.3420],\n",
      " |                  [-0.5112, -0.2324]], dtype=torch.float16)\n",
      " |  \n",
      " |  train(self, mode=True)\n",
      " |      Sets the module in training mode.\n",
      " |      \n",
      " |      This has any effect only on certain modules. See documentations of\n",
      " |      particular modules for details of their behaviors in training/evaluation\n",
      " |      mode, if they are affected, e.g. :class:`Dropout`, :class:`BatchNorm`,\n",
      " |      etc.\n",
      " |      \n",
      " |      Args:\n",
      " |          mode (bool): whether to set training mode (``True``) or evaluation\n",
      " |                       mode (``False``). Default: ``True``.\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  type(self, dst_type)\n",
      " |      Casts all parameters and buffers to :attr:`dst_type`.\n",
      " |      \n",
      " |      Arguments:\n",
      " |          dst_type (type or string): the desired type\n",
      " |      \n",
      " |      Returns:\n",
      " |          Module: self\n",
      " |  \n",
      " |  zero_grad(self)\n",
      " |      Sets gradients of all model parameters to zero.\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data descriptors inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  __dict__\n",
      " |      dictionary for instance variables (if defined)\n",
      " |  \n",
      " |  __weakref__\n",
      " |      list of weak references to the object (if defined)\n",
      " |  \n",
      " |  ----------------------------------------------------------------------\n",
      " |  Data and other attributes inherited from torch.nn.modules.module.Module:\n",
      " |  \n",
      " |  dump_patches = False\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(CosineSimilarity)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### V3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "from numba import jit, njit, prange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "@njit(parallel=True)\n",
    "def getSimilarWords(model, tokenizer, combinedOutputFolder, symptom, meanEmb, similarityThreshold = 0.3, numThreshold = 150000, numComp = 10000):\n",
    "        \n",
    "    output = []\n",
    "\n",
    "    symptomToken = tokenizer.encode(symptom)[1]\n",
    "\n",
    "    fileList = os.listdir(combinedOutputFolder)\n",
    "    \n",
    "    cos = CosineSimilarity(dim=1, eps=1e-6)\n",
    "\n",
    "    examineCount = 0\n",
    "\n",
    "    for i in prange(len(fileList)):\n",
    "\n",
    "        if examineCount >= numThreshold:\n",
    "            break\n",
    "\n",
    "\n",
    "        filename = os.path.join(combinedOutputFolder, f\"{i}.pkl\")\n",
    "        subDict = pickle.load(open(filename,'rb'))\n",
    "\n",
    "        IDList = subDict['id']\n",
    "        tokenList = subDict['token']\n",
    "        embList = subDict['emb']\n",
    "\n",
    "#         sim = np.round(cosine_similarity(embList, meanEmb.reshape(1,-1)).reshape(-1),4)\n",
    "\n",
    "        arrA = torch.from_numpy(meanEmb.reshape(1,-1)).cuda()\n",
    "        arrB = torch.from_numpy(embList).cuda()\n",
    "        \n",
    "        sim = cos(arrA,arrB).cpu().numpy().reshape(-1)\n",
    "        \n",
    "        sim = np.round(sim,4)\n",
    "\n",
    "        index= np.where([sim> similarityThreshold])[1]\n",
    "\n",
    "        tokenList_ = tokenList[index]\n",
    "        IDList_ = IDList[index]\n",
    "        simList = sim[index]\n",
    "\n",
    "        out = [(x,y,z) for x,y,z in zip(tokenList_, simList, IDList_)]\n",
    "\n",
    "        output += out\n",
    "\n",
    "        examineCount += numComp\n",
    "\n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "numComp = 10000\n",
    "numThreshold = 20000\n",
    "\n",
    "startTime= time.time()\n",
    "out2 = getSimilarWords(model, tokenizer, combinedOutputFolder, symptom, meanEmb, similarityThreshold = 0.3, numThreshold = numThreshold, numComp = numComp)\n",
    "\n",
    "print(f\"Time taken : {time.time() - startTime}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### V4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Client(n_workers=6, threads_per_worker=2, processes=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'http://128.91.252.35:8787/status'"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.dashboard_link"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getSimilarWords(model, tokenizer, combinedOutputFolder, symptom, meanEmb, similarityThreshold = 0.3, numThreshold = 150000, numComp = 10000):\n",
    "        \n",
    "    output = []\n",
    "\n",
    "    symptomToken = tokenizer.encode(symptom)[1]\n",
    "\n",
    "    fileList = os.listdir(combinedOutputFolder)\n",
    "    \n",
    "    cos = CosineSimilarity(dim=1, eps=1e-6)\n",
    "    \n",
    "    computeTaskPart_ = partial(computeTask_, symptom = '',combinedOutputFolder=combinedOutputFolder, \n",
    "                               meanEmb=meanEmb,similarityThreshold = similarityThreshold)\n",
    "\n",
    "    examineCount = 0\n",
    "    \n",
    "    totalSteps = int(numThreshold/numComp)\n",
    "    \n",
    "    mult = 10\n",
    "    \n",
    "    totalIters = totalSteps//mult if totalSteps%mult==0 else totalSteps//mult + 1\n",
    "    \n",
    "    \n",
    "    finalOutput = []\n",
    "    \n",
    "    for i in tqdm(range(totalIters)):\n",
    "        \n",
    "        start = i*mult\n",
    "        end = min((i+1)*mult, totalSteps)\n",
    "    \n",
    "        tasks = list(range(start,end))\n",
    "        \n",
    "        sent = client.map(computeTaskPart_, tasks)\n",
    "\n",
    "        result = client.gather(sent)\n",
    "\n",
    "        print(\"Tasks complete\")\n",
    "\n",
    "        result = list(itertools.chain.from_iterable(result))\n",
    "        \n",
    "        finalOutput.append(result)\n",
    "    \n",
    "    \n",
    "    finalOutput = list(itertools.chain.from_iterable(finalOutput))\n",
    "    \n",
    "    return finalOutput\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "def computeTask_(index, symptom, combinedOutputFolder,meanEmb, similarityThreshold):\n",
    "\n",
    "    symptomToken = tokenizer.encode(symptom)[1]\n",
    "\n",
    "    cos = CosineSimilarity(dim=1, eps=1e-6)\n",
    "\n",
    "    filename = os.path.join(combinedOutputFolder, f\"{index+6}.pkl\")\n",
    "    subDict = pickle.load(open(filename,'rb'))\n",
    "\n",
    "    IDList = subDict['id']\n",
    "    tokenList = subDict['token']\n",
    "    embList = subDict['emb']\n",
    "\n",
    "#         sim = np.round(cosine_similarity(embList, meanEmb.reshape(1,-1)).reshape(-1),4)\n",
    "\n",
    "    arrA = torch.from_numpy(meanEmb.reshape(1,-1))\n",
    "    arrB = torch.from_numpy(embList)\n",
    "\n",
    "#         arrA = torch.from_numpy(meanEmb.reshape(1,-1)).cuda()\n",
    "#         arrB = torch.from_numpy(embList).cuda()\n",
    "\n",
    "    sim = cos(arrA,arrB).cpu().numpy().reshape(-1)\n",
    "\n",
    "    sim = np.round(sim,4)\n",
    "\n",
    "    index= np.where([sim> similarityThreshold])[1]\n",
    "\n",
    "    tokenList_ = tokenList[index]\n",
    "    IDList_ = IDList[index]\n",
    "    simList = sim[index]\n",
    "\n",
    "    out = [(x,y,z) for x,y,z in zip(tokenList_, simList, IDList_)]\n",
    "\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [],
   "source": [
    "combinedOutputFolder = '/data2/roshansk/ADRModel_DataStore_10000/'\n",
    "numComp = 10000\n",
    "numThreshold = 250000\n",
    "\n",
    "\n",
    "file = 'fatigue_16342_Emb.npy'\n",
    "embList = np.load(os.path.join('EmbFolder/',file))\n",
    "meanEmb = np.mean(embList,0)\n",
    "symptom = ''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/3 [00:00<?, ?it/s]distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 17.29 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 17.32 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 17.35 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 17.38 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 17.41 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 17.43 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 91% memory usage. Pausing worker.  Process memory: 20.71 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 20.71 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 121% memory usage. Pausing worker.  Process memory: 27.28 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 27.28 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 114% memory usage. Pausing worker.  Process memory: 25.79 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 25.79 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 133% memory usage. Pausing worker.  Process memory: 30.00 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 30.00 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 153% memory usage. Pausing worker.  Process memory: 34.62 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 34.62 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 149% memory usage. Pausing worker.  Process memory: 33.61 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 33.61 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 54% memory usage. Resuming worker. Process memory: 12.19 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 54% memory usage. Resuming worker. Process memory: 12.21 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 54% memory usage. Resuming worker. Process memory: 12.22 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 54% memory usage. Resuming worker. Process memory: 12.24 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 54% memory usage. Resuming worker. Process memory: 12.24 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 54% memory usage. Resuming worker. Process memory: 12.25 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 15.95 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.40 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 17.38 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 17.81 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.23 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 80% memory usage. Pausing worker.  Process memory: 18.23 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 81% memory usage. Pausing worker.  Process memory: 18.24 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 81% memory usage. Pausing worker.  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 81% memory usage. Pausing worker.  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 81% memory usage. Pausing worker.  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 81% memory usage. Pausing worker.  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 18.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 69% memory usage. Resuming worker. Process memory: 15.62 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 69% memory usage. Resuming worker. Process memory: 15.62 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 69% memory usage. Resuming worker. Process memory: 15.58 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 69% memory usage. Resuming worker. Process memory: 15.58 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 69% memory usage. Resuming worker. Process memory: 15.54 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Worker is at 69% memory usage. Resuming worker. Process memory: 15.54 GB -- Worker memory limit: 22.52 GB\n",
      " 33%|███▎      | 1/3 [00:29<00:59, 29.57s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tasks complete\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.43 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.44 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.45 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.46 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.47 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.47 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.51 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.52 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.53 GB -- Worker memory limit: 22.52 GB\n",
      "distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 16.53 GB -- Worker memory limit: 22.52 GB\n",
      " 67%|██████▋   | 2/3 [02:32<00:57, 57.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tasks complete\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 3/3 [04:20<00:00, 86.98s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tasks complete\n",
      "Time taken : 260.9436800479889\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "startTime = time.time()\n",
    "\n",
    "result = getSimilarWords(model, tokenizer, combinedOutputFolder, symptom, \n",
    "                meanEmb, similarityThreshold = 0.3, numThreshold = numThreshold, numComp = numComp)\n",
    "\n",
    "print(f\"Time taken : {time.time() - startTime}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "distributed.client - ERROR - Failed to reconnect to scheduler after 10.00 seconds, closing client\n",
      "_GatheringFuture exception was never retrieved\n",
      "future: <_GatheringFuture finished exception=CancelledError()>\n",
      "concurrent.futures._base.CancelledError\n"
     ]
    }
   ],
   "source": [
    "client.shutdown()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4.333333333333333"
      ]
     },
     "execution_count": 80,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "260/60"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "25.98"
      ]
     },
     "execution_count": 81,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "4.33*6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "distributed.client - ERROR - Failed to reconnect to scheduler after 10.00 seconds, closing client\n",
      "_GatheringFuture exception was never retrieved\n",
      "future: <_GatheringFuture finished exception=CancelledError()>\n",
      "concurrent.futures._base.CancelledError\n"
     ]
    }
   ],
   "source": [
    "client.shutdown()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### New ADR Model  v4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": [
    "\n",
    "def computeTask_(index, symptom, combinedOutputFolder,meanEmb, similarityThreshold):\n",
    "\n",
    "    symptomToken = tokenizer.encode(symptom)[1]\n",
    "\n",
    "    cos = CosineSimilarity(dim=1, eps=1e-6)\n",
    "\n",
    "    filename = os.path.join(combinedOutputFolder, f\"{index+6}.pkl\")\n",
    "    subDict = pickle.load(open(filename,'rb'))\n",
    "\n",
    "    IDList = subDict['id']\n",
    "    tokenList = subDict['token']\n",
    "    embList = subDict['emb']\n",
    "\n",
    "#         sim = np.round(cosine_similarity(embList, meanEmb.reshape(1,-1)).reshape(-1),4)\n",
    "\n",
    "    arrA = torch.from_numpy(meanEmb.reshape(1,-1))\n",
    "    arrB = torch.from_numpy(embList)\n",
    "\n",
    "#         arrA = torch.from_numpy(meanEmb.reshape(1,-1)).cuda()\n",
    "#         arrB = torch.from_numpy(embList).cuda()\n",
    "\n",
    "    sim = cos(arrA,arrB).cpu().numpy().reshape(-1)\n",
    "\n",
    "    sim = np.round(sim,4)\n",
    "\n",
    "    index= np.where([sim> similarityThreshold])[1]\n",
    "\n",
    "    tokenList_ = tokenList[index]\n",
    "    IDList_ = IDList[index]\n",
    "    simList = sim[index]\n",
    "\n",
    "    out = [(x,y,z) for x,y,z in zip(tokenList_, simList, IDList_)]\n",
    "\n",
    "    return out\n",
    "\n",
    "\n",
    "class ADRModel(object):\n",
    "    \n",
    "    def __init__(self, df, model, tokenizer, graph, outputFolder, combinedOutputFolder, modelOutputFolder = './', queue=None, useMasterEmb = False, \n",
    "                 masterContrib = 0.5, embeddingType='last4sum',\n",
    "                 numThreshold= 10000, saveEveryDepth = False,\n",
    "                numComp = 10000):\n",
    "        \n",
    "        self.df = df\n",
    "        self.model = model\n",
    "        self.tokenizer = tokenizer\n",
    "        self.graph = graph\n",
    "        self.outputFolder = outputFolder\n",
    "        self.combinedOutputFolder = combinedOutputFolder\n",
    "        self.embeddingType = embeddingType\n",
    "        self.numThreshold = numThreshold\n",
    "        self.saveEveryDepth = saveEveryDepth\n",
    "        self.modelOutputFolder = modelOutputFolder\n",
    "        self.numComp = numComp\n",
    "        \n",
    "        if queue is None:\n",
    "            self.q = deque()\n",
    "        else:\n",
    "            self.q = queue\n",
    "            \n",
    "            \n",
    "        self.masterEmb = None\n",
    "        \n",
    "        self.useMasterEmb = useMasterEmb\n",
    "        self.masterContrib = masterContrib\n",
    "        \n",
    "        self.masterEmbList = []\n",
    "        \n",
    "        self.client = Client(n_workers=6, threads_per_worker=2, processes=False)\n",
    "        \n",
    "        self.generateStates()\n",
    "        \n",
    "        \n",
    "    def generateStates(self):\n",
    "        \n",
    "        \n",
    "        for i in tqdm(range(len(self.df))):\n",
    "            \n",
    "            if os.path.exists(os.path.join(self.outputFolder, f\"{i}.msh\")):\n",
    "                continue\n",
    "\n",
    "\n",
    "            tokens = self.tokenizer.encode(self.df.iloc[i]['message'].lower())\n",
    "            decoded = self.tokenizer.decode(tokens).split(\" \")\n",
    "            logits, hidden_states = self.model(torch.Tensor(tokens).unsqueeze(0).long())\n",
    "\n",
    "            hidden_states = torch.stack(hidden_states).squeeze(1).permute(1,0,2)\n",
    "\n",
    "            \n",
    "            if self.embeddingType == 'last4sum':\n",
    "                embedding = torch.sum(hidden_states[:,9:13,:],1)\n",
    "            elif self.embeddingType =='last4concat':\n",
    "                embedding = hidden_states[tokenIndex,9:13,:].reshape(-1)\n",
    "            elif self.embeddingType == 'secondlast':\n",
    "                embedding = hidden_states[tokenIndex,-2,:]\n",
    "            else:\n",
    "                embedding = hidden_states[tokenIndex,-1,:]\n",
    "                    \n",
    "                    \n",
    "            embedding = embedding.detach().cpu().numpy()\n",
    "            \n",
    "            marshal.dump(embedding.tolist(), open(os.path.join(self.outputFolder, f\"{i}.msh\"), 'wb'))\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "    def getSymptomEmbedding(self, symptom, subset = None):\n",
    "    \n",
    "        embeddingList = []\n",
    "        messageList = []\n",
    "\n",
    "#         if subset is not None:\n",
    "#             self.df = self.df.iloc[subset]\n",
    "\n",
    "#         if type(df) == pd.Series:\n",
    "#             self.df = pd.DataFrame(self.df).T\n",
    "\n",
    "#         symptomToken = self.tokenizer.encode(symptom)[1]\n",
    "        symptomToken = self.tokenizer.convert_tokens_to_ids(symptom)\n",
    "\n",
    "        for i in range(len(self.df)):\n",
    "\n",
    "            if symptomToken in self.tokenizer.encode(self.df.iloc[i]['message'].lower()):\n",
    "\n",
    "                tokens = self.tokenizer.encode(self.df.iloc[i]['message'].lower())\n",
    "                decoded = self.tokenizer.decode(tokens).split(\" \")\n",
    "\n",
    "                hidden_states = np.array(marshal.load( open(os.path.join(self.outputFolder, f\"{i}.msh\"), 'rb') ))\n",
    "\n",
    "                try:\n",
    "                    tokenIndex = tokens.index(symptomToken)\n",
    "                except:\n",
    "                    a= 1\n",
    "                    continue\n",
    "\n",
    " \n",
    "                embedding = hidden_states[tokenIndex,:]\n",
    "\n",
    "                embeddingList.append(embedding)\n",
    "                messageList.append(self.df.iloc[i]['message'].lower())\n",
    "\n",
    "                if len(embeddingList)==30:\n",
    "                    break\n",
    "\n",
    "\n",
    "\n",
    "        return embeddingList, messageList\n",
    "    \n",
    "    def getSimilarWords(self, symptom, meanEmb, similarityThreshold = 0.3):\n",
    "        \n",
    "        output = []\n",
    "\n",
    "        symptomToken = self.tokenizer.encode(symptom)[1]\n",
    "\n",
    "        fileList = os.listdir(self.combinedOutputFolder)\n",
    "\n",
    "        cos = CosineSimilarity(dim=1, eps=1e-6)\n",
    "\n",
    "        computeTaskPart_ = partial(computeTask_, symptom = symptom,combinedOutputFolder=self.combinedOutputFolder, \n",
    "                                   meanEmb=meanEmb,similarityThreshold = similarityThreshold)\n",
    "\n",
    "        examineCount = 0\n",
    "\n",
    "        totalSteps = int(self.numThreshold/self.numComp)\n",
    "\n",
    "        mult = 10\n",
    "\n",
    "        totalIters = totalSteps//mult if totalSteps%mult==0 else totalSteps//mult + 1\n",
    "\n",
    "\n",
    "        finalOutput = []\n",
    "\n",
    "        for i in tqdm(range(totalIters)):\n",
    "\n",
    "            start = i*mult\n",
    "            end = min((i+1)*mult, totalSteps)\n",
    "\n",
    "            tasks = list(range(start,end))\n",
    "\n",
    "            sent = self.client.map(computeTaskPart_, tasks)\n",
    "\n",
    "            result = self.client.gather(sent)\n",
    "\n",
    "            print(\"Tasks complete\")\n",
    "\n",
    "            result = list(itertools.chain.from_iterable(result))\n",
    "\n",
    "            finalOutput.append(result)\n",
    "\n",
    "\n",
    "        finalOutput = list(itertools.chain.from_iterable(finalOutput))\n",
    "\n",
    "        return finalOutput\n",
    "\n",
    "        \n",
    "    \n",
    "    \n",
    "    def getOutput(self, out):\n",
    "    \n",
    "        output = out\n",
    "\n",
    "        outMap = {}\n",
    "\n",
    "        for i in range(len(output)):\n",
    "            if output[i][0] in outMap:\n",
    "                outMap[output[i][0]].append(output[i][1])\n",
    "            else:\n",
    "                outMap[output[i][0]] = [output[i][1]]\n",
    "\n",
    "\n",
    "        outMap_ = {}\n",
    "\n",
    "        for i in range(len(output)):\n",
    "            if output[i][0] in outMap_:\n",
    "                outMap_[output[i][0]].append(output[i][2])\n",
    "            else:\n",
    "                outMap_[output[i][0]] = [output[i][2]]\n",
    "\n",
    "\n",
    "        outputDf = []\n",
    "\n",
    "        for key in outMap.keys():\n",
    "            length = len(outMap[key])\n",
    "            mean = np.mean(outMap[key])\n",
    "\n",
    "            outputDf.append([key, length, mean])\n",
    "\n",
    "        outputDf = pd.DataFrame(outputDf)\n",
    "        outputDf.columns = ['word','counts','mean_sim']\n",
    "        outputDf = outputDf.sort_values('mean_sim', ascending=False)\n",
    "\n",
    "        return outputDf, outMap, outMap_\n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    def exploreNode(self, word, depth, maxDepth = 3, topk = 5):\n",
    "\n",
    "    \n",
    "        self.graph.addNode(word,0,depth)\n",
    "\n",
    "        print(f\"Depth : {depth} Exploring {word}\")\n",
    "\n",
    "        if depth == maxDepth:\n",
    "            print(\"Reached max depth\")\n",
    "            return\n",
    "\n",
    "        keyWord = word\n",
    "\n",
    "        token = self.tokenizer.encode(keyWord)[1]\n",
    "\n",
    "        if self.graph[word].vector is None:\n",
    "\n",
    "            inEdgeList = self.graph[word].edges_in\n",
    "\n",
    "            if len(inEdgeList)==0:\n",
    "                textIDList = None\n",
    "            else:\n",
    "                textIDList = []\n",
    "\n",
    "                for edge in inEdgeList:\n",
    "                    textIDList.append(self.graph.edgeList[edge].textID)\n",
    "\n",
    "                textIDList = list(set(list(itertools.chain.from_iterable(textIDList))))\n",
    "\n",
    "            \n",
    "            embList,msgList = self.getSymptomEmbedding(keyWord, subset = textIDList)\n",
    "\n",
    "            meanEmb = np.array(embList)\n",
    "            meanEmb = np.mean(meanEmb,0)\n",
    "\n",
    "\n",
    "            self.graph[word].vector = meanEmb\n",
    "            \n",
    "            if self.masterEmb is None:\n",
    "                self.masterEmb = meanEmb\n",
    "            \n",
    "            dist = getCosineDist(meanEmb, self.masterEmb)\n",
    "            \n",
    "            self.graph[word].masterDist = dist\n",
    "\n",
    "        else:\n",
    "            meanEmb = self.graph[word].vector\n",
    "            \n",
    "            if self.masterEmb is None:\n",
    "                self.masterEmb = meanEmb\n",
    "                \n",
    "            dist = getCosineDist(meanEmb, self.masterEmb)\n",
    "            \n",
    "            self.graph[word].masterDist = dist\n",
    "\n",
    "\n",
    "        symptom_ =''\n",
    "        embList_ = meanEmb\n",
    "\n",
    "        if self.useMasterEmb:\n",
    "            \n",
    "            finalEmb = self.masterContrib*self.masterEmb + (1 - self.masterContrib)*meanEmb\n",
    "            \n",
    "            out = self.getSimilarWords( symptom_, finalEmb , similarityThreshold = 0.3)\n",
    "        else:\n",
    "            out = self.getSimilarWords( symptom_, meanEmb, similarityThreshold = 0.3)\n",
    "\n",
    "        outputDf, outMap, outMap_ = self.getOutput(out)\n",
    "\n",
    "        outputDf = outputDf[outputDf.word!=keyWord]\n",
    "    #     outputDf = outputDf[~outputDf.word.isin(list(graph.wordMap.keys()))]\n",
    "        outputDf = outputDf.sort_values('mean_sim', ascending=False)\n",
    "        outputDf = outputDf.head(topk)\n",
    "\n",
    "        outputDf = outputDf[outputDf.mean_sim>0.4]\n",
    "\n",
    "        print(outputDf)\n",
    "        print(\"-----------------------\")\n",
    "\n",
    "        for i in range(len(outputDf)):\n",
    "\n",
    "            word = outputDf.iloc[i]['word']\n",
    "            numCount = outputDf.iloc[i]['counts']\n",
    "            weight = outputDf.iloc[i]['mean_sim']\n",
    "            textIDs = outMap_[word]\n",
    "\n",
    "            wordList = set(self.graph.wordMap.keys())\n",
    "\n",
    "            self.graph.addNode(word,0,depth+1)\n",
    "            self.graph[word].textIDList.append(textIDs)\n",
    "            self.graph.addEdge(keyWord, word, numCount, weight, textIDs)\n",
    "\n",
    "            if word in wordList:\n",
    "                continue\n",
    "\n",
    "#             if \"#\" in word:\n",
    "#                 continue\n",
    "\n",
    "\n",
    "            self.q.append((word, depth+1))\n",
    "            \n",
    "            \n",
    "    def trainModel(self, maxDepth = 3, topk = 5):\n",
    "        \n",
    "        currDepth = 0\n",
    "        \n",
    "        while len(self.q)>0:\n",
    "            token, depth = self.q.popleft()\n",
    "            \n",
    "            if depth> currDepth:\n",
    "                \n",
    "                if self.saveEveryDepth:\n",
    "                    filepath = os.path.join( self.modelOutputFolder, f\"depth_{currDepth}.pkl\")\n",
    "                    self.saveModel(filepath)\n",
    "                \n",
    "                self.masterEmbList.append(self.masterEmb.copy())\n",
    "                self.getMeanEmbedding(depth-1)\n",
    "                currDepth += 1\n",
    "            \n",
    "            self.exploreNode(word = token, depth = depth, maxDepth=maxDepth, topk=topk)\n",
    "        \n",
    "        #Saving final model\n",
    "        filepath = os.path.join(self.modelOutputFolder, \"final.pkl\")\n",
    "        self.saveModel(filepath)\n",
    "\n",
    "\n",
    "            \n",
    "    def getMeanEmbedding(self, depth, topk = 3):\n",
    "        \n",
    "        candidates = self.graph.depthMap[depth]\n",
    "        \n",
    "        vals = [self.graph[x].masterDist for x in candidates]\n",
    "        \n",
    "        vals = [(x,y) for x,y in zip(candidates,vals)]\n",
    "        \n",
    "        vals = sorted(vals, key = lambda x : -x[1])\n",
    "        \n",
    "        meanEmb = self.masterEmb\n",
    "        \n",
    "        selectedWords = []\n",
    "        for i in range(min(topk, len(vals)) ):\n",
    "            meanEmb += self.graph[ vals[i][0] ].vector\n",
    "            selectedWords.append(vals[i][0])\n",
    "            \n",
    "        meanEmb = meanEmb/(topk+1)\n",
    "        \n",
    "        self.masterEmb = meanEmb\n",
    "        \n",
    "        for i in range(len(selectedWords)):\n",
    "            print(selectedWords[i])\n",
    "        print(\"Master Embedding updated.\")\n",
    "        print(\"-----------------\")\n",
    "        \n",
    "        \n",
    "    \n",
    "    def plotGraph(self):\n",
    "        \n",
    "        edgeList, nodeList, nodeValues, nodeCount, nodeText, nodeSize = getGraphComponents(self.graph)\n",
    "\n",
    "        G=nx.Graph()\n",
    "\n",
    "        G.add_nodes_from(nodeList)\n",
    "        G.add_edges_from(edgeList)\n",
    "\n",
    "        edge_trace, node_trace1, node_trace = getPlotlyComponents(G, nodeList, nodeSize, nodeValues, nodeText)\n",
    "\n",
    "\n",
    "        fig = go.Figure(data=[edge_trace, node_trace1, node_trace],\n",
    "             layout=go.Layout(\n",
    "                title='<br>Network graph made with Python',\n",
    "                titlefont_size=16,\n",
    "                showlegend=False,\n",
    "                hovermode='closest',\n",
    "                margin=dict(b=20,l=5,r=5,t=50),\n",
    "\n",
    "                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),\n",
    "                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False))\n",
    "                )\n",
    "        \n",
    "        fig.show()\n",
    "        \n",
    "        \n",
    "    def saveModel(self,filename):\n",
    "        \n",
    "        classDict = self.__dict__.copy()\n",
    "        classDict.pop('model')\n",
    "        classDict.pop('tokenizer')\n",
    "        classDict.pop('df')\n",
    "        \n",
    "        pickle.dump( classDict, open( filename, \"wb\" ) )\n",
    "        \n",
    "        \n",
    "    def loadModel(self, filename):\n",
    "        \n",
    "        classDict = pickle.load(open(filename, 'rb'))\n",
    "        \n",
    "        for key in list(classDict.keys()):\n",
    "            self.__dict__[key] = classDict[key]\n",
    "        \n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hidden": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "env"
  },
  "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.6.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}