Embedding_Analysis.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c7d28558",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "from transformers import GPT2Config, GPT2Tokenizer, BertModel, BertTokenizer, DistilBertModel, DistilBertTokenizer\n",
    "from transformers import AdamW, get_linear_schedule_with_warmup\n",
    "\n",
    "from utils.InductiveAttentionModels import GPT2InductiveAttentionHeadModel\n",
    "from utils.SequenceCrossEntropyLoss import SequenceCrossEntropyLoss\n",
    "\n",
    "import numpy as np\n",
    "import math\n",
    "import random\n",
    "import time\n",
    "import tqdm\n",
    "\n",
    "from annoy import AnnoyIndex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f896cdce",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_projector.weight', 'vocab_transform.weight', 'vocab_projector.bias', 'vocab_layer_norm.bias', 'vocab_layer_norm.weight', 'vocab_transform.bias']\n",
      "- This IS expected if you are initializing DistilBertModel 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 DistilBertModel 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 the model checkpoint at distilbert-base-uncased were not used when initializing DistilBertModel: ['vocab_projector.weight', 'vocab_transform.weight', 'vocab_projector.bias', 'vocab_layer_norm.bias', 'vocab_layer_norm.weight', 'vocab_transform.bias']\n",
      "- This IS expected if you are initializing DistilBertModel 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 DistilBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "50261\n"
     ]
    }
   ],
   "source": [
    "bert_tokenizer = DistilBertTokenizer.from_pretrained(\"distilbert-base-uncased\")\n",
    "bert_model_recall = DistilBertModel.from_pretrained('distilbert-base-uncased')\n",
    "bert_model_rerank = DistilBertModel.from_pretrained('distilbert-base-uncased')\n",
    "# bert_tokenizer = BertTokenizer.from_pretrained(\"bert-base-uncased\")\n",
    "# bert_model = BertModel.from_pretrained('bert-base-uncased')\n",
    "gpt_tokenizer = GPT2Tokenizer.from_pretrained(\"gpt2\")\n",
    "gpt2_model = GPT2InductiveAttentionHeadModel.from_pretrained('gpt2')\n",
    "\n",
    "# REC_TOKEN = \"R\"\n",
    "# REC_END_TOKEN = \"E\"\n",
    "REC_TOKEN = \"[REC]\"\n",
    "REC_END_TOKEN = \"[REC_END]\"\n",
    "SEP_TOKEN = \"[SEP]\"\n",
    "PLACEHOLDER_TOKEN = \"[MOVIE_ID]\"\n",
    "gpt_tokenizer.add_tokens([REC_TOKEN, REC_END_TOKEN, SEP_TOKEN, PLACEHOLDER_TOKEN])\n",
    "gpt2_model.resize_token_embeddings(len(gpt_tokenizer)) \n",
    "original_token_emb_size = gpt2_model.get_input_embeddings().weight.shape[0]\n",
    "print(original_token_emb_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a1ea96bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MovieRecDataset(Dataset):\n",
    "    def __init__(self, data, bert_tok, gpt2_tok):\n",
    "        self.data = data\n",
    "        self.bert_tok = bert_tok\n",
    "        self.gpt2_tok = gpt2_tok\n",
    "        self.turn_ending = torch.tensor([[628, 198]]) # end of turn, '\\n\\n\\n'\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        dialogue = self.data[index]\n",
    "        \n",
    "        dialogue_tokens = []\n",
    "        \n",
    "        for utterance, gt_ind in dialogue:\n",
    "            utt_tokens = self.gpt2_tok(utterance, return_tensors=\"pt\")['input_ids']\n",
    "            dialogue_tokens.append( ( torch.cat( (utt_tokens, self.turn_ending), dim=1), gt_ind) )\n",
    "            \n",
    "        role_ids = None\n",
    "        previous_role_ids = None\n",
    "        if role_ids == None:\n",
    "            role_ids = [ 0 if item[0] == 'B' else 1 for item, _ in dialogue]\n",
    "            previous_role_ids = role_ids\n",
    "        else:\n",
    "            role_ids = [ 0 if item[0] == 'B' else 1 for item, _ in dialogue]\n",
    "            if not np.array_equal(role_ids, previous_role_ids):\n",
    "                raise Exception(\"Role ids dont match between languages\")\n",
    "            previous_role_ids = role_ids\n",
    "        \n",
    "        return role_ids, dialogue_tokens\n",
    "    \n",
    "    def collate(self, unpacked_data):\n",
    "        return unpacked_data\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "21a0ea69",
   "metadata": {},
   "outputs": [],
   "source": [
    "items_db = torch.load('/local-scratch1/data/by2299/redial_full_movie_db_placeholder')\n",
    "\n",
    "def sample_ids_from_db(item_db,\n",
    "                       gt_id, # ground truth id\n",
    "                       num_samples, # num samples to return\n",
    "                       include_gt # if we want gt_id to be included\n",
    "                      ):\n",
    "    ids_2_sample_from = list(item_db.keys())\n",
    "    ids_2_sample_from.remove(gt_id)\n",
    "    if include_gt:\n",
    "        results = random.sample(ids_2_sample_from, num_samples-1)\n",
    "        results.append(gt_id)\n",
    "    else:\n",
    "        results = random.sample(ids_2_sample_from, num_samples)\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "59b56796",
   "metadata": {},
   "outputs": [],
   "source": [
    "class UniversalCRSModel(torch.nn.Module):\n",
    "    def __init__(self, \n",
    "                 language_model, # backbone of Pretrained LM such as GPT2\n",
    "                 encoder, # backbone of item encoder such as bert\n",
    "                 recall_encoder,\n",
    "                 lm_tokenizer, # language model tokenizer\n",
    "                 encoder_tokenizer, # item encoder tokenizer\n",
    "                 device, # Cuda device\n",
    "                 items_db, # {id:info}, information of all items to be recommended\n",
    "                 annoy_base_recall=None, # annoy index base of encoded recall embeddings of items\n",
    "                 annoy_base_rerank=None, # annoy index base of encoded rerank embeddings of items, for validation and inference only\n",
    "                 recall_item_dim=768, # dimension of each item to be stored in annoy base\n",
    "                 lm_trim_offset=100, # offset to trim language model wte inputs length = (1024-lm_trim_offset)\n",
    "                 rec_token_str=\"[REC]\", # special token indicating recommendation and used for recall\n",
    "                 rec_end_token_str=\"[REC_END]\", # special token indicating recommendation ended, conditional generation starts\n",
    "                 sep_token_str=\"[SEP]\",\n",
    "                 placeholder_token_str=\"[MOVIE_ID]\"\n",
    "                ):\n",
    "        super(UniversalCRSModel, self).__init__()\n",
    "        \n",
    "        #models and tokenizers\n",
    "        self.language_model = language_model\n",
    "        self.item_encoder = encoder\n",
    "        self.recall_encoder = recall_encoder\n",
    "        self.lm_tokenizer = lm_tokenizer\n",
    "        self.encoder_tokenizer = encoder_tokenizer\n",
    "        self.device = device\n",
    "        \n",
    "        # item db and annoy index base\n",
    "        self.items_db = items_db\n",
    "        self.annoy_base_recall = annoy_base_recall\n",
    "        self.annoy_base_rerank = annoy_base_rerank\n",
    "        \n",
    "        # hyperparameters\n",
    "        self.recall_item_dim = recall_item_dim\n",
    "        self.lm_trim_offset = lm_trim_offset\n",
    "        \n",
    "        #constants\n",
    "        self.REC_TOKEN_STR = rec_token_str\n",
    "        self.REC_END_TOKEN_STR = rec_end_token_str\n",
    "        self.SEP_TOKEN_STR = sep_token_str\n",
    "        self.PLACEHOLDER_TOKEN_STR = placeholder_token_str\n",
    "        \n",
    "        # map language model hidden states to a vector to query annoy-item-base for recall\n",
    "        self.recall_lm_query_mapper = torch.nn.Linear(self.language_model.config.n_embd, self.recall_item_dim) # default [768,768]\n",
    "        # map output of self.item_encoder to vectors to be stored in annoy-item-base \n",
    "        self.recall_item_wte_mapper = torch.nn.Linear(self.recall_encoder.config.hidden_size, self.recall_item_dim) # default [768,768]\n",
    "        # map output of self.item_encoder to a wte of self.language_model\n",
    "        self.rerank_item_wte_mapper = torch.nn.Linear(self.item_encoder.config.hidden_size, self.language_model.config.n_embd) # default [768,768]\n",
    "        # map language model hidden states of item wte to a one digit logit for softmax computation\n",
    "        self.rerank_logits_mapper = torch.nn.Linear(self.language_model.config.n_embd, 1) # default [768,1]\n",
    "    \n",
    "    def get_sep_token_wtes(self):\n",
    "        sep_token_input_ids = self.lm_tokenizer(self.SEP_TOKEN_STR, return_tensors=\"pt\")[\"input_ids\"].to(self.device)\n",
    "        return self.language_model.transformer.wte(sep_token_input_ids) # [1, 1, self.language_model.config.n_embd]\n",
    "\n",
    "    def get_rec_token_wtes(self):\n",
    "        rec_token_input_ids = self.lm_tokenizer(self.REC_TOKEN_STR, return_tensors=\"pt\")[\"input_ids\"].to(self.device)\n",
    "        return self.language_model.transformer.wte(rec_token_input_ids) # [1, 1, self.language_model.config.n_embd]\n",
    "    \n",
    "    def get_rec_end_token_wtes(self):\n",
    "        rec_end_token_input_ids = self.lm_tokenizer(self.REC_END_TOKEN_STR, return_tensors=\"pt\")[\"input_ids\"].to(self.device)\n",
    "        return self.language_model.transformer.wte(rec_end_token_input_ids) # [1, 1, self.language_model.config.n_embd]\n",
    "    \n",
    "    def get_movie_title(self, m_id):\n",
    "        title = self.items_db[m_id]\n",
    "        title = title.split('[SEP]')[0].strip()\n",
    "        return title\n",
    "    \n",
    "    # compute BERT encoded item hidden representation\n",
    "    # output can be passed to self.recall_item_wte_mapper or self.rerank_item_wte_mapper\n",
    "    def compute_encoded_embeddings_for_items(self, \n",
    "                                             encoder_to_use,\n",
    "                                             item_ids, # an array of ids, single id should be passed as [id]\n",
    "                                             items_db_to_use # item databse to use\n",
    "                                            ):\n",
    "        chunk_ids = item_ids\n",
    "        chunk_infos = [items_db_to_use[key] for key in chunk_ids ]\n",
    "        chunk_tokens = self.encoder_tokenizer(chunk_infos, padding=True, truncation=True, return_tensors=\"pt\")\n",
    "        chunk_input_ids = chunk_tokens['input_ids'].to(self.device)\n",
    "        chunk_attention_mask = chunk_tokens['attention_mask'].to(self.device)\n",
    "        chunk_hiddens = encoder_to_use(input_ids=chunk_input_ids, attention_mask=chunk_attention_mask).last_hidden_state\n",
    "        # TODO: analyze whether averaging or taking [CLS] gives better result.\n",
    "#             chunk_pooled = torch.mean(chunk_hiddens, dim=1).cpu().detach().numpy()\n",
    "#             chunk_pooled = chunk_hiddens[:,0,:].cpu().detach().numpy()\n",
    "\n",
    "        # average of non-padding tokens\n",
    "        expanded_mask_size = list(chunk_attention_mask.size())\n",
    "        expanded_mask_size.append(encoder_to_use.config.hidden_size)\n",
    "        expanded_mask = chunk_attention_mask.unsqueeze(-1).expand(expanded_mask_size)\n",
    "        chunk_masked = torch.mul(chunk_hiddens, expanded_mask) # [num_example, len, 768]\n",
    "        chunk_pooled = torch.sum(chunk_masked, dim=1) / torch.sum(chunk_attention_mask, dim=1).unsqueeze(-1)\n",
    "        \n",
    "        # [len(item_ids), encoder_to_use.config.hidden_size], del chunk_hiddens to free up GPU memory\n",
    "        return chunk_pooled, chunk_hiddens\n",
    "    \n",
    "    def annoy_base_constructor_emb_analysis(self, items_db=None, distance_type='angular', chunk_size=50, n_trees=10):\n",
    "        items_db_to_use = self.items_db if items_db == None else items_db\n",
    "        all_item_ids = list(items_db_to_use.keys())\n",
    "        \n",
    "        total_pooled = []\n",
    "        # break into chunks/batches for model concurrency\n",
    "        num_chunks = math.ceil(len(all_item_ids) / chunk_size)\n",
    "        for i in range(num_chunks):\n",
    "            chunk_ids = all_item_ids[i*chunk_size: (i+1)*chunk_size]\n",
    "            chunk_pooled, chunk_hiddens = self.compute_encoded_embeddings_for_items(self.item_encoder, chunk_ids, items_db_to_use)\n",
    "#             chunk_pooled = chunk_hiddens[:,0,:].cpu().detach().numpy()\n",
    "            chunk_pooled = chunk_pooled.cpu().detach().numpy()\n",
    "            del chunk_hiddens\n",
    "            total_pooled.append(chunk_pooled)\n",
    "        total_pooled = np.concatenate(total_pooled, axis=0)\n",
    "        \n",
    "        pooled_tensor = torch.tensor(total_pooled).to(self.device)\n",
    "        \n",
    "        #build recall annoy index\n",
    "#         annoy_base_recall = AnnoyIndex(self.recall_item_wte_mapper.out_features, distance_type)\n",
    "#         pooled_recall = self.recall_item_wte_mapper(pooled_tensor) # [len(items_db_to_use), self.recall_item_wte_mapper.out_features]\n",
    "#         pooled_recall = pooled_recall.cpu().detach().numpy()\n",
    "#         for i, vector in zip(all_item_ids, pooled_recall):\n",
    "#             annoy_base_recall.add_item(i, vector)\n",
    "#         annoy_base_recall.build(n_trees)\n",
    "        \n",
    "        #build rerank annoy index, for validation and inference only\n",
    "        annoy_base_rerank = AnnoyIndex(self.rerank_item_wte_mapper.out_features, distance_type)\n",
    "#         pooled_rerank = self.rerank_item_wte_mapper(pooled_tensor) # [len(items_db_to_use), self.recall_item_wte_mapper.out_features]\n",
    "        pooled_rerank = pooled_tensor.cpu().detach().numpy()\n",
    "        for i, vector in zip(all_item_ids, pooled_rerank):\n",
    "            annoy_base_rerank.add_item(i, vector)\n",
    "        annoy_base_rerank.build(n_trees)\n",
    "        \n",
    "        del pooled_tensor\n",
    "        \n",
    "#         self.annoy_base_recall = annoy_base_recall\n",
    "        self.annoy_base_rerank = annoy_base_rerank\n",
    "    \n",
    "    # annoy_base_constructor, constructs\n",
    "    def annoy_base_constructor(self, items_db=None, distance_type='angular', chunk_size=50, n_trees=10):\n",
    "        items_db_to_use = self.items_db if items_db == None else items_db\n",
    "        all_item_ids = list(items_db_to_use.keys())\n",
    "        \n",
    "        total_pooled = []\n",
    "        # break into chunks/batches for model concurrency\n",
    "        num_chunks = math.ceil(len(all_item_ids) / chunk_size)\n",
    "        for i in range(num_chunks):\n",
    "            chunk_ids = all_item_ids[i*chunk_size: (i+1)*chunk_size]\n",
    "            chunk_pooled, chunk_hiddens = self.compute_encoded_embeddings_for_items(self.recall_encoder, chunk_ids, items_db_to_use)\n",
    "            chunk_pooled = chunk_pooled.cpu().detach().numpy()\n",
    "            del chunk_hiddens\n",
    "            total_pooled.append(chunk_pooled)\n",
    "        total_pooled = np.concatenate(total_pooled, axis=0)\n",
    "        \n",
    "        pooled_tensor = torch.tensor(total_pooled).to(self.device)\n",
    "        \n",
    "        #build recall annoy index\n",
    "        annoy_base_recall = AnnoyIndex(self.recall_item_wte_mapper.out_features, distance_type)\n",
    "        pooled_recall = self.recall_item_wte_mapper(pooled_tensor) # [len(items_db_to_use), self.recall_item_wte_mapper.out_features]\n",
    "        pooled_recall = pooled_recall.cpu().detach().numpy()\n",
    "        for i, vector in zip(all_item_ids, pooled_recall):\n",
    "            annoy_base_recall.add_item(i, vector)\n",
    "        annoy_base_recall.build(n_trees)\n",
    "        \n",
    "        total_pooled = []\n",
    "        # break into chunks/batches for model concurrency\n",
    "        num_chunks = math.ceil(len(all_item_ids) / chunk_size)\n",
    "        for i in range(num_chunks):\n",
    "            chunk_ids = all_item_ids[i*chunk_size: (i+1)*chunk_size]\n",
    "            chunk_pooled, chunk_hiddens = self.compute_encoded_embeddings_for_items(self.item_encoder, chunk_ids, items_db_to_use)\n",
    "            chunk_pooled = chunk_pooled.cpu().detach().numpy()\n",
    "            del chunk_hiddens\n",
    "            total_pooled.append(chunk_pooled)\n",
    "        total_pooled = np.concatenate(total_pooled, axis=0)\n",
    "        \n",
    "        pooled_tensor = torch.tensor(total_pooled).to(self.device)\n",
    "        \n",
    "        #build rerank annoy index, for validation and inference only\n",
    "        annoy_base_rerank = AnnoyIndex(self.rerank_item_wte_mapper.out_features, distance_type)\n",
    "        pooled_rerank = self.rerank_item_wte_mapper(pooled_tensor) # [len(items_db_to_use), self.recall_item_wte_mapper.out_features]\n",
    "        pooled_rerank = pooled_rerank.cpu().detach().numpy()\n",
    "        for i, vector in zip(all_item_ids, pooled_rerank):\n",
    "            annoy_base_rerank.add_item(i, vector)\n",
    "        annoy_base_rerank.build(n_trees)\n",
    "        \n",
    "        del pooled_tensor\n",
    "        \n",
    "        self.annoy_base_recall = annoy_base_recall\n",
    "        self.annoy_base_rerank = annoy_base_rerank\n",
    "    \n",
    "    def annoy_loader(self, path, annoy_type, distance_type=\"angular\"):\n",
    "        if annoy_type == \"recall\":\n",
    "            annoy_base = AnnoyIndex(self.recall_item_wte_mapper.out_features, distance_type)\n",
    "            annoy_base.load(path)\n",
    "            return annoy_base\n",
    "        elif annoy_type == \"rerank\":\n",
    "            annoy_base = AnnoyIndex(self.rerank_item_wte_mapper.out_features, distance_type)\n",
    "            annoy_base.load(path)\n",
    "            return annoy_base\n",
    "        else:\n",
    "            return None\n",
    "    \n",
    "    def lm_expand_wtes_with_items_annoy_base(self):\n",
    "        all_item_ids = list(self.items_db.keys())\n",
    "        total_pooled = []\n",
    "        for i in all_item_ids:\n",
    "            total_pooled.append(self.annoy_base_rerank.get_item_vector(i))\n",
    "        total_pooled = np.asarray(total_pooled) # [len(all_item_ids), 768]\n",
    "        pooled_tensor = torch.tensor(total_pooled, dtype=torch.float).to(self.device)\n",
    "        \n",
    "        old_embeddings = self.language_model.get_input_embeddings()\n",
    "        item_id_2_lm_token_id = {}\n",
    "        for k in all_item_ids:\n",
    "            item_id_2_lm_token_id[k] = len(item_id_2_lm_token_id) + old_embeddings.weight.shape[0]\n",
    "        new_embeddings = torch.cat([old_embeddings.weight, pooled_tensor], 0)\n",
    "        new_embeddings = torch.nn.Embedding.from_pretrained(new_embeddings)\n",
    "        self.language_model.set_input_embeddings(new_embeddings)\n",
    "        self.language_model.to(device)\n",
    "        return item_id_2_lm_token_id\n",
    "    \n",
    "    def lm_restore_wtes(self, original_token_emb_size):\n",
    "        old_embeddings = self.language_model.get_input_embeddings()\n",
    "        new_embeddings = torch.nn.Embedding(original_token_emb_size, old_embeddings.weight.size()[1])\n",
    "        new_embeddings.to(self.device, dtype=old_embeddings.weight.dtype)\n",
    "        new_embeddings.weight.data[:original_token_emb_size, :] = old_embeddings.weight.data[:original_token_emb_size, :]\n",
    "        self.language_model.set_input_embeddings(new_embeddings)\n",
    "        self.language_model.to(self.device)\n",
    "        assert(self.language_model.get_input_embeddings().weight.shape[0] == original_token_emb_size)\n",
    "    \n",
    "    def trim_lm_wtes(self, wtes):\n",
    "        trimmed_wtes = wtes\n",
    "        if trimmed_wtes.shape[1] > self.language_model.config.n_positions:\n",
    "            trimmed_wtes = trimmed_wtes[:,-self.language_model.config.n_positions + self.lm_trim_offset:,:]\n",
    "        return trimmed_wtes # [batch, self.language_model.config.n_positions - self.lm_trim_offset, self.language_model.config.n_embd]\n",
    "    \n",
    "    def trim_positional_ids(self, p_ids, num_items_wtes):\n",
    "        trimmed_ids = p_ids\n",
    "        if trimmed_ids.shape[1] > self.language_model.config.n_positions:\n",
    "            past_ids = trimmed_ids[:,:self.language_model.config.n_positions - self.lm_trim_offset - num_items_wtes]\n",
    "#             past_ids = trimmed_ids[:, self.lm_trim_offset + num_items_wtes:self.language_model.config.n_positions]\n",
    "            item_ids = trimmed_ids[:,-num_items_wtes:]\n",
    "            trimmed_ids = torch.cat((past_ids, item_ids), dim=1)\n",
    "        return trimmed_ids # [batch, self.language_model.config.n_positions - self.lm_trim_offset]\n",
    "    \n",
    "    def compute_inductive_attention_mask(self, length_language, length_rerank_items_wtes):\n",
    "        total_length = length_language + length_rerank_items_wtes\n",
    "        language_mask_to_add = torch.zeros((length_language, total_length), dtype=torch.float, device=self.device)\n",
    "        items_mask_to_add = torch.ones((length_rerank_items_wtes, total_length), dtype=torch.float, device=self.device)\n",
    "        combined_mask_to_add = torch.cat((language_mask_to_add, items_mask_to_add), dim=0)\n",
    "        return combined_mask_to_add #[total_length, total_length]\n",
    "    \n",
    "    def forward_pure_language_turn(self, \n",
    "                                   past_wtes, # past word token embeddings, [1, len, 768]\n",
    "                                   current_tokens # tokens of current turn conversation, [1, len]\n",
    "                                  ):\n",
    "        train_logits, train_targets = None, None\n",
    "        current_wtes = self.language_model.transformer.wte(current_tokens)\n",
    "        \n",
    "        if past_wtes == None:\n",
    "            lm_outputs = self.language_model(inputs_embeds=current_wtes)\n",
    "            train_logits = lm_outputs.logits[:, :-1, :]\n",
    "            train_targets = current_tokens[:,1:]\n",
    "        else:\n",
    "            all_wtes = torch.cat((past_wtes, current_wtes), dim=1)\n",
    "            all_wtes = self.trim_lm_wtes(all_wtes)\n",
    "            lm_outputs = self.language_model(inputs_embeds=all_wtes)\n",
    "            train_logits = lm_outputs.logits[:, -current_wtes.shape[1]:-1, :] # skip the last one\n",
    "            train_targets = current_tokens[:,1:]\n",
    "        \n",
    "        # torch.Size([batch, len_cur, lm_vocab]), torch.Size([batch, len_cur]), torch.Size([batch, len_past+len_cur, lm_emb(768)])\n",
    "        # if past_wtes == None, len_cur - 1\n",
    "        return train_logits, train_targets\n",
    "        \n",
    "    def forward_recall(self, \n",
    "                       past_wtes, # past word token embeddings, [1, len, 768]\n",
    "                       current_tokens, # tokens of current turn conversation, [1, len]\n",
    "                       gt_item_id, # id, ex. 0\n",
    "                       num_samples # num examples to sample for training, including groud truth id\n",
    "                      ):\n",
    "        # recall step 1. construct LM sequence output\n",
    "        # LM input composition: [past_wtes, REC_wtes, gt_item_wte, (gt_item_info_wtes), REC_END_wtes, current_wtes ]\n",
    "        \n",
    "        REC_wtes = self.get_rec_token_wtes() # [1, 1, self.language_model.config.n_embd]\n",
    "        gt_item_wte, _ = self.compute_encoded_embeddings_for_items(self.recall_encoder, [gt_item_id], self.items_db)\n",
    "        gt_item_wte = self.rerank_item_wte_mapper(gt_item_wte) # [1, self.rerank_item_wte_mapper.out_features]\n",
    "        \n",
    "        REC_END_wtes = self.get_rec_end_token_wtes() # [1, 1, self.language_model.config.n_embd]\n",
    "        current_wtes = self.language_model.transformer.wte(current_tokens) #[1, current_tokens.shape[1], self.language_model.config.n_embd]\n",
    "        \n",
    "        REC_wtes_len = REC_wtes.shape[1] # 1 by default\n",
    "        gt_item_wte_len = gt_item_wte.shape[0] # 1 by default\n",
    "        REC_END_wtes_len = REC_END_wtes.shape[1] # 1 by default\n",
    "        current_wtes_len = current_wtes.shape[1]\n",
    "        \n",
    "        lm_wte_inputs = torch.cat(\n",
    "            (past_wtes, # [batch (1), len, self.language_model.config.n_embd]\n",
    "             REC_wtes,\n",
    "             gt_item_wte.unsqueeze(0), # reshape to [1,1,self.rerank_item_wte_mapper.out_features]\n",
    "             REC_END_wtes,\n",
    "             current_wtes # [batch (1), len, self.language_model.config.n_embd]\n",
    "            ),\n",
    "            dim=1\n",
    "        )\n",
    "        lm_wte_inputs = self.trim_lm_wtes(lm_wte_inputs) # trim for len > self.language_model.config.n_positions\n",
    "        \n",
    "        # recall step 2. get gpt output logits and hidden states\n",
    "        lm_outputs = self.language_model(inputs_embeds=lm_wte_inputs, output_hidden_states=True)\n",
    "        \n",
    "        # recall step 3. pull logits (recall, rec_token and language logits of current turn) and compute\n",
    "        \n",
    "        # recall logit(s)\n",
    "        rec_token_start_index = -current_wtes_len-REC_END_wtes_len-gt_item_wte_len-REC_wtes_len\n",
    "        rec_token_end_index = -current_wtes_len-REC_END_wtes_len-gt_item_wte_len\n",
    "        # [batch (1), REC_wtes_len, self.language_model.config.n_embd]\n",
    "        rec_token_hidden = lm_outputs.hidden_states[-1][:, rec_token_start_index:rec_token_end_index, :]\n",
    "        # [batch (1), self.recall_lm_query_mapper.out_features]\n",
    "        rec_query_vector = self.recall_lm_query_mapper(rec_token_hidden).squeeze(1)\n",
    "        \n",
    "        # sample num_samples item ids to train recall with \"recommendation as classification\"\n",
    "        sampled_item_ids = sample_ids_from_db(self.items_db, gt_item_id, num_samples, include_gt=True)\n",
    "        gt_item_id_index = sampled_item_ids.index(gt_item_id)\n",
    "        \n",
    "        # [num_samples, self.item_encoder.config.hidden_size]\n",
    "        encoded_items_embeddings, _ = self.compute_encoded_embeddings_for_items(self.recall_encoder, sampled_item_ids, self.items_db)\n",
    "        # to compute dot product with rec_query_vector\n",
    "        items_key_vectors = self.recall_item_wte_mapper(encoded_items_embeddings) # [num_samples, self.recall_item_wte_mapper.out_features]\n",
    "        expanded_rec_query_vector = rec_query_vector.expand(items_key_vectors.shape[0], rec_query_vector.shape[1]) # [num_samples, self.recall_item_wte_mapper.out_features]\n",
    "        recall_logits = torch.sum(expanded_rec_query_vector * items_key_vectors, dim=1) # torch.size([num_samples])\n",
    "        \n",
    "        # REC_TOKEN prediction and future sentence prediction\n",
    "        # hidden rep of the token that's right before REC_TOKEN\n",
    "        token_before_REC_logits = lm_outputs.logits[:, rec_token_start_index-1:rec_token_end_index-1, :]\n",
    "        REC_targets = self.lm_tokenizer(self.REC_TOKEN_STR, return_tensors=\"pt\")['input_ids'].to(self.device) # [1, 1]\n",
    "        \n",
    "        #language logits and targets\n",
    "        current_language_logits = lm_outputs.logits[:, -current_wtes_len:-1, :]\n",
    "        current_language_targets = current_tokens[:,1:]\n",
    "        \n",
    "        # REC token and language, their logits and targets\n",
    "        # [batch, current_wtes_len+REC_wtes_len, lm_vocab]\n",
    "        all_wte_logits = torch.cat((token_before_REC_logits, current_language_logits), dim=1)\n",
    "        # [current_wtes_len+REC_wtes_len, lm_vocab]\n",
    "        all_wte_targets = torch.cat((REC_targets, current_language_targets), dim=1)\n",
    "        \n",
    "        # torch.size([num_samples]), id, [batch, current_wtes_len+REC_wtes_len, lm_vocab], [current_wtes_len+REC_wtes_len, lm_vocab]\n",
    "        return recall_logits, gt_item_id_index, all_wte_logits, all_wte_targets\n",
    "        \n",
    "    def forward_rerank(self,\n",
    "                       past_wtes, # past word token embeddings, [1, len, 768]\n",
    "                       gt_item_id, # tokens of current turn conversation, [1, len]\n",
    "                       num_samples, # num examples to sample for training, including groud truth id\n",
    "                       rerank_items_chunk_size=10, # batch size for encoder GPU computation\n",
    "                      ):\n",
    "        # REC wte\n",
    "        REC_wtes = self.get_rec_token_wtes() # [batch (1), 1, self.language_model.config.n_embd]\n",
    "        \n",
    "        #  items wtes to compute rerank loss\n",
    "        # sample rerank examples\n",
    "        sampled_item_ids = sample_ids_from_db(self.items_db, gt_item_id, num_samples, include_gt=True)\n",
    "        gt_item_id_index = sampled_item_ids.index(gt_item_id)\n",
    "        # compute item wtes by batch\n",
    "        num_chunks = math.ceil(len(sampled_item_ids) / rerank_items_chunk_size)\n",
    "        total_wtes = []\n",
    "        for i in range(num_chunks):\n",
    "            chunk_ids = sampled_item_ids[i*rerank_items_chunk_size: (i+1)*rerank_items_chunk_size]\n",
    "            chunk_pooled, _ = self.compute_encoded_embeddings_for_items(self.item_encoder, chunk_ids, self.items_db) # [rerank_items_chunk_size, self.item_encoder.config.hidden_size]\n",
    "            chunk_wtes = self.rerank_item_wte_mapper(chunk_pooled)\n",
    "            total_wtes.append(chunk_wtes)\n",
    "        total_wtes = torch.cat(total_wtes, dim=0) # [num_samples, self.language_model.config.n_embd]\n",
    "        \n",
    "        past_wtes_len = past_wtes.shape[1]\n",
    "        REC_wtes_len = REC_wtes.shape[1] # 1 by default\n",
    "        total_wtes_len = total_wtes.shape[0]\n",
    "        \n",
    "        # compute positional ids, all rerank item wte should have the same positional encoding id 0\n",
    "        position_ids = torch.arange(0, past_wtes_len + REC_wtes_len, dtype=torch.long, device=self.device)\n",
    "        items_position_ids = torch.zeros(total_wtes.shape[0], dtype=torch.long, device=device)\n",
    "#         items_position_ids = torch.tensor([1023] * total_wtes.shape[0], dtype=torch.long, device=device)\n",
    "        combined_position_ids = torch.cat((position_ids, items_position_ids), dim=0)\n",
    "        combined_position_ids = combined_position_ids.unsqueeze(0) # [1, past_wtes_len+REC_wtes_len+total_wtes_len]\n",
    "        \n",
    "        # compute concatenated lm wtes\n",
    "        lm_wte_inputs = torch.cat(\n",
    "            (past_wtes, # [batch (1), len, self.language_model.config.n_embd]\n",
    "             REC_wtes, # [batch (1), 1, self.language_model.config.n_embd]\n",
    "             total_wtes.unsqueeze(0), # [1, num_samples, self.language_model.config.n_embd]\n",
    "            ),\n",
    "            dim=1\n",
    "        ) # [1, past_len + REC_wtes_len + num_samples, self.language_model.config.n_embd]\n",
    "\n",
    "        # trim sequence to smaller length (len < self.language_model.config.n_positions-self.lm_trim_offset)\n",
    "        combined_position_ids_trimmed = self.trim_positional_ids(combined_position_ids, total_wtes_len) # [1, len]\n",
    "        lm_wte_inputs_trimmed = self.trim_lm_wtes(lm_wte_inputs) # [1, len, self.language_model.config.n_embd]\n",
    "        assert(combined_position_ids.shape[1] == lm_wte_inputs.shape[1])\n",
    "        \n",
    "        # compute inductive attention mask\n",
    "        #     Order of recommended items shouldn't affect their score, thus every item \n",
    "        # should have full attention over the entire sequence: they should know each other and the entire\n",
    "        # conversation history\n",
    "        inductive_attention_mask = self.compute_inductive_attention_mask(\n",
    "            lm_wte_inputs_trimmed.shape[1]-total_wtes.shape[0], \n",
    "            total_wtes.shape[0]\n",
    "        )\n",
    "        rerank_lm_outputs = self.language_model(inputs_embeds=lm_wte_inputs_trimmed,\n",
    "                  inductive_attention_mask=inductive_attention_mask,\n",
    "                  position_ids=combined_position_ids_trimmed,\n",
    "                  output_hidden_states=True)\n",
    "        \n",
    "        rerank_lm_hidden = rerank_lm_outputs.hidden_states[-1][:, -total_wtes.shape[0]:, :]\n",
    "        rerank_logits = self.rerank_logits_mapper(rerank_lm_hidden).squeeze() # torch.Size([num_samples])\n",
    "        \n",
    "        return rerank_logits, gt_item_id_index\n",
    "    \n",
    "    def validation_perform_recall(self, past_wtes, topk):\n",
    "        REC_wtes = self.get_rec_token_wtes()\n",
    "        lm_wte_inputs = torch.cat(\n",
    "            (past_wtes, # [batch (1), len, self.language_model.config.n_embd]\n",
    "             REC_wtes # [1, 1, self.language_model.config.n_embd]\n",
    "            ),\n",
    "            dim=1\n",
    "        )\n",
    "        lm_wte_inputs = self.trim_lm_wtes(lm_wte_inputs) # trim for len > self.language_model.config.n_positions\n",
    "        lm_outputs = self.language_model(inputs_embeds=lm_wte_inputs, output_hidden_states=True)\n",
    "        \n",
    "        rec_token_hidden = lm_outputs.hidden_states[-1][:, -1, :]\n",
    "        # [batch (1), self.recall_lm_query_mapper.out_features]\n",
    "        rec_query_vector = self.recall_lm_query_mapper(rec_token_hidden).squeeze(0) # [768]\n",
    "        rec_query_vector = rec_query_vector.cpu().detach().numpy()\n",
    "        recall_results = self.annoy_base_recall.get_nns_by_vector(rec_query_vector, topk)\n",
    "        return recall_results\n",
    "    \n",
    "    def validation_perform_rerank(self, past_wtes, recalled_ids):\n",
    "        REC_wtes = self.get_rec_token_wtes()\n",
    "        \n",
    "        total_wtes = [ self.annoy_base_rerank.get_item_vector(r_id) for r_id in recalled_ids]\n",
    "        total_wtes = [ torch.tensor(wte).reshape(-1, self.language_model.config.n_embd).to(self.device) for wte in total_wtes]\n",
    "        total_wtes = torch.cat(total_wtes, dim=0) # [len(recalled_ids), 768]\n",
    "        \n",
    "        past_wtes_len = past_wtes.shape[1]\n",
    "        REC_wtes_len = REC_wtes.shape[1] # 1 by default\n",
    "        total_wtes_len = total_wtes.shape[0]\n",
    "        \n",
    "        # compute positional ids, all rerank item wte should have the same positional encoding id 0\n",
    "        position_ids = torch.arange(0, past_wtes_len + REC_wtes_len, dtype=torch.long, device=self.device)\n",
    "        items_position_ids = torch.zeros(total_wtes.shape[0], dtype=torch.long, device=device)\n",
    "#         items_position_ids = torch.tensor([1023] * total_wtes.shape[0], dtype=torch.long, device=device)\n",
    "        combined_position_ids = torch.cat((position_ids, items_position_ids), dim=0)\n",
    "        combined_position_ids = combined_position_ids.unsqueeze(0) # [1, past_wtes_len+REC_wtes_len+total_wtes_len]\n",
    "        \n",
    "        # compute concatenated lm wtes\n",
    "        lm_wte_inputs = torch.cat(\n",
    "            (past_wtes, # [batch (1), len, self.language_model.config.n_embd]\n",
    "             REC_wtes, # [batch (1), 1, self.language_model.config.n_embd]\n",
    "             total_wtes.unsqueeze(0), # [1, num_samples, self.language_model.config.n_embd]\n",
    "            ),\n",
    "            dim=1\n",
    "        ) # [1, past_len + REC_wtes_len + num_samples, self.language_model.config.n_embd]\n",
    "\n",
    "        # trim sequence to smaller length (len < self.language_model.config.n_positions-self.lm_trim_offset)\n",
    "        combined_position_ids_trimmed = self.trim_positional_ids(combined_position_ids, total_wtes_len) # [1, len]\n",
    "        lm_wte_inputs_trimmed = self.trim_lm_wtes(lm_wte_inputs) # [1, len, self.language_model.config.n_embd]\n",
    "        assert(combined_position_ids.shape[1] == lm_wte_inputs.shape[1])\n",
    "        \n",
    "        inductive_attention_mask = self.compute_inductive_attention_mask(\n",
    "            lm_wte_inputs_trimmed.shape[1]-total_wtes.shape[0], \n",
    "            total_wtes.shape[0]\n",
    "        )\n",
    "        rerank_lm_outputs = self.language_model(inputs_embeds=lm_wte_inputs_trimmed,\n",
    "                  inductive_attention_mask=inductive_attention_mask,\n",
    "                  position_ids=combined_position_ids_trimmed,\n",
    "                  output_hidden_states=True)\n",
    "        \n",
    "        rerank_lm_hidden = rerank_lm_outputs.hidden_states[-1][:, -total_wtes.shape[0]:, :]\n",
    "        rerank_logits = self.rerank_logits_mapper(rerank_lm_hidden).squeeze()\n",
    "        \n",
    "        return rerank_logits\n",
    "    \n",
    "    # TODO: Modify\n",
    "    def validation_generate_sentence(self, past_wtes, recommended_id):\n",
    "        if recommended_id == None: #  pure language\n",
    "            lm_wte_inputs = self.trim_lm_wtes(past_wtes)\n",
    "            lm_outputs = self.language_model(inputs_embeds=lm_wte_inputs, output_hidden_states=True)\n",
    "            generated = model.language_model.generate(\n",
    "                past_key_values=lm_outputs.past_key_values, \n",
    "                max_length=50, \n",
    "                num_return_sequences=1, \n",
    "                do_sample=True, \n",
    "                eos_token_id=198,\n",
    "                pad_token_id=198\n",
    "            )\n",
    "            generated_sen = gpt_tokenizer.decode(generated[0], skip_special_tokens=True)\n",
    "            return generated_sen, None\n",
    "        else:\n",
    "            REC_wtes = self.get_rec_token_wtes() # [1, 1, self.language_model.config.n_embd]\n",
    "            gt_item_wte, _ = self.compute_encoded_embeddings_for_items([recommended_id], self.items_db)\n",
    "            gt_item_wte = self.rerank_item_wte_mapper(gt_item_wte) # [1, self.rerank_item_wte_mapper.out_features]\n",
    "            REC_END_wtes = self.get_rec_end_token_wtes() # [1, 1, self.language_model.config.n_embd]\n",
    "\n",
    "            REC_wtes_len = REC_wtes.shape[1] # 1 by default\n",
    "            gt_item_wte_len = gt_item_wte.shape[0] # 1 by default\n",
    "            REC_END_wtes_len = REC_END_wtes.shape[1] # 1 by default\n",
    "\n",
    "            lm_wte_inputs = torch.cat(\n",
    "                (past_wtes, # [batch (1), len, self.language_model.config.n_embd]\n",
    "                 REC_wtes,\n",
    "                 gt_item_wte.unsqueeze(0), # reshape to [1,1,self.rerank_item_wte_mapper.out_features]\n",
    "                 REC_END_wtes\n",
    "                ),\n",
    "                dim=1\n",
    "            )\n",
    "            lm_wte_inputs = self.trim_lm_wtes(lm_wte_inputs) # trim for len > self.language_model.config.n_positions\n",
    "            lm_outputs = self.language_model(inputs_embeds=lm_wte_inputs, output_hidden_states=True)\n",
    "            generated = model.language_model.generate(\n",
    "                past_key_values=lm_outputs.past_key_values, \n",
    "                max_length=50, \n",
    "                num_return_sequences=1, \n",
    "                do_sample=True, \n",
    "                eos_token_id=198,\n",
    "                pad_token_id=198\n",
    "            )\n",
    "            generated_sen = gpt_tokenizer.decode(generated[0], skip_special_tokens=True)\n",
    "            return generated_sen, self.get_movie_title(recommended_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6a424c72",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(0)\n",
    "model = UniversalCRSModel(gpt2_model, bert_model_rerank, bert_model_recall, gpt_tokenizer, bert_tokenizer, device, items_db, rec_token_str=REC_TOKEN, rec_end_token_str=REC_END_TOKEN)\n",
    "# model.load_state_dict(torch.load(\"/local-scratch1/data/by2299/CRS_Redial_Train_1st_normal_lang_sep_annoy_7.pt\"))\n",
    "\n",
    "model.to(device)\n",
    "pass\n",
    "# model.recall_item_wte_mapper.out_features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d8369177",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.annoy_base_constructor()\n",
    "model.annoy_base_constructor_emb_analysis()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "08340717",
   "metadata": {},
   "outputs": [],
   "source": [
    "movie_2_genres = {}\n",
    "for k, v in items_db.items():\n",
    "    genres = v.split(\"[SEP]\")[3].strip().split(\",\")\n",
    "#     if len(genres) != 1 or genres[0] == '': continue\n",
    "    if '' in genres:\n",
    "        continue\n",
    "    for i in range(len(genres)):\n",
    "        genres[i] = genres[i].strip()\n",
    "        if 'Sci-Fi' in genres[i]:\n",
    "            genres[i] = 'Science Fiction'\n",
    "        if 'Musical' in genres[i]:\n",
    "            genres[i] = 'Music'\n",
    "    movie_2_genres[k] = genres"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "292a2e34",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Animation': 0,\n",
       " 'Action': 1,\n",
       " 'Adventure': 2,\n",
       " 'Comedy': 3,\n",
       " 'Drama': 4,\n",
       " 'Crime': 5,\n",
       " 'Fantasy': 6,\n",
       " 'Horror': 7,\n",
       " 'Science Fiction': 8,\n",
       " 'Thriller': 9,\n",
       " 'Mystery': 10,\n",
       " 'Romance': 11,\n",
       " 'Music': 12,\n",
       " 'Family': 13,\n",
       " 'Biography': 14,\n",
       " 'Sport': 15,\n",
       " 'History': 16,\n",
       " 'War': 17,\n",
       " 'Film-Noir': 18,\n",
       " 'Documentary': 19,\n",
       " 'TV Movie': 20,\n",
       " 'Western': 21,\n",
       " 'Foreign': 22,\n",
       " 'Short': 23,\n",
       " 'Adult': 24}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "distinct_genres = {}\n",
    "for k, genres in movie_2_genres.items():\n",
    "    for g in genres:\n",
    "        if g.strip() not in distinct_genres:\n",
    "            distinct_genres[g.strip()] = len(distinct_genres)\n",
    "distinct_genres"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9b3856bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "wtes = []\n",
    "genre_labels = []\n",
    "for k, genres in movie_2_genres.items():\n",
    "    wte = model.annoy_base_rerank.get_item_vector(k)\n",
    "    wtes.append(wte)\n",
    "    genre_vector = [0] * len(distinct_genres)\n",
    "    for g in genres:\n",
    "        genre_vector[distinct_genres[g]] = 1\n",
    "    genre_labels.append(genre_vector)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b87004ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "portion = 4800\n",
    "train_wtes = wtes[:portion]; train_labels = genre_labels[:portion]\n",
    "test_wtes = wtes[portion:]; test_labels = genre_labels[portion:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a006f170",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiClassModel(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MultiClassModel, self).__init__()\n",
    "        self.core = torch.nn.Linear(768, len(distinct_genres))\n",
    "    def forward(self, vectors):\n",
    "        return self.core(vectors)\n",
    "model = MultiClassModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6a8d133d",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = torch.nn.BCEWithLogitsLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.03)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "e2d73044",
   "metadata": {},
   "outputs": [],
   "source": [
    "#eval\n",
    "def eval_func(features_test, labels_test):\n",
    "    total = 0; correct = 0\n",
    "    features_test = torch.tensor( features_test, dtype=torch.float )\n",
    "    labels_test = torch.tensor( labels_test, dtype=torch.float )\n",
    "    preds = model.forward(features_test)\n",
    "    sig = torch.nn.Sigmoid()\n",
    "    preds = sig(preds)\n",
    "    for pred, gt in zip(preds, labels_test):\n",
    "        c_vec = [0] * len(distinct_genres)\n",
    "        for i in range(len(gt)):\n",
    "#             if pred[i] >= 0.5:\n",
    "#                 c_vec[i] = 1\n",
    "#             else:\n",
    "#                 c_vec[i] = 0\n",
    "#         print(gt)\n",
    "#         print(c_vec)\n",
    "#         print()\n",
    "            if gt[i] == 1:\n",
    "                total += 1\n",
    "                if pred[i] >= 0.5:\n",
    "                    correct += 1\n",
    "#             elif gt[i] == 0 and pred[i] < 0.5:\n",
    "#                 correct += 1\n",
    "    return correct, total\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7fab34df",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, ratio 0.5061728395061729\n",
      "Epoch 1, ratio 0.5653021442495126\n",
      "Epoch 2, ratio 0.5912930474333983\n",
      "Epoch 3, ratio 0.6140350877192983\n",
      "Epoch 4, ratio 0.6231319038336582\n",
      "Epoch 5, ratio 0.6367771280051981\n",
      "Epoch 6, ratio 0.6432748538011696\n",
      "Epoch 7, ratio 0.6510721247563352\n",
      "Epoch 8, ratio 0.6588693957115009\n",
      "Epoch 9, ratio 0.6653671215074723\n",
      "Epoch 10, ratio 0.6686159844054581\n",
      "Epoch 11, ratio 0.6757634827810266\n",
      "Epoch 12, ratio 0.6848602988953866\n",
      "Epoch 13, ratio 0.6900584795321637\n",
      "Epoch 14, ratio 0.6939571150097466\n",
      "Epoch 15, ratio 0.6952566601689408\n",
      "Epoch 16, ratio 0.6985055230669266\n",
      "Epoch 17, ratio 0.7030539311241065\n",
      "Epoch 18, ratio 0.7069525666016894\n",
      "Epoch 19, ratio 0.7108512020792722\n",
      "Epoch 20, ratio 0.7166991552956465\n",
      "Epoch 21, ratio 0.7179987004548408\n",
      "Epoch 22, ratio 0.7179987004548408\n",
      "Epoch 23, ratio 0.7205977907732294\n",
      "Epoch 24, ratio 0.723196881091618\n",
      "Epoch 25, ratio 0.7257959714100065\n",
      "Epoch 26, ratio 0.7264457439896036\n",
      "Epoch 27, ratio 0.7277452891487979\n",
      "Epoch 28, ratio 0.7283950617283951\n",
      "Epoch 29, ratio 0.7309941520467836\n",
      "Epoch 30, ratio 0.7348927875243665\n",
      "Epoch 31, ratio 0.7355425601039636\n",
      "Epoch 32, ratio 0.7355425601039636\n",
      "Epoch 33, ratio 0.7394411955815464\n",
      "Epoch 34, ratio 0.7407407407407407\n",
      "Epoch 35, ratio 0.7426900584795322\n",
      "Epoch 36, ratio 0.742040285899935\n",
      "Epoch 37, ratio 0.7413905133203379\n",
      "Epoch 38, ratio 0.7400909681611436\n",
      "Epoch 39, ratio 0.7407407407407407\n",
      "Epoch 40, ratio 0.7407407407407407\n",
      "Epoch 41, ratio 0.7394411955815464\n",
      "Epoch 42, ratio 0.7355425601039636\n",
      "Epoch 43, ratio 0.7348927875243665\n",
      "Epoch 44, ratio 0.7348927875243665\n",
      "Epoch 45, ratio 0.7348927875243665\n",
      "Epoch 46, ratio 0.7355425601039636\n",
      "Epoch 47, ratio 0.737491877842755\n",
      "Epoch 48, ratio 0.7381416504223521\n",
      "Epoch 49, ratio 0.7394411955815464\n"
     ]
    }
   ],
   "source": [
    "#train\n",
    "batch_size = 100\n",
    "epochs = 50\n",
    "for e in range(epochs):\n",
    "    model.train()\n",
    "    for i in range(int(portion/batch_size)):\n",
    "        features_train = torch.tensor( train_wtes[i*batch_size : (i+1)*batch_size], dtype=torch.float )\n",
    "        labels_train = torch.tensor( train_labels[i*batch_size : (i+1)*batch_size], dtype=torch.float )\n",
    "        outputs = model.forward(features_train)\n",
    "        loss = criterion(outputs, labels_train)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    model.eval()\n",
    "    correct, total = eval_func(test_wtes, test_labels)\n",
    "    print(f\"Epoch {e}, ratio {correct/total}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1860211b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#base 0.743\n",
    "# ep0 0.871\n",
    "# ep1 0.860\n",
    "# ep2 0.793\n",
    "# ep3 0.782\n",
    "# ep4 0.719"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "2acd347d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, ratio 0.2014294996751137\n",
      "Epoch 1, ratio 0.2378167641325536\n",
      "Epoch 2, ratio 0.28914879792072773\n",
      "Epoch 3, ratio 0.3398310591293047\n",
      "Epoch 4, ratio 0.3742690058479532\n",
      "Epoch 5, ratio 0.39246263807667314\n",
      "Epoch 6, ratio 0.41325536062378165\n",
      "Epoch 7, ratio 0.4301494476933073\n",
      "Epoch 8, ratio 0.44834307992202727\n",
      "Epoch 9, ratio 0.45938921377517866\n",
      "Epoch 10, ratio 0.47043534762833006\n",
      "Epoch 11, ratio 0.4821312540610786\n",
      "Epoch 12, ratio 0.48927875243664715\n",
      "Epoch 13, ratio 0.49577647823261856\n",
      "Epoch 14, ratio 0.5029239766081871\n",
      "Epoch 15, ratio 0.50682261208577\n",
      "Epoch 16, ratio 0.5139701104613386\n",
      "Epoch 17, ratio 0.5224171539961013\n",
      "Epoch 18, ratio 0.5282651072124757\n",
      "Epoch 19, ratio 0.5308641975308642\n",
      "Epoch 20, ratio 0.5386614684860299\n",
      "Epoch 21, ratio 0.5432098765432098\n",
      "Epoch 22, ratio 0.5477582846003899\n",
      "Epoch 23, ratio 0.5516569200779727\n",
      "Epoch 24, ratio 0.557504873294347\n",
      "Epoch 25, ratio 0.5594541910331384\n",
      "Epoch 26, ratio 0.5607537361923327\n",
      "Epoch 27, ratio 0.5627030539311241\n",
      "Epoch 28, ratio 0.5659519168291098\n",
      "Epoch 29, ratio 0.5672514619883041\n",
      "Epoch 30, ratio 0.5672514619883041\n",
      "Epoch 31, ratio 0.5679012345679012\n",
      "Epoch 32, ratio 0.5705003248862898\n",
      "Epoch 33, ratio 0.5724496426250812\n",
      "Epoch 34, ratio 0.5730994152046783\n",
      "Epoch 35, ratio 0.5776478232618584\n",
      "Epoch 36, ratio 0.5789473684210527\n",
      "Epoch 37, ratio 0.5782975958414555\n",
      "Epoch 38, ratio 0.5802469135802469\n",
      "Epoch 39, ratio 0.5802469135802469\n",
      "Epoch 40, ratio 0.5841455490578298\n",
      "Epoch 41, ratio 0.5867446393762183\n",
      "Epoch 42, ratio 0.5880441845354126\n",
      "Epoch 43, ratio 0.5906432748538012\n",
      "Epoch 44, ratio 0.5912930474333983\n",
      "Epoch 45, ratio 0.5925925925925926\n",
      "Epoch 46, ratio 0.5932423651721898\n",
      "Epoch 47, ratio 0.5958414554905783\n",
      "Epoch 48, ratio 0.5977907732293697\n",
      "Epoch 49, ratio 0.599090318388564\n"
     ]
    }
   ],
   "source": [
    "#train\n",
    "batch_size = 100\n",
    "epochs = 50\n",
    "for e in range(epochs):\n",
    "    model.train()\n",
    "    for i in range(int(4800/batch_size)):\n",
    "        features_train = torch.tensor( train_wtes[i*100 : (i+1)*100], dtype=torch.float )\n",
    "        labels_train = torch.tensor( train_labels[i*100 : (i+1)*100], dtype=torch.float )\n",
    "        outputs = model.forward(features_train)\n",
    "        loss = criterion(outputs, labels_train)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    model.eval()\n",
    "    correct, total = eval_func(test_wtes, test_labels)\n",
    "    print(f\"Epoch {e}, ratio {correct/total}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8463cac",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6a21a1c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "single_genre = {}\n",
    "for k, v in items_db.items():\n",
    "    if len(v.split(\"[SEP]\")) < 5: continue\n",
    "    genres = v.split(\"[SEP]\")[3].strip().split(\",\")\n",
    "    if genres[0] == \"Documentary\":\n",
    "        continue\n",
    "    if len(genres) == 1 and genres[0] != '':\n",
    "        cleaned_genre = genres[0]\n",
    "        if genres[0] == 'Music':\n",
    "            cleaned_genre = 'Musical'\n",
    "        if genres[0] == 'Sci-Fi':\n",
    "            cleaned_genre = 'Science Fiction'\n",
    "        single_genre[k] = cleaned_genre"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5abcee10",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Counter({'Drama': 192,\n",
       "         'Comedy': 218,\n",
       "         'Romance': 5,\n",
       "         'Horror': 69,\n",
       "         'Fantasy': 3,\n",
       "         'Thriller': 17,\n",
       "         'Animation': 5,\n",
       "         'Science Fiction': 6,\n",
       "         'Mystery': 4,\n",
       "         'Action': 7,\n",
       "         'Western': 3,\n",
       "         'Adventure': 3,\n",
       "         'Adult': 1,\n",
       "         'Short': 3,\n",
       "         'Musical': 3,\n",
       "         'Family': 5,\n",
       "         'Crime': 1})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from collections import Counter\n",
    "Counter(single_genre.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "24d2a115",
   "metadata": {},
   "outputs": [],
   "source": [
    "wtes = []\n",
    "genre_labels = []\n",
    "for k,v in single_genre.items():\n",
    "    wte = model.annoy_base_rerank.get_item_vector(k)\n",
    "    wtes.append(wte)\n",
    "    genre_labels.append(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "581b8703",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.cluster import KMeans\n",
    "from collections import defaultdict\n",
    "sample_times = 20\n",
    "avg_acc = []\n",
    "for rs in range(sample_times):\n",
    "    kmeans = KMeans(n_clusters=3, random_state=rs).fit(wtes)\n",
    "    id_2_main_genre = {}\n",
    "    for i in range(len(wtes)):\n",
    "        cluster_id = kmeans.labels_[i]\n",
    "        genre_label = genre_labels[i]\n",
    "        if cluster_id not in id_2_main_genre:\n",
    "            id_2_main_genre[cluster_id] = defaultdict(int)\n",
    "        id_2_main_genre[cluster_id][genre_label] += 1\n",
    "    class_name = []\n",
    "    acc = []\n",
    "    for key in id_2_main_genre.keys():\n",
    "        cluster = id_2_main_genre[key]\n",
    "        cnter = Counter(cluster)\n",
    "        majority = cnter.most_common(1)[0]\n",
    "        total_cnt = sum(cluster.values())\n",
    "        main_acc = majority[1] / total_cnt\n",
    "        class_name.append(majority[0])\n",
    "        acc.append(main_acc)\n",
    "    avg_acc.append(np.mean(acc))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "59b5f16e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.49274422012473573"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.mean(avg_acc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6421a1f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# no training\n",
    "3: 0.49274422012473573\n",
    "4: 0.5136649549291342\n",
    "5: 0.5738269301657036\n",
    "    \n",
    "#no genre, no plot\n",
    "# 3: 0.6044874826639764\n",
    "# 4: 0.6664281020406956\n",
    "# 5: 0.689556041130807\n",
    "3: 0.5554836134916321\n",
    "4: 0.589166330158624\n",
    "5: 0.6060811410293794\n",
    "\n",
    "# sep annoy\n",
    "3: 0.694558391935933\n",
    "4: 0.7253916217797459\n",
    "5: 0.7375563891288288"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7481b4fe",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "d56ba71e",
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (3318023016.py, line 1)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  File \u001b[0;32m\"/tmp/ipykernel_559100/3318023016.py\"\u001b[0;36m, line \u001b[0;32m1\u001b[0m\n\u001b[0;31m    old:\u001b[0m\n\u001b[0m        ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "# without Documentary\n",
    "old:\n",
    "    3, 0.4936322713287639\n",
    "    4, 0.5144524780708778\n",
    "    5, 0.5736495385033692\n",
    "\n",
    "old CLS:\n",
    "    3, 0.43372707112345416\n",
    "    4, 0.45186175234274845\n",
    "    5, 0.490816081925875\n",
    "    \n",
    "trained:\n",
    "    3, 0.6068449328399261\n",
    "    4, 0.7047642984671186\n",
    "    5, 0.7459912269026595\n",
    "    \n",
    "# with documentary\n",
    "old:\n",
    "    3, 0.5419993871238592\n",
    "    4, 0.5654777377625189\n",
    "    5, 0.5780730722362244\n",
    "\n",
    "old CLS:\n",
    "    3, 0.5253262398313966\n",
    "    4, 0.5285883323843146\n",
    "    5, 0.512210301337264\n",
    "\n",
    "trained:\n",
    "    3, 0.5557598070972909\n",
    "    4, 0.6259810961551399\n",
    "    5, 0.7178085603955927"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "aa996433",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "id_2_main_genre = {}\n",
    "for i in range(len(wtes)):\n",
    "    cluster_id = kmeans.labels_[i]\n",
    "    genre_label = genre_labels[i]\n",
    "    if cluster_id not in id_2_main_genre:\n",
    "        id_2_main_genre[cluster_id] = defaultdict(int)\n",
    "    id_2_main_genre[cluster_id][genre_label] += 1\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "577ecdca",
   "metadata": {},
   "outputs": [],
   "source": [
    "class_name = []\n",
    "acc = []\n",
    "for rs in range(20):\n",
    "for key in id_2_main_genre.keys():\n",
    "    cluster = id_2_main_genre[key]\n",
    "    cnter = Counter(cluster)\n",
    "    majority = cnter.most_common(1)[0]\n",
    "    total_cnt = sum(cluster.values())\n",
    "    main_acc = majority[1] / total_cnt\n",
    "    class_name.append(majority[0])\n",
    "    acc.append(main_acc)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}