{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/gscratch/raivn/ethans/miniconda3/envs/llms_12.1/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "<frozen importlib._bootstrap>:241: RuntimeWarning: pyarrow.lib.IpcWriteOptions size changed, may indicate binary incompatibility. Expected 72 from C header, got 88 from PyObject\n",
      "<frozen importlib._bootstrap>:241: RuntimeWarning: pyarrow.lib.IpcReadOptions size changed, may indicate binary incompatibility. Expected 96 from C header, got 104 from PyObject\n",
      "<frozen importlib._bootstrap>:241: RuntimeWarning: pyarrow._fs.FileInfo size changed, may indicate binary incompatibility. Expected 64 from C header, got 88 from PyObject\n",
      "<frozen importlib._bootstrap>:241: RuntimeWarning: pyarrow._fs.FileSelector size changed, may indicate binary incompatibility. Expected 48 from C header, got 72 from PyObject\n",
      "2024-05-30 01:35:17.813978: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2024-05-30 01:35:20.452213: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-05-30 01:35:41.833487: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import copy\n",
    "import json\n",
    "import pickle\n",
    "import os\n",
    "import random\n",
    "import re\n",
    "import string\n",
    "import math\n",
    "from datetime import datetime\n",
    "\n",
    "import evaluate\n",
    "import torch\n",
    "import numpy as np\n",
    "from datasets import load_dataset\n",
    "from transformers import LlamaTokenizer\n",
    "from tqdm import tqdm\n",
    "\n",
    "from eval import *\n",
    "from superposed.llama.metrics import *\n",
    "from superposed.llama.generation import Llama\n",
    "from superposed.llama.superposed_generation import SuperposedLlama\n",
    "from superposed.llama.tokenizer import Tokenizer\n",
    "from superposed.ngrams.ngram_models import make_models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parameters: {'alpha': 0.54, 'temp': 0.06, 'n_drafts': 3, 'prompt_len': 15, 'n_token_sample': 9, 'n_token_consider': 32000, 'mixing_method': 'sample_new_weights_with_score', 'smoothing': 'geom', 'sample_tokens': 0, 'sample_beams': 0, 'i_weights': [0.01, 0.04, 0.15, 0.18, 0.12], 'i_length': [1, 2, 3, 4, 5]}\n"
     ]
    }
   ],
   "source": [
    "# Params\n",
    "param_file = \"../../params/p15_d3_mixed.json\"\n",
    "with open(param_file, \"r\") as f:\n",
    "    params = json.load(f)\n",
    "    print(f\"Parameters: {params}\")\n",
    "alpha = params[\"alpha\"]\n",
    "temp = params[\"temp\"]\n",
    "n_drafts = params[\"n_drafts\"]\n",
    "prompt_len = params[\"prompt_len\"]\n",
    "n_token_sample = params[\"n_token_sample\"]\n",
    "i_weights = params[\"i_weights\"]\n",
    "i_length = params[\"i_length\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Making bigram...\n",
      "1310800\n",
      "Making trigram...\n",
      "671088728\n",
      "Making fourgram...\n",
      "2684354648\n",
      "Making fivegram...\n",
      "5368709200\n",
      "Making sixgram...\n",
      "5368709200\n"
     ]
    }
   ],
   "source": [
    "ngrams = make_models(\"../../ckpts-200k\", bigram=True, trigram=True, fourgram=True, fivegram=True, sixgram=True, sevengram=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "sup_device = torch.device(\"cuda:0\")\n",
    "reg_device = torch.device(\"cuda:1\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> initializing model parallel with size 1\n",
      "> initializing ddp with size 1\n",
      "> initializing pipeline with size 1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/gscratch/raivn/ethans/miniconda3/envs/llms_12.1/lib/python3.11/site-packages/torch/__init__.py:614: UserWarning: torch.set_default_tensor_type() is deprecated as of PyTorch 2.1, please use torch.set_default_dtype() and torch.set_default_device() as alternatives. (Triggered internally at ../torch/csrc/tensor/python_tensor.cpp:451.)\n",
      "  _C._set_default_tensor_type(t)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded in 22.07 seconds\n",
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "weight_path = \"../../7B/\"\n",
    "sup_model = SuperposedLlama.build(ckpt_dir=weight_path, \n",
    "                                 tokenizer_path=f'{weight_path}/tokenizer.model', \n",
    "                                 max_seq_len=1000, \n",
    "                                 max_batch_size=16,\n",
    "                                 device=sup_device,\n",
    "                                 model_parallel_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "Loaded in 22.76 seconds\n"
     ]
    }
   ],
   "source": [
    "reg_model = Llama.build(ckpt_dir=weight_path, \n",
    "                    tokenizer_path=f'{weight_path}/tokenizer.model', \n",
    "                    max_seq_len=1000, \n",
    "                    max_batch_size=16,\n",
    "                    device=reg_device,\n",
    "                    model_parallel_size=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = Tokenizer(f\"{weight_path}/tokenizer.model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length: 7993\n"
     ]
    }
   ],
   "source": [
    "trivia_path = \"../../../datasets/qa/wikipedia-dev.json\"\n",
    "with open(trivia_path, \"r\") as f:\n",
    "    triviaqa = json.load(f)[\"Data\"]\n",
    "print(f\"Length: {len(triviaqa)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.set_default_dtype(torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_types = [\"superposed\", \"regular\"]\n",
    "model_type = model_types[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/triviaqa/default.yaml\n",
    "def evaluate_trivia(model_type, question, max_gen_len):\n",
    "    question = \"Question: \" + question + \"\\nAnswer:\"\n",
    "    text_len = len(question) # for truncating\n",
    "    prompt_len = len(tokenizer.encode([question], True, False)[0]) # for model\n",
    "    if model_type == \"regular\":\n",
    "        input = [question for _ in range(n_drafts)]\n",
    "        sequences, _ = evaluate_nucleus_losses(data=input,\n",
    "                                               model=reg_model,\n",
    "                                               tokenizer=tokenizer,\n",
    "                                               prompt_len=prompt_len,\n",
    "                                               max_gen_len=max_gen_len,\n",
    "                                               temp=0.6, # Set to 0 for greedy\n",
    "                                               bsz=8,\n",
    "                                               marker=False)\n",
    "        n_pd, seq_len = sequences.shape\n",
    "    elif model_type == \"superposed\":\n",
    "        sequences, _ = evaluate_mixed_losses(data=[question],\n",
    "                                                   model=sup_model,\n",
    "                                                   tokenizer=tokenizer,\n",
    "                                                   prompt_len=prompt_len,\n",
    "                                                   max_gen_len=max_gen_len,\n",
    "                                                   alpha=alpha,\n",
    "                                                   temp=temp,\n",
    "                                                   n_drafts=n_drafts,\n",
    "                                                   n_token_sample=n_token_sample,\n",
    "                                                   smoothing=None, # greedy\n",
    "                                                   bsz=8,\n",
    "                                                   i_weights=i_weights,\n",
    "                                                   i_length=i_length,\n",
    "                                                   ngrams=ngrams,\n",
    "                                                   marker=False)\n",
    "        n_p, n_d, seq_len = sequences.shape\n",
    "    # Process results\n",
    "    sequences = sequences.reshape(-1, seq_len).tolist()\n",
    "    for d_idx in range(len(sequences)):\n",
    "        draft = sequences[d_idx]\n",
    "        if -1 in draft:\n",
    "            draft = draft[:draft.index(-1)]\n",
    "        sequences[d_idx] = draft\n",
    "    decoded_seq = tokenizer.decode(sequences)\n",
    "    answers = []\n",
    "    for s in decoded_seq:\n",
    "        # print(s)\n",
    "        answers.append(re.split(\"[,.\\n]\", s[text_len:].strip())[0])\n",
    "    return answers\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "questions = {}\n",
    "predictions = {}\n",
    "print(f\"Precision from 1 to {n_drafts}\")\n",
    "for sample in tqdm(triviaqa):\n",
    "    # Adaptively select generation length\n",
    "    longest = 0\n",
    "    shortest = 1000\n",
    "    total = 0\n",
    "    for answer in sample[\"Answer\"][\"Aliases\"]:\n",
    "        tmp = tokenizer.encode([answer], False, False)[0]\n",
    "        if len(tmp) > longest:\n",
    "            longest = len(tmp)\n",
    "        if len(tmp) < shortest:\n",
    "            shortest = len(tmp)\n",
    "        total += len(tmp)\n",
    "    # Evaluation code\n",
    "    id = sample[\"QuestionId\"]\n",
    "    question = sample[\"Question\"]\n",
    "    answer = evaluate_trivia(model_type, question, max_gen_len=longest + 3)\n",
    "    predictions[id] = answer\n",
    "    questions[id] = question"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save precisions\n",
    "precisions = {}\n",
    "for i in range(1, n_drafts+1):\n",
    "    prec = str(i)\n",
    "    responses = {k: v[:i] for k, v in predictions.items()}\n",
    "    precisions[prec] = responses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print some results\n",
    "counter = 0\n",
    "for k in predictions:\n",
    "    if counter >= 10:\n",
    "        break\n",
    "    print(questions[k])\n",
    "    print(predictions[k])\n",
    "    counter += 1\n",
    "    print(\"================\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save results\n",
    "os.makedirs(\"../../trivia/\", exist_ok=True)\n",
    "for prec in range(1, n_drafts+1):\n",
    "    out_path = f\"../nucleus_extra/trivia_extra/ngram_4trivia_{model_type}_{prec}_4.json\"\n",
    "    with open(out_path, \"w\") as f:\n",
    "        json.dump(precisions[str(prec)], f, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}