extract_lifespan_data_rag_text.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "46ebdfed-13ee-4464-91db-bc1e7f4b6ca2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "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, TextLoader\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 text file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc5a926e-e205-4863-9c8a-acce0debe720",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_file = \"data/study_1.txt\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fa5818f8-b63f-4f3b-91ce-1ee37b8bee56",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Step 1: Load your raw text document using TextLoader\n",
    "loader = TextLoader(data_file)\n",
    "loaded_doc = loader.load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "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": 8,
   "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": 22,
   "id": "5655e72b-1f47-4066-b8cb-f1024a1faf27",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1392\n"
     ]
    }
   ],
   "source": [
    "n_chunks = 6\n",
    "\n",
    "query = \"\"\"\n",
    "Extract all the modifications / treatments vs. no modification / no treatment and point out the change in median and maximum lifespan for each modification / treatment.\n",
    "\"\"\"\n",
    "docs = db.similarity_search(query, k=n_chunks)\n",
    "\n",
    "prompt_context = f\"{docs[0].page_content}\"\n",
    "\n",
    "print(len(prompt_context))\n",
    "\n",
    "for i in range(1, n_chunks):\n",
    "    prompt_context = prompt_context + f\"\\n\\n{docs[i].page_content}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "f91aa8f2-0d21-41d2-879f-43e2b79e464d",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt_template = r\"\"\"In the following, separated by ```, you will find parts of a study on longevity. \n",
    "Extract all the modifications / treatments vs. no modification / no treatment / control group and point out the change in median and maximum lifespan for \n",
    "each modification / treatment. \n",
    "If you don't know the answer, then say that you don't know.\n",
    "\n",
    "```\n",
    "{prompt_context}\n",
    "```\n",
    "\"\"\".format(prompt_context=prompt_context)\n",
    "\n",
    "\n",
    "formatting = \"\"\"\n",
    "\n",
    "Return the result in JSON format without any further comments. Structure the JSON in the following way: \n",
    "{\n",
    "    treatment_1: { \n",
    "        median_lifespan_no_treament: x_median, \n",
    "        median_lifespan_treament: y_median, \n",
    "        maximum_lifespan_no_treatment: x_max, \n",
    "        maximum_lifespan_treatment: y_max \n",
    "    }, \n",
    "    treatment_2: { \n",
    "        median_lifespan_no_treament: x_median, \n",
    "        median_lifespan_treament: y_median, \n",
    "        maximum_lifespan_no_treatment: x_max, \n",
    "        maximum_lifespan_treatment: y_max \n",
    "    } \n",
    "}\n",
    "Replace treatment_1, treatment_2, etc. by the real treatment or intervention.\"\"\"\n",
    "\n",
    "\n",
    "prompt_template = prompt_template + formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "aa74222c-5ab5-43ff-bae9-65fd01bf69f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```json\n",
      "{\n",
      "    \"ENU-treated Apcdel/+\": {\n",
      "        \"median_lifespan_no_treatment\": 255,\n",
      "        \"median_lifespan_treatment\": 117,\n",
      "        \"maximum_lifespan_no_treatment\": \"unknown\",\n",
      "        \"maximum_lifespan_treatment\": \"unknown\"\n",
      "    },\n",
      "    \"ENU-treated Egr1+/−, Apc del/+\": {\n",
      "        \"median_lifespan_no_treatment\": 179,\n",
      "        \"median_lifespan_treatment\": 114,\n",
      "        \"maximum_lifespan_no_treatment\": \"unknown\",\n",
      "        \"maximum_lifespan_treatment\": \"unknown\"\n",
      "    },\n",
      "    \"ENU-treated Tp53+/−, Apc del/+\": {\n",
      "        \"median_lifespan_no_treatment\": 144,\n",
      "        \"median_lifespan_treatment\": 101,\n",
      "        \"maximum_lifespan_no_treatment\": \"unknown\",\n",
      "        \"maximum_lifespan_treatment\": \"unknown\"\n",
      "    },\n",
      "    \"ENU-treated triple heterozygous (Egr1+/−, Tp53+/−, Apc del/+)\": {\n",
      "        \"median_lifespan_no_treatment\": 178,\n",
      "        \"median_lifespan_treatment\": 97,\n",
      "        \"maximum_lifespan_no_treatment\": \"unknown\",\n",
      "        \"maximum_lifespan_treatment\": \"unknown\"\n",
      "    },\n",
      "    \"ENU-treated Apcdel/+ recipients\": {\n",
      "        \"median_lifespan_no_treatment\": 101,\n",
      "        \"median_lifespan_treatment\": 84,\n",
      "        \"maximum_lifespan_no_treatment\": \"unknown\",\n",
      "        \"maximum_lifespan_treatment\": \"unknown\"\n",
      "    },\n",
      "    \"ENU-treated Tp53+/− donor BM cells\": {\n",
      "        \"median_lifespan_no_treatment\": 141,\n",
      "        \"median_lifespan_treatment\": 112,\n",
      "        \"maximum_lifespan_no_treatment\": \"unknown\",\n",
      "        \"maximum_lifespan_treatment\": \"unknown\"\n",
      "    }\n",
      "}\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "get_completion(prompt_template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed24e6a",
   "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
}