entity_linker_dspy_rag.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a20dac2-faf0-41f5-8524-44eaf008b630",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import json\n",
    "import requests\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = 'openai-key'\n",
    "\n",
    "client = OpenAI()\n",
    "\n",
    "FIELDS = \"entities\"\n",
    "HOST = \"nl.diffbot.com\"\n",
    "TOKEN = \"diffbot-token\" \n",
    "\n",
    "import requests\n",
    "\n",
    "def get_request(payload):\n",
    "    res = requests.post(\"https://{}/v1/?fields={}&token={}\".format(HOST, FIELDS, TOKEN), json=payload)\n",
    "    ret = None\n",
    "    try:\n",
    "        ret = res.json()\n",
    "    except:\n",
    "        print(\"Bad response: \" + res.text)\n",
    "        print(res.status_code)\n",
    "        print(res.headers)\n",
    "    return ret\n",
    "\n",
    "def get_entity_types_diffbot(res):\n",
    "    entity_types = []\n",
    "    if res['entities'] == []:\n",
    "        entity_types = \"No valid entity type found from the answer.\"\n",
    "    else:\n",
    "        answer_name = res['entities'][0]['name']\n",
    "        answer_types = res['entities'][0]['allTypes']\n",
    "        for i in range(len(answer_types)): \n",
    "            entity_type = answer_types[i]['name']\n",
    "            entity_types.append(entity_type)\n",
    "        \n",
    "    return entity_types\n",
    "\n",
    "def type_evaluator_by_llm(question):\n",
    "    response = client.chat.completions.create(\n",
    "        model=\"gpt-4-turbo\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": \"You are a helpful assistant who determines the correct entity type that the answer of the question should belongs to. Possible entity types are organization, person, product, country, city, article. And if the question is a yes-or-no type of question, answer: yes/no type. Yes/No type questions start with: Is, Was, Do, Did, etc.\"},\n",
    "            {\"role\": \"user\", \"content\": question},\n",
    "        ]\n",
    "    )\n",
    "    answer = response.choices[0].message.content\n",
    "    question_type = answer\n",
    "    return question_type\n",
    "\n",
    "def entity_linker(answer_vector, question):\n",
    "    if answer_vector.lower() == \"none\":\n",
    "        return \"The entity type of the answer should be: None\", \"Entity type(s) of the answer: None\", \"entity type matched!\"\n",
    "    \n",
    "    # fact-check validity of vector-based answer via Diffbot Natural Language API\n",
    "    type_to_be_checked = answer_vector\n",
    "    content = \"Please verify the type of: \" + type_to_be_checked\n",
    "    \n",
    "    res = get_request({\n",
    "        \"content\": content,\n",
    "        \"lang\": \"en\",\n",
    "        \"format\": \"plain text with title\"\n",
    "    })\n",
    "    \n",
    "    entity_list_diffbot = get_entity_types_diffbot(res)\n",
    "    question_type = type_evaluator_by_llm(question)\n",
    "    correct_question_type = \"The entity type of the answer should be: \" + type_evaluator_by_llm(question)\n",
    "    \n",
    "    if entity_list_diffbot != [] and \"yes\" not in question_type.lower():\n",
    "        type_matches = any(question_type.lower() in entity.lower() for entity in entity_list_diffbot)\n",
    "        original_answer_type = \"Entity type(s) of the answer: \"\n",
    "        if isinstance(entity_list_diffbot, list):\n",
    "            original_answer_type += ', '.join(entity_list_diffbot)\n",
    "        else:\n",
    "            original_answer_type += entity_list_diffbot\n",
    "            \n",
    "        if type_matches:\n",
    "            type_status = \"entity type matched!\"\n",
    "        else:\n",
    "            type_status = \"entity type not matched!\"\n",
    "\n",
    "    elif \"yes\" or \"no\" in question_type.lower():\n",
    "        original_answer_type = \"Entity type(s) of the answer: \" + \"yes/no type\"\n",
    "        type_status = \"entity type not required for the answer!\"\n",
    "        \n",
    "    else:\n",
    "        original_answer_type = \"Entity type(s) of the answer: invalid.\"\n",
    "        type_status = \"entity type not matched!\"\n",
    "        \n",
    "    return correct_question_type, original_answer_type, type_status\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 241,
   "id": "c28b2a3a-21f6-4a5b-922e-021b279c37ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: person',\n",
       " 'Entity type(s) of the answer: person',\n",
       " 'entity type matched!')"
      ]
     },
     "execution_count": 241,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 1\n",
    "entity_linker(answer_vector = \"Sam Altman, Greg Brockman, Ilya Sutskever, Andrej Karpathy, Wojciech Zaremba\", question = \"Who are the other founders Elon Musk co-found OpenAI with?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 243,
   "id": "3e3e993f-cc6b-4f2b-b070-cb1238533274",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: None',\n",
       " 'Entity type(s) of the answer: None',\n",
       " 'entity type matched!')"
      ]
     },
     "execution_count": 243,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 2\n",
    "entity_linker(answer_vector = \"None\", question = \"Who are the other founders Elon Musk co-found SpaceX with?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 248,
   "id": "a6b3c3aa-89f7-4a39-a3db-87037f40c126",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: yes/no type',\n",
       " 'Entity type(s) of the answer: yes/no type',\n",
       " 'entity type not required for the answer!')"
      ]
     },
     "execution_count": 248,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 3\n",
    "entity_linker(answer_vector = \"No\", question = \"Did Elon co-found The Boring Company with other founders?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 249,
   "id": "66f59ee1-6449-4741-bed9-a47010ea82c7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: yes/no type',\n",
       " 'Entity type(s) of the answer: yes/no type',\n",
       " 'entity type not required for the answer!')"
      ]
     },
     "execution_count": 249,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 4\n",
    "entity_linker(answer_vector = \"No\", question = \"Is The Boring Company one of the companies Elon co-founded with other founders?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 274,
   "id": "7ff84b7e-9ede-4b7d-943a-0a568c3147e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: organization',\n",
       " 'Entity type(s) of the answer: No valid entity type found from the answer.',\n",
       " 'entity type not matched!')"
      ]
     },
     "execution_count": 274,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 5\n",
    "entity_linker(answer_vector = \"sdalfkfjgfldkger\", question = \"List out the companies that Elon Musk co-founded with other founders.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 271,
   "id": "b29aeb74-e799-44ff-88f2-32178f519b3a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: organization',\n",
       " 'Entity type(s) of the answer: organization',\n",
       " 'entity type matched!')"
      ]
     },
     "execution_count": 271,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 6\n",
    "entity_linker(answer_vector = \"The Boring Company\", question = \"List out one of the companies that Elon Musk founded himself.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 276,
   "id": "1bdc89bf-fcaa-4c00-bbc4-7616193170df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: organization',\n",
       " 'Entity type(s) of the answer: No valid entity type found from the answer.',\n",
       " 'entity type not matched!')"
      ]
     },
     "execution_count": 276,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 7\n",
    "entity_linker(answer_vector = \"PetrolCar\", question = \"List out one of the companies that Elon Musk founded himself.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 277,
   "id": "9c36de64-97c0-45e4-a3c5-8c9d33ad8978",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: organization',\n",
       " 'Entity type(s) of the answer: organization',\n",
       " 'entity type matched!')"
      ]
     },
     "execution_count": 277,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 8\n",
    "entity_linker(answer_vector = \"SpaceX\", question = \"List out one of the companies that Elon Musk founded himself.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 279,
   "id": "2725da91-dc19-44f3-beac-0b88cee48831",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('The entity type of the answer should be: organization',\n",
       " 'Entity type(s) of the answer: No valid entity type found from the answer.',\n",
       " 'entity type not matched!')"
      ]
     },
     "execution_count": 279,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test 9\n",
    "entity_linker(answer_vector = \"EarthY\", question = \"List out one of the companies that Elon Musk founded himself.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18af48fe-7f02-4bb3-bfa0-c2a095e083b2",
   "metadata": {},
   "source": [
    "# Update Vanilla DSPy RAG with Entity Linker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "58cfa28d-9ad1-4a9d-95ef-a516a552c27e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GenerateAnswer(dspy.Signature):\n",
    "    \"\"\"Answer questions with short factoid answers.\"\"\"\n",
    "    context = dspy.InputField(desc=\"may contain relevant facts\")\n",
    "    question = dspy.InputField()\n",
    "    answer = dspy.OutputField(desc=\"often between 1 and 5 words\")\n",
    "\n",
    "class AnswerTypeValidityCheck(dspy.Signature):\n",
    "    \"\"\"Check if the generated answer has the valid type(s) and valid information.\"\"\"\n",
    "    entity_type = dspy.InputField(desc=\"check valid entity type for the answer\")\n",
    "    question = dspy.InputField()\n",
    "    answer = dspy.InputField()\n",
    "    type_status = dspy.OutputField(desc=\"entity type matched or not\")\n",
    "\n",
    "class vanilla_dspy_rag(dspy.Module):\n",
    "    \n",
    "    # we set num_passages=1 to avoid the same passage being repeatedly retrieved multiple times\n",
    "    def __init__(self, num_passages=3):\n",
    "        super().__init__()\n",
    "        self.retrieve = dspy.Retrieve(k=num_passages)\n",
    "        self.generate_answer = dspy.ChainOfThought(GenerateAnswer)\n",
    "        self.check_answer_type = dspy.ChainOfThought(AnswerTypeValidityCheck)\n",
    "    \n",
    "    def forward(self, question):\n",
    "        # Step 1: Retrieve context based on the question\n",
    "        context = self.retrieve(question).passages\n",
    "        \n",
    "        # Step 2: Generate an answer based on the context and question\n",
    "        prediction = self.generate_answer(context=context, question=question)\n",
    "        answer = prediction.answer\n",
    "        \n",
    "        # Step 3: Validate the answer type using the entity_linker function\n",
    "        correct_question_type, original_answer_type, type_status = entity_linker(answer, question)\n",
    "        \n",
    "        # Optional: You can use the AnswerTypeValidityCheck signature for validation if needed\n",
    "        validation = self.check_answer_type(entity_type=correct_question_type, question=question, answer=answer).type_status\n",
    "        \n",
    "        return dspy.Prediction(context=context, answer=answer, type_status=type_status)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "361dec8e-a8e7-4a84-93d9-af3d636d459f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Prediction(\n",
       "    context=[\"== Notable members ==\\nMaye Musk (née Haldeman; born 1948), model and dietitian. She has appeared on the cover of several magazines, including a Time magazine health edition, Women's Day, and Vogue; Elon's mother.\\nElon Musk (born 1971), entrepreneur and business magnate. Variously CEO, CTO, and/or Chairman of SpaceX, Tesla, X Corp., and Neuralink. He was Time Magazine's 2021 Person of the Year.\\nKimbal Musk (born 1972), entrepreneur, philanthropist, and restaurateur. He founded Zip2 in 1998 with his brother Elon and later sold it to Compaq for $307 million. He is the co-founder and chairman of Big Green.\\nTosca Musk (born 1974), filmmaker and sister of Elon; she is the co-founder of Passionflix, an OTT entertainment streaming platform and production company.\\nLyndon Rive (born 1977), businessman who co-founded SolarCity and served as its CEO until 2017. He is Elon's cousin through his mother Kaye Rive, Maye Musk's twin.\\n\\n\\n== References ==\", \"== Notable members ==\\nMaye Musk (née Haldeman; born 1948), model and dietitian. She has appeared on the cover of several magazines, including a Time magazine health edition, Women's Day, and Vogue; Elon's mother.\\nElon Musk (born 1971), entrepreneur and business magnate. Variously CEO, CTO, and/or Chairman of SpaceX, Tesla, X Corp., and Neuralink. He was Time Magazine's 2021 Person of the Year.\\nKimbal Musk (born 1972), entrepreneur, philanthropist, and restaurateur. He founded Zip2 in 1998 with his brother Elon and later sold it to Compaq for $307 million. He is the co-founder and chairman of Big Green.\\nTosca Musk (born 1974), filmmaker and sister of Elon; she is the co-founder of Passionflix, an OTT entertainment streaming platform and production company.\\nLyndon Rive (born 1977), businessman who co-founded SolarCity and served as its CEO until 2017. He is Elon's cousin through his mother Kaye Rive, Maye Musk's twin.\\n\\n\\n== References ==\", \"== Notable members ==\\nMaye Musk (née Haldeman; born 1948), model and dietitian. She has appeared on the cover of several magazines, including a Time magazine health edition, Women's Day, and Vogue; Elon's mother.\\nElon Musk (born 1971), entrepreneur and business magnate. Variously CEO, CTO, and/or Chairman of SpaceX, Tesla, X Corp., and Neuralink. He was Time Magazine's 2021 Person of the Year.\\nKimbal Musk (born 1972), entrepreneur, philanthropist, and restaurateur. He founded Zip2 in 1998 with his brother Elon and later sold it to Compaq for $307 million. He is the co-founder and chairman of Big Green.\\nTosca Musk (born 1974), filmmaker and sister of Elon; she is the co-founder of Passionflix, an OTT entertainment streaming platform and production company.\\nLyndon Rive (born 1977), businessman who co-founded SolarCity and served as its CEO until 2017. He is Elon's cousin through his mother Kaye Rive, Maye Musk's twin.\\n\\n\\n== References ==\"],\n",
       "    answer='PayPal'\n",
       ")"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# nomic embedding + llama3\n",
    "question_org_elon_cofounded_1b = \"What are the other founders Elon Musk co-found SpaceX with?\"\n",
    "vanilla_dspy_rag_q1b = vanilla_dspy_rag_nomic_llama3.forward(question=question_org_elon_cofounded_1b)\n",
    "vanilla_dspy_rag_q1b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "589d9215-f712-47dc-97b1-d52f9f2ed816",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated Answer: Elon Musk co-founded SpaceX with Jim Cantrell.\n",
      "Entity Type Status: entity type matched!\n"
     ]
    }
   ],
   "source": [
    "import dspy\n",
    "dspy.settings.configure(lm=lm_llama3, rm=nomic_rm)\n",
    "vanilla_dspy_rag_nomic_llama3 = vanilla_dspy_rag()\n",
    "question_org_elon_cofounded_1b = \"What are the other founders Elon Musk co-founded SpaceX with?\"\n",
    "vanilla_dspy_rag_q1b = vanilla_dspy_rag_nomic_llama3.forward(question=question_org_elon_cofounded_1b)\n",
    "\n",
    "print(\"Generated Answer:\", vanilla_dspy_rag_q1b.answer)\n",
    "print(\"Entity Type Status:\", vanilla_dspy_rag_q1b.type_status)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "310b9035-4b96-4482-b480-2b27449385b8",
   "metadata": {},
   "source": [
    "With entity linker, we finally don't get the answer 'Paypal' (type: organization) anymore, but we're still seeing hallucination as Elon Musk is the sole founder of SpaceX, so we need to further bring in ground truth from the knowledge graph and see how well the optimizer can make LLM stick to the ground truth from the knowledge graph"
   ]
  }
 ],
 "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
}