5_best_of_n.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "import torch\n",
    "from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM\n",
    "from transformers import LogitsProcessor\n",
    "from sentence_transformers import CrossEncoder\n",
    "from typing import Iterable\n",
    "import envs\n",
    "import pandas as pd\n",
    "import string\n",
    "import numpy\n",
    "from functools import partial\n",
    "from leaderboard import SummaryGenerator, EvaluationModel, run_eval\n",
    "\n",
    "MODEL_NAME = \"mistralai/Mistral-7B-Instruct-v0.1\"\n",
    "# MODEL_NAME = \"mistral_dpo_5k\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-17 14:23:02,331 - INFO - We will use 90% of the memory on device 0 for storing the model, and 10% for the buffer to avoid OOM. You can set `max_memory` in to a higher value to use more memory (at your own risk).\n",
      "Loading checkpoint shards: 100%|███████████████████████████████████████████████████| 2/2 [00:07<00:00,  3.71s/it]\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n",
    "model = AutoModelForCausalLM.from_pretrained(MODEL_NAME,\n",
    "                                             device_map=\"auto\",\n",
    "                                             torch_dtype=\"auto\",\n",
    "                                             attn_implementation=\"flash_attention_2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.chat_template = open(\"mistral_template.jinja\", \"r\").read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/rogger/miniconda3/lib/python3.11/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n",
      "/home/rogger/miniconda3/lib/python3.11/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n",
      "2024-06-17 14:23:18,541 - INFO - Use pytorch device: cuda\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import CrossEncoder\n",
    "scorer = CrossEncoder(\"vectara/hallucination_evaluation_model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gen_func(source, scorer, debug=False):\n",
    "    messages = [{\"role\": \"system\", \"content\": envs.SYSTEM_PROMPT},\n",
    "                {\"role\": \"user\", \"content\": envs.USER_PROMPT.format(passage=source)}]\n",
    "    # good_word_ids = tokenizer.encode(source + \"\\n\" + \"\\n\".join(string.printable) + \"\\n\" + tokenizer.eos_token)\n",
    "    # white_list_processor = WhiteListLogitsProcessor(good_word_ids)\n",
    "    input_ids = tokenizer.apply_chat_template(messages, \n",
    "                                              add_generation_prompt=True, \n",
    "                                              return_tensors=\"pt\").to(\"cuda\")\n",
    "    out = model.generate(input_ids, \n",
    "                         do_sample=True, \n",
    "                         temperature=0.7,\n",
    "                         top_p=0.95,\n",
    "                         max_new_tokens=512,\n",
    "                         num_return_sequences=10,\n",
    "                         pad_token_id=tokenizer.eos_token_id)\n",
    "    texts = tokenizer.batch_decode(out[:, len(input_ids[0]):], skip_special_tokens=True)\n",
    "    # Best of N\n",
    "    scores = scorer.predict([[source, summ] for summ in texts], show_progress_bar=False)\n",
    "    if debug:\n",
    "        print(tokenizer.decode(input_ids[0]))\n",
    "        print(scores)\n",
    "    idx = numpy.argmax(scores)\n",
    "    return texts[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<s> [INST] You are a chat bot answering questions using data. You must stick to the answers provided solely by the text in the passage provided. \n",
      "\n",
      "You are asked the question 'Provide a concise summary of the following passage, covering the core pieces of information described': \n",
      "Passage:\n",
      "The first vaccine for Ebola was approved by the FDA in 2019 in the US, five years after the initial outbreak in 2014. To produce the vaccine, scientists had to sequence the DNA of Ebola, then identify possible vaccines, and finally show successful clinical trials. Scientists say a vaccine for COVID-19 is unlikely to be ready this year, although clinical trials have already started. [/INST]\n",
      "[0.9990177  0.99417984 0.9993338  0.9990607  0.99931836 0.99933136\n",
      " 0.9990609  0.9993253  0.99923015 0.99903166]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'The first Ebola vaccine was approved by the FDA in the US in 2019, five years after the initial outbreak in 2014. Scientists had to sequence the DNA of Ebola, identify possible vaccines, and show successful clinical trials to produce the vaccine. A vaccine for COVID-19 is unlikely to be ready this year, although clinical trials have already started.'"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gen_func(\"The first vaccine for Ebola was approved by the FDA in 2019 in the US, five years after the initial outbreak in 2014. To produce the vaccine, scientists had to sequence the DNA of Ebola, then identify possible vaccines, and finally show successful clinical trials. Scientists say a vaccine for COVID-19 is unlikely to be ready this year, although clinical trials have already started.\", \n",
    "         scorer=scorer,\n",
    "         debug=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "summ = SummaryGenerator()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 39%|███████████████████████████▊                                           | 394/1006 [43:08<3:18:45, 19.49s/it]"
     ]
    }
   ],
   "source": [
    "df = summ.generate_summaries(pd.read_csv(\"leaderboard_dataset.csv\"), partial(gen_func, scorer=scorer))\n",
    "df.to_csv(\"generated.csv\", index=False)"
   ]
  }
 ],
 "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
}