generate_embeddings.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "from neo4j import Query, GraphDatabase, RoutingControl, Result \n",
    "from graphdatascience import GraphDataScience\n",
    "from langchain_openai import OpenAIEmbeddings, ChatOpenAI\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.output_parsers import StrOutputParser"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Change this\n",
    "os.environ['OPENAI_API_KEY'] = \"\" \n",
    "DB_ULR = \"neo4j://localhost:7687\"\n",
    "DB_USER = \"neo4j\"\n",
    "DB_PASS = \"test1234\"\n",
    "DB_NAME = \"sdtm\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "driver = GraphDatabase.driver(DB_ULR, auth=(DB_USER, DB_PASS))\n",
    "driver.verify_connectivity()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds = GraphDataScience(driver)\n",
    "gds.set_database(DB_NAME)\n",
    "gds.version()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Utility\n",
    "def split_dataframe(df, chunk_size = 5000): \n",
    "    chunks = list()\n",
    "    num_chunks = len(df) // chunk_size + 1\n",
    "    for i in range(num_chunks):\n",
    "        chunks.append(df[i*chunk_size:(i+1)*chunk_size])\n",
    "    return chunks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate embeddings\n",
    "![image](./images/generate_embeddings.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pharmacological_group = gds.run_cypher('''\n",
    "    match (n:PharmacologicalSubgroup)\n",
    "    return n.id as id, n.name as name \n",
    "''')\n",
    "df_pharmacological_group.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pharmacological_group.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = OpenAIEmbeddings(\n",
    "    deployment=\"your-embeddings-deployment-name\",\n",
    "    model='text-embedding-3-small', # 1536\n",
    "    #openai_api_base=\"https://your-endpoint.openai.azure.com/\",\n",
    "    #openai_api_type=\"azure\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pharmacological_group['embedding'] = df_pharmacological_group['name'].apply( lambda name: embeddings.embed_documents([name])[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pharmacological_group.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Store embeddings and index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pharmacological_group = gds.run_cypher('''\n",
    "    unwind $data as row\n",
    "    match (n:PharmacologicalSubgroup{id: row.id})\n",
    "    set n.embedding = row.embedding                                                                             \n",
    "''',\n",
    "params = { 'data': df_pharmacological_group.to_dict(orient='records') })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dimension=1536\n",
    "index_name = 'pharmacological'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher('''\n",
    "    CREATE VECTOR INDEX $index_name if not exists\n",
    "    for (n:PharmacologicalSubgroup) on (n.embedding)\n",
    "    OPTIONS {\n",
    "        indexConfig: {\n",
    "            `vector.dimensions`: $dimension,\n",
    "            `vector.similarity_function`: 'cosine'       \n",
    "        }\n",
    "    }\n",
    "''',\n",
    "    params = { 'index_name': index_name, 'dimension': dimension }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify that index is online and populated \n",
    "gds.run_cypher('''\n",
    "    show index yield name, state, populationPercent, type, labelsOrTypes, properties\n",
    "    where name = $index_name\n",
    "    return *\n",
    "''',\n",
    "    params = { 'index_name': index_name, 'dimension': dimension }\n",
    ").head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Vector search\n",
    "\n",
    "![image](./images/vector_search.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = 'What drugs are used to treat type 2 diabetes?'\n",
    "q_vector = embeddings.embed_query(text=question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher(''' \n",
    "    call db.index.vector.queryNodes($index_name, 5, $vector) yield node, score\n",
    "    return node.id as id, node.name as name, score\n",
    "''',\n",
    "    params = {'index_name': index_name, 'vector': q_vector}\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_group_and_drugs = gds.run_cypher(''' \n",
    "    call db.index.vector.queryNodes($index_name, 5, $vector) yield node, score\n",
    "    match (node)-[:DrugClassHierarchy*1..10]->(d:Drug)\n",
    "    return node.id as id, node.name as name, score, collect(d.name) as drugs\n",
    "''',\n",
    "    params = {'index_name': index_name, 'vector': q_vector}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_group_and_drugs.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Agmented Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", temperature=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "promt = ChatPromptTemplate.from_template('''\n",
    "    Answer the question based only on the following context: {data}\n",
    "                                         \n",
    "    Question: {question}\n",
    "''')\n",
    "chain = promt | llm | StrOutputParser()\n",
    "message =chain.invoke({ 'question': question, 'data': df_group_and_drugs[['name','drugs']].to_dict(orient='records') })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "notebooks",
   "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": 2
}