notebooks/test_beauty_small_direct_group5_part1.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../')\n",
    "\n",
    "import collections\n",
    "import os\n",
    "import random\n",
    "from pathlib import Path\n",
    "import logging\n",
    "import shutil\n",
    "import time\n",
    "from packaging import version\n",
    "from collections import defaultdict\n",
    "\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import gzip\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn.parallel import DistributedDataParallel as DDP\n",
    "import torch.distributed as dist\n",
    "import torch.backends.cudnn as cudnn\n",
    "\n",
    "from src.param import parse_args\n",
    "from src.utils import LossMeter\n",
    "from src.dist_utils import reduce_dict\n",
    "from transformers import T5Tokenizer, T5TokenizerFast\n",
    "from src.tokenization import P5Tokenizer, P5TokenizerFast\n",
    "from src.pretrain_model import P5Pretraining\n",
    "\n",
    "_use_native_amp = False\n",
    "_use_apex = False\n",
    "\n",
    "# Check if Pytorch version >= 1.6 to switch between Native AMP and Apex\n",
    "if version.parse(torch.__version__) < version.parse(\"1.6\"):\n",
    "    from transormers.file_utils import is_apex_available\n",
    "    if is_apex_available():\n",
    "        from apex import amp\n",
    "    _use_apex = True\n",
    "else:\n",
    "    _use_native_amp = True\n",
    "    from torch.cuda.amp import autocast\n",
    "\n",
    "from src.trainer_base import TrainerBase\n",
    "\n",
    "import pickle\n",
    "\n",
    "def load_pickle(filename):\n",
    "    with open(filename, \"rb\") as f:\n",
    "        return pickle.load(f)\n",
    "\n",
    "\n",
    "def save_pickle(data, filename):\n",
    "    with open(filename, \"wb\") as f:\n",
    "        pickle.dump(data, f, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "        \n",
    "import json\n",
    "\n",
    "def load_json(file_path):\n",
    "    with open(file_path, \"r\") as f:\n",
    "        return json.load(f)\n",
    "    \n",
    "def ReadLineFromFile(path):\n",
    "    lines = []\n",
    "    with open(path,'r') as fd:\n",
    "        for line in fd:\n",
    "            lines.append(line.rstrip('\\n'))\n",
    "    return lines\n",
    "\n",
    "def parse(path):\n",
    "    g = gzip.open(path, 'r')\n",
    "    for l in g:\n",
    "        yield eval(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['rating_loss', 'sequential_loss', 'explanation_loss', 'review_loss', 'traditional_loss']\n",
      "Process Launching at GPU 2\n",
      "{'distributed': False, 'multiGPU': True, 'fp16': True, 'train': 'beauty', 'valid': 'beauty', 'test': 'beauty', 'batch_size': 100, 'optim': 'adamw', 'warmup_ratio': 0.05, 'lr': 0.001, 'num_workers': 4, 'clip_grad_norm': 1.0, 'losses': 'rating,sequential,explanation,review,traditional', 'backbone': 't5-small', 'output': 'snap/beauty-small', 'epoch': 10, 'local_rank': 0, 'comment': '', 'train_topk': -1, 'valid_topk': -1, 'dropout': 0.1, 'tokenizer': 'p5', 'max_text_length': 512, 'do_lower_case': False, 'word_mask_rate': 0.15, 'gen_max_length': 64, 'weight_decay': 0.01, 'adam_eps': 1e-06, 'gradient_accumulation_steps': 1, 'seed': 2022, 'whole_word_embed': True, 'world_size': 4, 'LOSSES_NAME': ['rating_loss', 'sequential_loss', 'explanation_loss', 'review_loss', 'traditional_loss', 'total_loss'], 'gpu': 2, 'rank': 2, 'run_name': 'Jul31_16-50_GPU4_beauty_t5-small_ratingsequentialexplanationreviewtraditional'}\n"
     ]
    }
   ],
   "source": [
    "class DotDict(dict):\n",
    "    def __init__(self, **kwds):\n",
    "        self.update(kwds)\n",
    "        self.__dict__ = self\n",
    "        \n",
    "args = DotDict()\n",
    "\n",
    "args.distributed = False\n",
    "args.multiGPU = True\n",
    "args.fp16 = True\n",
    "args.train = \"beauty\"\n",
    "args.valid = \"beauty\"\n",
    "args.test = \"beauty\"\n",
    "args.batch_size = 100\n",
    "args.optim = 'adamw' \n",
    "args.warmup_ratio = 0.05\n",
    "args.lr = 1e-3\n",
    "args.num_workers = 4\n",
    "args.clip_grad_norm = 1.0\n",
    "args.losses = 'rating,sequential,explanation,review,traditional'\n",
    "args.backbone = 't5-small' # small or base\n",
    "args.output = 'snap/beauty-small'\n",
    "args.epoch = 10\n",
    "args.local_rank = 0\n",
    "\n",
    "args.comment = ''\n",
    "args.train_topk = -1\n",
    "args.valid_topk = -1\n",
    "args.dropout = 0.1\n",
    "\n",
    "args.tokenizer = 'p5'\n",
    "args.max_text_length = 512\n",
    "args.do_lower_case = False\n",
    "args.word_mask_rate = 0.15\n",
    "args.gen_max_length = 64\n",
    "\n",
    "args.weight_decay = 0.01\n",
    "args.adam_eps = 1e-6\n",
    "args.gradient_accumulation_steps = 1\n",
    "\n",
    "'''\n",
    "Set seeds\n",
    "'''\n",
    "args.seed = 2022\n",
    "torch.manual_seed(args.seed)\n",
    "random.seed(args.seed)\n",
    "np.random.seed(args.seed)\n",
    "\n",
    "'''\n",
    "Whole word embedding\n",
    "'''\n",
    "args.whole_word_embed = True\n",
    "\n",
    "cudnn.benchmark = True\n",
    "ngpus_per_node = torch.cuda.device_count()\n",
    "args.world_size = ngpus_per_node\n",
    "\n",
    "LOSSES_NAME = [f'{name}_loss' for name in args.losses.split(',')]\n",
    "if args.local_rank in [0, -1]:\n",
    "    print(LOSSES_NAME)\n",
    "LOSSES_NAME.append('total_loss') # total loss\n",
    "\n",
    "args.LOSSES_NAME = LOSSES_NAME\n",
    "\n",
    "gpu = 2 # Change GPU ID\n",
    "args.gpu = gpu\n",
    "args.rank = gpu\n",
    "print(f'Process Launching at GPU {gpu}')\n",
    "\n",
    "torch.cuda.set_device('cuda:{}'.format(gpu))\n",
    "\n",
    "comments = []\n",
    "dsets = []\n",
    "if 'toys' in args.train:\n",
    "    dsets.append('toys')\n",
    "if 'beauty' in args.train:\n",
    "    dsets.append('beauty')\n",
    "if 'sports' in args.train:\n",
    "    dsets.append('sports')\n",
    "comments.append(''.join(dsets))\n",
    "if args.backbone:\n",
    "    comments.append(args.backbone)\n",
    "comments.append(''.join(args.losses.split(',')))\n",
    "if args.comment != '':\n",
    "    comments.append(args.comment)\n",
    "comment = '_'.join(comments)\n",
    "\n",
    "from datetime import datetime\n",
    "current_time = datetime.now().strftime('%b%d_%H-%M')\n",
    "\n",
    "if args.local_rank in [0, -1]:\n",
    "    run_name = f'{current_time}_GPU{args.world_size}'\n",
    "    if len(comments) > 0:\n",
    "        run_name += f'_{comment}'\n",
    "    args.run_name = run_name\n",
    "    print(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_config(args):\n",
    "    from transformers import T5Config, BartConfig\n",
    "\n",
    "    if 't5' in args.backbone:\n",
    "        config_class = T5Config\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "    config = config_class.from_pretrained(args.backbone)\n",
    "    config.dropout_rate = args.dropout\n",
    "    config.dropout = args.dropout\n",
    "    config.attention_dropout = args.dropout\n",
    "    config.activation_dropout = args.dropout\n",
    "    config.losses = args.losses\n",
    "\n",
    "    return config\n",
    "\n",
    "\n",
    "def create_tokenizer(args):\n",
    "    from transformers import T5Tokenizer, T5TokenizerFast\n",
    "    from src.tokenization import P5Tokenizer, P5TokenizerFast\n",
    "\n",
    "    if 'p5' in args.tokenizer:\n",
    "        tokenizer_class = P5Tokenizer\n",
    "\n",
    "    tokenizer_name = args.backbone\n",
    "    \n",
    "    tokenizer = tokenizer_class.from_pretrained(\n",
    "        tokenizer_name,\n",
    "        max_length=args.max_text_length,\n",
    "        do_lower_case=args.do_lower_case,\n",
    "    )\n",
    "\n",
    "    print(tokenizer_class, tokenizer_name)\n",
    "    \n",
    "    return tokenizer\n",
    "\n",
    "\n",
    "def create_model(model_class, config=None):\n",
    "    print(f'Building Model at GPU {args.gpu}')\n",
    "\n",
    "    model_name = args.backbone\n",
    "\n",
    "    model = model_class.from_pretrained(\n",
    "        model_name,\n",
    "        config=config\n",
    "    )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'src.tokenization.P5Tokenizer'> t5-small\n",
      "Building Model at GPU 2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of P5Pretraining were not initialized from the model checkpoint at t5-small and are newly initialized: ['encoder.whole_word_embeddings.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "config = create_config(args)\n",
    "\n",
    "if args.tokenizer is None:\n",
    "    args.tokenizer = args.backbone\n",
    "    \n",
    "tokenizer = create_tokenizer(args)\n",
    "\n",
    "model_class = P5Pretraining\n",
    "model = create_model(model_class, config)\n",
    "\n",
    "model = model.cuda()\n",
    "\n",
    "if 'p5' in args.tokenizer:\n",
    "    model.resize_token_embeddings(tokenizer.vocab_size)\n",
    "    \n",
    "model.tokenizer = tokenizer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model loaded from  ../snap/beauty-small-group5-part1.pth\n",
      "<All keys matched successfully>\n"
     ]
    }
   ],
   "source": [
    "args.load = \"../snap/beauty-small-group5-part1.pth\"\n",
    "\n",
    "# Load Checkpoint\n",
    "from src.utils import load_state_dict, LossMeter, set_global_logging_level\n",
    "from pprint import pprint\n",
    "\n",
    "def load_checkpoint(ckpt_path):\n",
    "    state_dict = load_state_dict(ckpt_path, 'cpu')\n",
    "    results = model.load_state_dict(state_dict, strict=False)\n",
    "    print('Model loaded from ', ckpt_path)\n",
    "    pprint(results)\n",
    "\n",
    "ckpt_path = args.load\n",
    "load_checkpoint(ckpt_path)\n",
    "\n",
    "from src.all_amazon_templates import all_tasks as task_templates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check Test Split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_splits = load_pickle('../data/beauty/rating_splits_augmented.pkl')\n",
    "test_review_data = data_splits['test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "19850"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(test_review_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'reviewerID': 'A2QKXW3LDQ66P5',\n",
       " 'asin': 'B005X2F7KI',\n",
       " 'reviewerName': 'stephanie',\n",
       " 'helpful': [5, 6],\n",
       " 'reviewText': 'Absolutely great product.  I bought this for my fourteen year old niece for Christmas and of course I had to try it out, then I tried another one, and another one and another one.  So much fun!  I even contemplated keeping a few for myself!',\n",
       " 'overall': 5.0,\n",
       " 'summary': 'Perfect!',\n",
       " 'unixReviewTime': 1352937600,\n",
       " 'reviewTime': '11 15, 2012',\n",
       " 'explanation': 'Absolutely great product',\n",
       " 'feature': 'product'}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_review_data[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "22363\n",
      "12101\n"
     ]
    }
   ],
   "source": [
    "data_maps = load_json(os.path.join('../data', 'beauty', 'datamaps.json'))\n",
    "print(len(data_maps['user2id'])) # number of users\n",
    "print(len(data_maps['item2id'])) # number of items"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Test P5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader, Dataset, Sampler\n",
    "from src.pretrain_data import get_loader\n",
    "from evaluate.utils import rouge_score, bleu_score, unique_sentence_percent, root_mean_square_error, mean_absolute_error, feature_detect, feature_matching_ratio, feature_coverage_ratio, feature_diversity\n",
    "from evaluate.metrics4rec import evaluate_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from multiprocessing import Pool\n",
    "import math\n",
    "from torch.utils.data.distributed import DistributedSampler\n",
    "from copy import deepcopy\n",
    "\n",
    "def load_json(file_path):\n",
    "    with open(file_path, \"r\") as f:\n",
    "        return json.load(f)\n",
    "\n",
    "def load_pickle(filename):\n",
    "    with open(filename, \"rb\") as f:\n",
    "        return pickle.load(f)\n",
    "\n",
    "def ReadLineFromFile(path):\n",
    "    lines = []\n",
    "    with open(path,'r') as fd:\n",
    "        for line in fd:\n",
    "            lines.append(line.rstrip('\\n'))\n",
    "    return lines\n",
    "\n",
    "def parse(path):\n",
    "    g = gzip.open(path, 'r')\n",
    "    for l in g:\n",
    "        yield eval(l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class P5_Amazon_Dataset(Dataset):\n",
    "    def __init__(self, all_tasks, task_list, tokenizer, args, sample_numbers, mode='train', split='toys', rating_augment=False, sample_type='random'): \n",
    "        self.all_tasks = all_tasks\n",
    "        self.task_list = task_list\n",
    "        self.tokenizer = tokenizer\n",
    "        self.args = args\n",
    "        self.sample_numbers = sample_numbers\n",
    "        self.split = split\n",
    "        self.rating_augment = rating_augment\n",
    "        self.sample_type = sample_type\n",
    "        \n",
    "        print('Data sources: ', split.split(','))\n",
    "        self.mode = mode\n",
    "        assert self.mode == 'test'\n",
    "        if self.mode == 'test':\n",
    "            self.review_data = load_pickle(os.path.join('data', split, 'review_splits.pkl'))['test']\n",
    "            self.exp_data = load_pickle(os.path.join('data', split, 'exp_splits.pkl'))['test']\n",
    "            if self.rating_augment:\n",
    "                self.rating_data = load_pickle(os.path.join('data', split, 'rating_splits_augmented.pkl'))['test']\n",
    "            else:\n",
    "                self.rating_data = self.review_data\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "            \n",
    "        self.sequential_data = ReadLineFromFile(os.path.join('data', split, 'sequential_data.txt'))\n",
    "        item_count = defaultdict(int)\n",
    "        user_items = defaultdict()\n",
    "\n",
    "        for line in self.sequential_data:\n",
    "            user, items = line.strip().split(' ', 1)\n",
    "            items = items.split(' ')\n",
    "            items = [int(item) for item in items]\n",
    "            user_items[user] = items\n",
    "            for item in items:\n",
    "                item_count[item] += 1\n",
    "                \n",
    "        self.all_item = list(item_count.keys())\n",
    "        count = list(item_count.values())\n",
    "        sum_value = np.sum([x for x in count])\n",
    "        self.probability = [value / sum_value for value in count]\n",
    "        self.user_items = user_items\n",
    "        \n",
    "        if self.mode == 'test':\n",
    "            self.negative_samples = ReadLineFromFile(os.path.join('data', split, 'negative_samples.txt'))\n",
    "            \n",
    "        datamaps = load_json(os.path.join('data', split, 'datamaps.json'))\n",
    "        self.user2id = datamaps['user2id']\n",
    "        self.item2id = datamaps['item2id']\n",
    "        self.user_list = list(datamaps['user2id'].keys())\n",
    "        self.item_list = list(datamaps['item2id'].keys())\n",
    "        self.id2item = datamaps['id2item']\n",
    "        \n",
    "        self.user_id2name = load_pickle(os.path.join('data', split, 'user_id2name.pkl'))\n",
    "        \n",
    "        self.meta_data = []\n",
    "        for meta in parse(os.path.join('data', split, 'meta.json.gz')):\n",
    "            self.meta_data.append(meta)\n",
    "        self.meta_dict = {}\n",
    "        for i, meta_item in enumerate(self.meta_data):\n",
    "            self.meta_dict[meta_item['asin']] = i\n",
    "            \n",
    "        print('compute_datum_info')\n",
    "        self.total_length = 0\n",
    "        self.datum_info = []\n",
    "        self.compute_datum_info()\n",
    "        \n",
    "    def compute_datum_info(self):\n",
    "        curr = 0\n",
    "        assert 'traditional' in self.task_list.keys()\n",
    "        key = 'traditional'\n",
    "        assert 0 < int(self.task_list[key].split('-')[1]) <= 4\n",
    "        self.total_length += len(self.user2id) * 100\n",
    "        for i in range(self.total_length - curr):\n",
    "            self.datum_info.append((i + curr, key, i // 100, i % 100))\n",
    "        curr = self.total_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.total_length\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        \n",
    "        out_dict = {}\n",
    "        out_dict['args'] = self.args\n",
    "        \n",
    "        loss_weight = 1.0\n",
    "        \n",
    "        datum_info_idx = self.datum_info[idx]\n",
    "        assert datum_info_idx[0] == idx\n",
    "        if len(datum_info_idx) == 4:\n",
    "            task_name = datum_info_idx[1]\n",
    "            datum_idx = datum_info_idx[2]\n",
    "            candidate_item_idx = datum_info_idx[3]\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "            \n",
    "        if task_name == 'traditional':\n",
    "            sequential_datum = self.sequential_data[datum_idx]\n",
    "            sequence = sequential_datum.split()\n",
    "            user_id = sequence[0]\n",
    "            user_desc = self.user_id2name[user_id]\n",
    "            assert self.mode == 'test'\n",
    "            if candidate_item_idx == 0:\n",
    "                target_item = sequence[-1]\n",
    "            \n",
    "            task_candidates = self.task_list[task_name]\n",
    "            task_template = self.all_tasks['traditional'][task_candidates]\n",
    "            assert task_template['task'] == 'traditional'\n",
    "            \n",
    "            if task_template['id'] == '5-1':\n",
    "                if candidate_item_idx == 0:\n",
    "                    source_text = task_template['source'].format(user_id, target_item)\n",
    "                    target_text = task_template['target'].format('yes')\n",
    "                else:\n",
    "                    assert user_id == self.negative_samples[int(user_id)-1].split(' ', 1)[0]\n",
    "                    candidate_samples = self.negative_samples[int(user_id)-1].split(' ', 1)[1].split(' ')\n",
    "                    source_text = task_template['source'].format(user_id, candidate_samples[candidate_item_idx-1])\n",
    "                    target_text = task_template['target'].format('no')\n",
    "            elif task_template['id'] == '5-2':\n",
    "                if candidate_item_idx == 0:\n",
    "                    source_text = task_template['source'].format(target_item, user_desc)\n",
    "                    target_text = task_template['target'].format('yes')\n",
    "                else:\n",
    "                    assert user_id == self.negative_samples[int(user_id)-1].split(' ', 1)[0]\n",
    "                    candidate_samples = self.negative_samples[int(user_id)-1].split(' ', 1)[1].split(' ')\n",
    "                    source_text = task_template['source'].format(candidate_samples[candidate_item_idx-1], user_desc)\n",
    "                    target_text = task_template['target'].format('no')\n",
    "            elif task_template['id'] == '5-3':\n",
    "                if candidate_item_idx == 0:\n",
    "                    if 'title' in self.meta_data[self.meta_dict[self.id2item[target_item]]]:\n",
    "                        title = self.meta_data[self.meta_dict[self.id2item[target_item]]]['title']\n",
    "                    else:\n",
    "                        title = 'unknown title'\n",
    "                    source_text = task_template['source'].format(user_desc, title)\n",
    "                    target_text = task_template['target'].format('yes')\n",
    "                else:\n",
    "                    assert user_id == self.negative_samples[int(user_id)-1].split(' ', 1)[0]\n",
    "                    candidate_samples = self.negative_samples[int(user_id)-1].split(' ', 1)[1].split(' ')\n",
    "                    if 'title' in self.meta_data[self.meta_dict[self.id2item[candidate_samples[candidate_item_idx-1]]]]:\n",
    "                        title = self.meta_data[self.meta_dict[self.id2item[candidate_samples[candidate_item_idx-1]]]]['title']\n",
    "                    else:\n",
    "                        title = 'unknown title'\n",
    "                    source_text = task_template['source'].format(user_desc, title)\n",
    "                    target_text = task_template['target'].format('no')\n",
    "            elif task_template['id'] == '5-4':\n",
    "                if candidate_item_idx == 0:\n",
    "                    if 'title' in self.meta_data[self.meta_dict[self.id2item[target_item]]]:\n",
    "                        title = self.meta_data[self.meta_dict[self.id2item[target_item]]]['title']\n",
    "                    else:\n",
    "                        title = 'unknown title'\n",
    "                    source_text = task_template['source'].format(user_id, title)\n",
    "                    target_text = task_template['target'].format('yes')\n",
    "                else:\n",
    "                    assert user_id == self.negative_samples[int(user_id)-1].split(' ', 1)[0]\n",
    "                    candidate_samples = self.negative_samples[int(user_id)-1].split(' ', 1)[1].split(' ')\n",
    "                    if 'title' in self.meta_data[self.meta_dict[self.id2item[candidate_samples[candidate_item_idx-1]]]]:\n",
    "                        title = self.meta_data[self.meta_dict[self.id2item[candidate_samples[candidate_item_idx-1]]]]['title']\n",
    "                    else:\n",
    "                        title = 'unknown title'\n",
    "                    source_text = task_template['source'].format(user_id, title)\n",
    "                    target_text = task_template['target'].format('no')\n",
    "            else:\n",
    "                raise NotImplementedError\n",
    "            \n",
    "        input_ids = self.tokenizer.encode(\n",
    "                source_text, padding=True, truncation=True, max_length=self.args.max_text_length)\n",
    "        tokenized_text = self.tokenizer.tokenize(source_text)\n",
    "        whole_word_ids = self.calculate_whole_word_ids(tokenized_text, input_ids)\n",
    "        assert len(whole_word_ids) == len(input_ids)\n",
    "        \n",
    "        target_ids = self.tokenizer.encode(\n",
    "                target_text, padding=True, truncation=True, max_length=self.args.gen_max_length)\n",
    "\n",
    "        out_dict['input_ids'] = torch.LongTensor(input_ids)\n",
    "        out_dict['input_length'] = len(input_ids)\n",
    "        out_dict['whole_word_ids'] = torch.LongTensor(whole_word_ids)\n",
    "        out_dict['target_ids'] = torch.LongTensor(target_ids)\n",
    "        out_dict['target_length'] = len(target_ids)\n",
    "\n",
    "        out_dict['source_text'] = source_text\n",
    "        out_dict['tokenized_text'] = tokenized_text\n",
    "        out_dict['target_text'] = target_text\n",
    "\n",
    "        out_dict['task'] = task_template['task']\n",
    "\n",
    "        out_dict['loss_weight'] = loss_weight\n",
    "\n",
    "        return out_dict\n",
    "    \n",
    "    def calculate_whole_word_ids(self, tokenized_text, input_ids):\n",
    "        whole_word_ids = []\n",
    "        curr = 0\n",
    "        for i in range(len(tokenized_text)):\n",
    "            if tokenized_text[i].startswith('▁'):\n",
    "                curr += 1\n",
    "                whole_word_ids.append(curr)\n",
    "            else:\n",
    "                whole_word_ids.append(curr)\n",
    "        last_item = whole_word_ids[len(input_ids) - 2]\n",
    "        return whole_word_ids[:len(input_ids) - 1] + [0] ## the added [0] is for </s>\n",
    "    \n",
    "    def collate_fn(self, batch):\n",
    "        batch_entry = {}\n",
    "\n",
    "        B = len(batch)\n",
    "\n",
    "        args = self.args\n",
    "\n",
    "        S_W_L = max(entry['input_length'] for entry in batch)\n",
    "        T_W_L = max(entry['target_length'] for entry in batch)\n",
    "\n",
    "        input_ids = torch.ones(B, S_W_L, dtype=torch.long) * self.tokenizer.pad_token_id\n",
    "        whole_word_ids = torch.ones(B, S_W_L, dtype=torch.long) * self.tokenizer.pad_token_id\n",
    "        target_ids = torch.ones(B, T_W_L, dtype=torch.long) * self.tokenizer.pad_token_id\n",
    "\n",
    "        loss_weights = torch.ones(B, dtype=torch.float)\n",
    "\n",
    "        tasks = []\n",
    "        source_text = []\n",
    "        tokenized_text = []\n",
    "        target_text = []\n",
    "\n",
    "        for i, entry in enumerate(batch):\n",
    "            input_ids[i, :entry['input_length']] = entry['input_ids']\n",
    "            whole_word_ids[i, :entry['input_length']] = entry['whole_word_ids']\n",
    "            target_ids[i, :entry['target_length']] = entry['target_ids']\n",
    "\n",
    "            if 'task' in entry:\n",
    "                tasks.append(entry['task'])\n",
    "\n",
    "            if 'source_text' in entry:\n",
    "                source_text.append(entry['source_text'])\n",
    "                \n",
    "            if 'tokenized_text' in entry:\n",
    "                tokenized_text.append(entry['tokenized_text'])\n",
    "                \n",
    "            if 'target_text' in entry:\n",
    "                target_text.append(entry['target_text'])\n",
    "\n",
    "            if 'loss_weight' in entry:\n",
    "                ## length-aware loss normalization\n",
    "                if entry['target_length'] > 0:\n",
    "                    loss_weights[i] = entry['loss_weight'] / entry['target_length']\n",
    "                else:\n",
    "                    loss_weights[i] = entry['loss_weight']\n",
    "\n",
    "        assert 't5' in args.backbone\n",
    "        word_mask = target_ids != self.tokenizer.pad_token_id\n",
    "        target_ids[~word_mask] = -100\n",
    "        batch_entry['task'] = tasks\n",
    "\n",
    "        batch_entry['source_text'] = source_text\n",
    "        batch_entry['target_text'] = target_text\n",
    "\n",
    "        batch_entry['input_ids'] = input_ids\n",
    "        batch_entry['whole_word_ids'] = whole_word_ids\n",
    "        batch_entry['target_ids'] = target_ids\n",
    "\n",
    "        batch_entry['loss_weights'] = loss_weights\n",
    "\n",
    "        return batch_entry\n",
    "    \n",
    "\n",
    "def get_loader(args, task_list, sample_numbers, split='toys', mode='train', \n",
    "               batch_size=16, workers=4, distributed=False):\n",
    "\n",
    "    if 't5' in args.backbone:\n",
    "        tokenizer = P5Tokenizer.from_pretrained(\n",
    "            args.backbone, \n",
    "            max_length=args.max_text_length, \n",
    "            do_lower_case=args.do_lower_case)\n",
    "\n",
    "    if split == 'yelp':\n",
    "        from all_yelp_templates import all_tasks as task_templates\n",
    "        \n",
    "        dataset = P5_Yelp_Dataset(\n",
    "            task_templates,\n",
    "            task_list,\n",
    "            tokenizer,\n",
    "            args,\n",
    "            sample_numbers,\n",
    "            mode=mode,\n",
    "            split=split,\n",
    "            rating_augment=False\n",
    "        )\n",
    "    else:\n",
    "        from all_amazon_templates import all_tasks as task_templates\n",
    "\n",
    "        dataset = P5_Amazon_Dataset(\n",
    "            task_templates,\n",
    "            task_list,\n",
    "            tokenizer,\n",
    "            args,\n",
    "            sample_numbers,\n",
    "            mode=mode,\n",
    "            split=split,\n",
    "            rating_augment=False\n",
    "        )\n",
    "\n",
    "    if distributed:\n",
    "        sampler = DistributedSampler(dataset)\n",
    "    else:\n",
    "        sampler = None\n",
    "\n",
    "    if mode == 'train':\n",
    "        loader = DataLoader(\n",
    "            dataset, batch_size=batch_size, shuffle=(sampler is None),\n",
    "            num_workers=workers, pin_memory=True, sampler=sampler,\n",
    "            collate_fn=dataset.collate_fn)\n",
    "    else:\n",
    "        loader = DataLoader(\n",
    "            dataset,\n",
    "            batch_size=batch_size,\n",
    "            num_workers=workers, pin_memory=True,\n",
    "            sampler=sampler,\n",
    "            shuffle=None if (sampler is not None) else False,\n",
    "            collate_fn=dataset.collate_fn,\n",
    "            drop_last=False)\n",
    "        \n",
    "    return loader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluation - Traditional (Prompt 5-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data sources:  ['beauty']\n",
      "compute_datum_info\n",
      "22363\n"
     ]
    }
   ],
   "source": [
    "test_task_list = {'traditional': '5-1'\n",
    "}\n",
    "test_sample_numbers = {'traditional': 100}\n",
    "\n",
    "zeroshot_test_loader = get_loader(\n",
    "        args,\n",
    "        test_task_list,\n",
    "        test_sample_numbers,\n",
    "        split=args.test,\n",
    "        mode='test', \n",
    "        batch_size=args.batch_size,\n",
    "        workers=args.num_workers,\n",
    "        distributed=args.distributed\n",
    ")\n",
    "print(len(zeroshot_test_loader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "m = nn.Softmax(dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "22363it [12:43, 29.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "NDCG@1\tRec@1\tHits@1\tPrec@1\tMAP@1\tMRR@1\n",
      "0.0600\t0.0600\t0.0600\t0.0600\t0.0600\t0.0600\n",
      "\n",
      "NDCG@5\tRec@5\tHits@5\tPrec@5\tMAP@5\tMRR@5\n",
      "0.1316\t0.2021\t0.2021\t0.0404\t0.1086\t0.1086\n",
      "\n",
      "NDCG@10\tRec@10\tHits@10\tPrec@10\tMAP@10\tMRR@10\n",
      "0.1670\t0.3121\t0.3121\t0.0312\t0.1230\t0.1230\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('\\nNDCG@10\\tRec@10\\tHits@10\\tPrec@10\\tMAP@10\\tMRR@10\\n0.1670\\t0.3121\\t0.3121\\t0.0312\\t0.1230\\t0.1230',\n",
       " {'ndcg': 0.16695003951866888,\n",
       "  'map': 0.12298331143633652,\n",
       "  'recall': 0.3121227026785315,\n",
       "  'precision': 0.03121227026785718,\n",
       "  'mrr': 0.12298331143633652,\n",
       "  'hit': 0.3121227026785315})"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_info = []\n",
    "for i, batch in tqdm(enumerate(zeroshot_test_loader)):\n",
    "    with torch.no_grad():\n",
    "        beam_outputs = model.generate(\n",
    "            batch['input_ids'].to('cuda'), \n",
    "            max_length=50, \n",
    "            num_beams=1, \n",
    "            no_repeat_ngram_size=2, \n",
    "            num_return_sequences=1, \n",
    "            early_stopping=True,\n",
    "            output_scores=True,\n",
    "            return_dict_in_generate=True\n",
    "        )\n",
    "        gen_yes_probs = m(beam_outputs.scores[1][:, [4273, 150]])[:, 0].cpu() # 4273 -> 'yes', 150 -> 'no'\n",
    "        sorted, indices = torch.sort(gen_yes_probs, descending=True)\n",
    "        new_info = {}\n",
    "        all_item_ids = [_.split(' ?')[0].split('item_')[-1] for _ in batch['source_text']]\n",
    "        new_info['target_item'] = all_item_ids[0]\n",
    "        new_info['gen_item_list'] = [all_item_ids[_] for _ in indices[:20]]\n",
    "        all_info.append(new_info)\n",
    "        \n",
    "gt = {}\n",
    "ui_scores = {}\n",
    "for i, info in enumerate(all_info):\n",
    "    gt[i] = [int(info['target_item'])]\n",
    "    pred_dict = {}\n",
    "    for j in range(len(info['gen_item_list'])):\n",
    "        try:\n",
    "            pred_dict[int(info['gen_item_list'][j])] = -(j+1)\n",
    "        except:\n",
    "            pass\n",
    "    ui_scores[i] = pred_dict\n",
    "    \n",
    "evaluate_all(ui_scores, gt, 1)\n",
    "evaluate_all(ui_scores, gt, 5)\n",
    "evaluate_all(ui_scores, gt, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "22363it [07:10, 51.93it/s]\n"
     ]
    }
   ],
   "source": [
    "all_item_ids = []\n",
    "for i, batch in tqdm(enumerate(zeroshot_test_loader)):\n",
    "    all_item_ids.append([_.split(' ?')[0].split('item_')[-1] for _ in batch['source_text']])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Evaluation - Traditional (Prompt 5-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data sources:  ['beauty']\n",
      "compute_datum_info\n",
      "22363\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "22363it [17:20, 21.48it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "NDCG@1\tRec@1\tHits@1\tPrec@1\tMAP@1\tMRR@1\n",
      "0.0862\t0.0862\t0.0862\t0.0862\t0.0862\t0.0862\n",
      "\n",
      "NDCG@5\tRec@5\tHits@5\tPrec@5\tMAP@5\tMRR@5\n",
      "0.1673\t0.2448\t0.2448\t0.0490\t0.1419\t0.1419\n",
      "\n",
      "NDCG@10\tRec@10\tHits@10\tPrec@10\tMAP@10\tMRR@10\n",
      "0.1993\t0.3441\t0.3441\t0.0344\t0.1550\t0.1550\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "('\\nNDCG@10\\tRec@10\\tHits@10\\tPrec@10\\tMAP@10\\tMRR@10\\n0.1993\\t0.3441\\t0.3441\\t0.0344\\t0.1550\\t0.1550',\n",
       " {'ndcg': 0.19929096493386037,\n",
       "  'map': 0.15497093768121806,\n",
       "  'recall': 0.34405044045968786,\n",
       "  'precision': 0.03440504404597355,\n",
       "  'mrr': 0.15497093768121806,\n",
       "  'hit': 0.34405044045968786})"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_task_list = {'traditional': '5-4'\n",
    "}\n",
    "test_sample_numbers = {'traditional': 100}\n",
    "\n",
    "zeroshot_test_loader = get_loader(\n",
    "        args,\n",
    "        test_task_list,\n",
    "        test_sample_numbers,\n",
    "        split=args.test,\n",
    "        mode='test', \n",
    "        batch_size=args.batch_size,\n",
    "        workers=args.num_workers,\n",
    "        distributed=args.distributed\n",
    ")\n",
    "print(len(zeroshot_test_loader))\n",
    "\n",
    "m = nn.Softmax(dim=1)\n",
    "\n",
    "all_info = []\n",
    "for i, batch in tqdm(enumerate(zeroshot_test_loader)):\n",
    "    with torch.no_grad():\n",
    "        beam_outputs = model.generate(\n",
    "            batch['input_ids'].to('cuda'), \n",
    "            max_length=50, \n",
    "            num_beams=1, \n",
    "            no_repeat_ngram_size=2, \n",
    "            num_return_sequences=1, \n",
    "            early_stopping=True,\n",
    "            output_scores=True,\n",
    "            return_dict_in_generate=True\n",
    "        )\n",
    "        gen_yes_probs = m(beam_outputs.scores[1][:, [4273, 150]])[:, 0].cpu() # 4273 -> 'yes', 150 -> 'no'\n",
    "        sorted, indices = torch.sort(gen_yes_probs, descending=True)\n",
    "        new_info = {}\n",
    "        new_info['target_item'] = all_item_ids[i][0]\n",
    "        new_info['gen_item_list'] = [all_item_ids[i][_] for _ in indices[:20]]\n",
    "        all_info.append(new_info)\n",
    "        \n",
    "gt = {}\n",
    "ui_scores = {}\n",
    "for i, info in enumerate(all_info):\n",
    "    gt[i] = [int(info['target_item'])]\n",
    "    pred_dict = {}\n",
    "    for j in range(len(info['gen_item_list'])):\n",
    "        try:\n",
    "            pred_dict[int(info['gen_item_list'][j])] = -(j+1)\n",
    "        except:\n",
    "            pass\n",
    "    ui_scores[i] = pred_dict\n",
    "    \n",
    "evaluate_all(ui_scores, gt, 1)\n",
    "evaluate_all(ui_scores, gt, 5)\n",
    "evaluate_all(ui_scores, gt, 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pt-1.10",
   "language": "python",
   "name": "pt-1.10"
  },
  "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}