extract_lifespan_data_rag_pdf.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "46ebdfed-13ee-4464-91db-bc1e7f4b6ca2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "import requests\n",
    "\n",
    "from openai import OpenAI\n",
    "from dotenv import load_dotenv\n",
    "from langchain.document_loaders import PyPDFLoader\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_community.vectorstores import FAISS\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f6e4dcb-dde7-4bd6-a113-52b755c9ce2a",
   "metadata": {},
   "source": [
    "## Load the data from the PDF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "cc5a926e-e205-4863-9c8a-acce0debe720",
   "metadata": {},
   "outputs": [],
   "source": [
    "# data_file = \"./data/test_study_1.pdf\"\n",
    "data_file = \"./data/test_study_3.pdf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "fa5818f8-b63f-4f3b-91ce-1ee37b8bee56",
   "metadata": {},
   "outputs": [],
   "source": [
    "doc = PyPDFLoader(data_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "54b3371f-1f0f-4947-ac75-e7bc5528e709",
   "metadata": {},
   "outputs": [],
   "source": [
    "loaded_doc = doc.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "5fae5b47-7bbb-4ded-9546-2a6d56b92474",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the Text Splitter \n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size = 1500,\n",
    "    chunk_overlap = 800\n",
    ")\n",
    "\n",
    "#Create a split of the document using the text splitter\n",
    "splits = text_splitter.split_documents(loaded_doc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "344cbec2-00c3-40e5-a6f0-b7d45b251a6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = FAISS.from_documents(splits, OpenAIEmbeddings())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fed7676-3b1d-4cb2-a90a-b09fa6d46916",
   "metadata": {},
   "source": [
    "## Setup GPT-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "3db183cc-9b5d-478e-8797-81cad160d0e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = OpenAI()\n",
    "\n",
    "def get_completion(message):\n",
    "\n",
    "    chat_completion = client.chat.completions.create(\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"user\",\n",
    "                \"content\": message\n",
    "            }\n",
    "        ],\n",
    "        model=\"gpt-4o\",\n",
    "        temperature=0\n",
    "    )\n",
    "\n",
    "    print(chat_completion.choices[0].message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f9e1f9f-5bcb-48a2-b86f-9a3139dbdf76",
   "metadata": {},
   "source": [
    "## Search the document"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "5655e72b-1f47-4066-b8cb-f1024a1faf27",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What are median lifespan and maximum lifespan in the treatment group vs the control group?\"\n",
    "docs = db.similarity_search(query, k=2)\n",
    "\n",
    "prompt_context = f\"{docs[0].page_content}\\n\\n{docs[1].page_content}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "f91aa8f2-0d21-41d2-879f-43e2b79e464d",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = \"\"\"In the following you will find a study on study on longevity, separated by ```. \n",
    "Please extract both the median and the maximum lifespan increase or decrease of the treatment group vs. the control group. \n",
    "If you don't know the answer, then say that you don't know.\n",
    "\n",
    "Return the result in json format, where \"median_lifespan_change\" and \"maximum_lifespan_change\" are both keys.\n",
    "\n",
    "```\n",
    "{prompt_context}\n",
    "```\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "aa74222c-5ab5-43ff-bae9-65fd01bf69f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```json\n",
      "{\n",
      "  \"median_lifespan_change\": \"I don't know\",\n",
      "  \"maximum_lifespan_change\": \"I don't know\"\n",
      "}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "get_completion(prompt_template.format(prompt_context=prompt_context))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb7eacc6-c9cc-4d30-a1b5-0987320da314",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}