run_mello.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8dd6e7ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "\n",
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e84644ce",
   "metadata": {},
   "source": [
    "#### Set up OpenAI API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1588cf93",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "\n",
    "openai.api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "\n",
    "def call_gpt(cur_prompt, stop):\n",
    "    ans = openai.Completion.create(\n",
    "                model=\"text-davinci-003\",\n",
    "                max_tokens=256,\n",
    "                stop=stop,\n",
    "                prompt=cur_prompt,\n",
    "                temperature=0)\n",
    "    returned = ans['choices'][0]['text']\n",
    "    return returned"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "964a79bd",
   "metadata": {},
   "source": [
    "#### Functions for retrieval models (Contriever)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7cd5d62b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mean_pooling(token_embeddings, mask):\n",
    "    token_embeddings = token_embeddings.masked_fill(~mask[..., None].bool(), 0.)\n",
    "    sentence_embeddings = token_embeddings.sum(dim=1) / mask.sum(dim=1)[..., None]\n",
    "    return sentence_embeddings\n",
    "\n",
    "def get_sent_embeddings(sents, contriever, tok, BSZ=32):    \n",
    "    all_embs = []\n",
    "    for i in tqdm(range(0, len(sents), BSZ)):\n",
    "        sent_batch = sents[i:i+BSZ]\n",
    "        inputs = tok(sent_batch, padding=True, truncation=True, return_tensors='pt').to(\"cuda\")\n",
    "        with torch.no_grad():\n",
    "            outputs = contriever(**inputs)\n",
    "            embeddings = mean_pooling(outputs[0], inputs['attention_mask'])\n",
    "        all_embs.append(embeddings.cpu())\n",
    "    all_embs = torch.vstack(all_embs)\n",
    "    return all_embs\n",
    "\n",
    "def retrieve_facts(query, fact_embs, contriever, tok, k=1):\n",
    "    inputs = tok([query], padding=True, truncation=True, return_tensors='pt').to(\"cuda\")\n",
    "    with torch.no_grad():\n",
    "        outputs = contriever(**inputs)\n",
    "        query_emb = mean_pooling(outputs[0], inputs['attention_mask']).cpu()\n",
    "    sim = (query_emb @ fact_embs.T)[0]\n",
    "    knn = sim.topk(k, largest=True)\n",
    "    return knn.indices"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "771712cb",
   "metadata": {},
   "source": [
    "#### Load dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "469cccc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('datasets/MQuAKE-CF-3k.json', 'r') as f:\n",
    "    dataset = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ded19987",
   "metadata": {},
   "source": [
    "#### Build a memory index which contains all the edits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f098daf5",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_facts = set()\n",
    "for d in dataset:\n",
    "    for r in d[\"requested_rewrite\"]:\n",
    "        new_facts.add(f'{r[\"prompt\"].format(r[\"subject\"])} {r[\"target_new\"][\"str\"]}')\n",
    "new_facts = list(new_facts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7a060c5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "contriever = AutoModel.from_pretrained(\"facebook/contriever-msmarco\").cuda()\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"facebook/contriever-msmarco\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f4eaa2a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 88/88 [00:01<00:00, 46.75it/s]\n"
     ]
    }
   ],
   "source": [
    "embs = get_sent_embeddings(new_facts, contriever, tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "24d605e8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The name of the current head of state in United States of America is Norodom Sihamoni\n"
     ]
    }
   ],
   "source": [
    "# Run test for retrieval index\n",
    "fact_ids = retrieve_facts(\"Who is the president of the US?\", embs, contriever, tokenizer)\n",
    "print(new_facts[fact_ids[0]])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0937cc8",
   "metadata": {},
   "source": [
    "#### Run MeLLo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "39b75fa6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# read prompts\n",
    "with open('prompts/MeLLo-prompt.txt', 'r') as f:\n",
    "    task_prompt = f.read()\n",
    "stop = [\"Retrieved fact:\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "ffdefa49",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [02:18<00:00, 13.83s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Multi-hop acc = 0.6 (6 / 10)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Run MeLLo on the first T (T=10) examples\n",
    "T = 10\n",
    "\n",
    "cor = 0\n",
    "tot = 0\n",
    "\n",
    "for d in tqdm(dataset[:10]):\n",
    "    tot += 1\n",
    "    for q in d[\"questions\"]:\n",
    "        found_ans = False\n",
    "        prompt = task_prompt + \"\\n\\nQustion: \" + q\n",
    "        for i in range(4):\n",
    "            # prompt the model to generate a subquestion and a tentative answer\n",
    "            gen = call_gpt(prompt, stop)\n",
    "            last_sent = gen.strip().split('\\n')[-1]\n",
    "            \n",
    "            # if final answer is there, get the answer and exit\n",
    "            if last_sent.startswith('Final answer: '):\n",
    "                found_ans = True\n",
    "                ans = last_sent[len(\"Final answer: \"):]\n",
    "                break\n",
    "            \n",
    "            # otherwise, extract the generated subquestion\n",
    "            if len(gen.strip().split('\\n')) < 2:\n",
    "                break # failed case\n",
    "            subquestion = gen.strip().split('\\n')[-2]\n",
    "            if not subquestion.startswith('Subquestion: '):\n",
    "                break # failed case\n",
    "            subquestion = subquestion[len(\"Subquestion: \"):]\n",
    "            \n",
    "            # retrieve an edited fact using the generated subquestion\n",
    "            fact_ids = retrieve_facts(subquestion, embs, contriever, tokenizer)\n",
    "            fact_sent = new_facts[fact_ids[0]]\n",
    "            \n",
    "            # put the retrieved fact at the end of the prompt, the model self-checks if it contradicts\n",
    "            prompt = prompt + gen + 'Retrieved fact: ' + fact_sent + '.'\n",
    "            \n",
    "        prompt = prompt + gen\n",
    "        \n",
    "        if not found_ans:\n",
    "            continue\n",
    "        # if the answer is correct\n",
    "        if ans == d[\"new_answer\"] or ans in d[\"new_answer_alias\"]:\n",
    "            cor += 1\n",
    "            break\n",
    "\n",
    "print(f'Multi-hop acc = {cor / tot} ({cor} / {tot})')"
   ]
  }
 ],
 "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.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}