notebooks/demo_notebook_for_yes_score.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "604a02ae-c353-4405-8aef-3a91d59cdca9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import transformers\n",
    "from tqdm import tqdm\n",
    "from datasets import load_dataset\n",
    "from transformers import LlamaForCausalLM, LlamaTokenizer, pipeline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06e2a148-df55-454b-bc4c-cb12bb2841ab",
   "metadata": {},
   "source": [
    "# load models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c324434-e38e-44a2-8b5b-636b1c8e39ed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "score_model_name = \"meta-llama/Llama-2-13b-chat-hf\"\n",
    "qa_model_name = \"deepset/roberta-base-squad2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f3b0c99-eb93-4271-b0a5-90a9c709695e",
   "metadata": {},
   "outputs": [],
   "source": [
    "score_model = LlamaForCausalLM.from_pretrained(score_model_name, device_map=device)\n",
    "score_tokenizer = LlamaTokenizer.from_pretrained(score_model_name, device_map=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23285cae-c2fc-41db-8e66-bb66b7f02c9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_model = AutoModelForQuestionAnswering.from_pretrained(qa_model_name)\n",
    "qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed8f43c1-1da0-4fcc-be60-117f4a88eb3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3c8a483-596b-4989-9f11-b96bbd3dc124",
   "metadata": {},
   "source": [
    "# load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f784444-401b-49f2-9f36-5c1ea261b82f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "squad = load_dataset(\"rajpurkar/squad_v2\", split=\"train[:5000]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f715dc-7870-4be5-9d5d-9b28517889e9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "squad = squad.train_test_split(test_size=0.2)\n",
    "data = squad[\"train\"]\n",
    "print(len(data))\n",
    "data[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd229691-b179-418c-b129-94138f97210c",
   "metadata": {},
   "source": [
    "# define methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f48d7f01-f04c-4b3d-9848-a5e8600d9187",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def make_yes_no_prompt(context: str, question: str, response: str) -> str:\n",
    "    return f\"\"\"Context: {context}\n",
    "\n",
    "Question: {question}\n",
    "\n",
    "Response: {response}\n",
    "\n",
    "Based on the given Context and Question, answer this question:\n",
    "\n",
    "Is the provided Response correct? Answer only Yes or No.\n",
    "\n",
    "Answer:\n",
    "    \"\"\"\n",
    "\n",
    "def yes_score_calculation(outputs, input_length, tokenizer):\n",
    "    generated_tokens = outputs.sequences[:, input_length:]\n",
    "\n",
    "    # 1. find the index (idx) of the first character-based token.\n",
    "    for idx, tok in enumerate(generated_tokens[0]):\n",
    "        next_token_str = tokenizer.decode(tok, skip_special_tokens=True)\n",
    "        n_letters = sum(c.isalpha() for c in next_token_str)\n",
    "        if n_letters != len(next_token_str):\n",
    "            continue\n",
    "        break\n",
    "    \n",
    "    # 2a. do preselection on high probabilities (out of 32k tokens)\n",
    "    probs_all = torch.nn.functional.softmax(outputs.logits[idx][0], dim=-1)\n",
    "    indices = torch.argwhere(probs_all > 0.001)\n",
    "    indices = indices[:, -1]\n",
    "    tokens_max = tokenizer.batch_decode(indices, skip_special_tokens=True)\n",
    "    probs_max = probs_all[probs_all > 0.001]\n",
    "    \n",
    "    # 2b. find yes/no probabilities\n",
    "    next_token_dict = {str(t): p for t, p in zip(tokens_max, probs_max)}\n",
    "    yes_prob = next_token_dict.get(\"Yes\", 0.)\n",
    "    no_prob = next_token_dict.get(\"No\", 0.)\n",
    "    \n",
    "    # 3. calculate and return yes/no confidence score\n",
    "    yes_score = yes_prob / (yes_prob + no_prob) if yes_prob != 0 or no_prob != 0 else 0.5\n",
    "    return yes_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e371c9a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_histogram(scores, title):\n",
    "    plt.hist(scores, range=(0, 1.0), bins=50)\n",
    "    plt.xlabel(\"Yes Score\")\n",
    "    plt.ylabel(\"Number of Questions\")\n",
    "    plt.title(title)\n",
    "    plt.clf()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1b8a8f7-d0b2-4838-8bc5-13b93f162383",
   "metadata": {},
   "source": [
    "# scores for accurate answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f7e43ea-b705-42f7-beb2-0319b02129e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "accurate_scores = []\n",
    "\n",
    "for i, row in tqdm(enumerate(data)):\n",
    "    \n",
    "    if len(row['answers']['text']) < 1:\n",
    "        break\n",
    "    response = row['answers']['text'][0]\n",
    "     \n",
    "    # 1. make the yes/no prompt\n",
    "    prompt = make_yes_no_prompt(row['context'], row['question'], response)\n",
    "    input_ids = score_tokenizer(prompt, return_tensors=\"pt\").input_ids.to(device)\n",
    "    input_length = input_ids.shape[1]\n",
    "\n",
    "    # 2. generate the yes/no answer\n",
    "    #    be sure to generate output with options output_logits=True, \n",
    "    #    and return_dict_in_generate=True\n",
    "    outputs = score_model.generate(input_ids, output_logits=True, return_dict_in_generate=True, max_new_tokens=5)\n",
    "\n",
    "    # 3. calculate the yes-score \n",
    "    yes_score = yes_score_calculation(outputs, input_length, tokenizer)\n",
    "    accurate_scores.append(yes_score)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b016580-f919-4869-8828-b28fe35f8b16",
   "metadata": {},
   "outputs": [],
   "source": [
    "title = \"Histogram of Yes Scores Correct Answers\"\n",
    "\n",
    "plot_histogram(accurate_scores, title)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e605505a-e7d9-4c2f-b23e-456f1e0cadde",
   "metadata": {},
   "source": [
    "# scores for answers given by roberta qa model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f425d72-e7d3-4cd6-a03c-65aa44b9dac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_scores = []\n",
    "\n",
    "for i, row in tqdm(enumerate(data)):\n",
    "\n",
    "    nlp = pipeline('question-answering', model=qa_model_name, tokenizer=qa_model_name)\n",
    "    QA_input = {\n",
    "        'question': row[\"question\"],\n",
    "        'context': row['context']\n",
    "    }\n",
    "    response = nlp(QA_input)\n",
    "\n",
    "    \n",
    "    prompt = make_yes_no_prompt(row['context'], row['question'], response)\n",
    "    input_ids = score_tokenizer(prompt, return_tensors=\"pt\").input_ids.to(device)\n",
    "    input_length = input_ids.shape[1]\n",
    "\n",
    "    outputs = score_model.generate(input_ids, output_logits=True, return_dict_in_generate=True, max_new_tokens=5)\n",
    "\n",
    "    yes_score = yes_score_calculation(outputs, input_length, score_tokenizer)\n",
    "    qa_scores.append(yes_score)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a1d5e28",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}