Add Graph Retriever and Generator code, Add input data, Update requirements

app.py

@@ -1,9 +1,13 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+import os
+from rag import local_retriever, global_retriever
 
 """
 For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
 """
+
+
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 
@@ -11,33 +15,37 @@ def respond(
     message,
     history: list[tuple[str, str]],
     system_message,
-    max_tokens,
-    temperature,
+    search_strategy,
     top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
+    if search_strategy == "Global":
+        return global_retriever(message, 2, "multiple paragraphs")
+    else:
+        messages = [{"role": "system", "content": system_message}]
+
+        for val in history:
+            if val[0]:
+                messages.append({"role": "user", "content": val[0]})
+            if val[1]:
+                messages.append({"role": "assistant", "content": val[1]})
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+        messages.append({"role": "user", "content": message})
 
-    messages.append({"role": "user", "content": message})
+        response = ""
 
-    response = ""
+        for message in client.chat_completion(
+            messages,
+            max_tokens=2048,
+            stream=True,
+            temperature=1.0,
+            top_p=top_p,
+        ):
+            token = message.choices[0].delta.content
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+            response += token
+
+        return response
 
-        response += token
-        yield response
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -45,9 +53,11 @@ For information on how to customize the ChatInterface, peruse the gradio docs: h
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+        gr.Textbox(
+            value="You are a medical assistant Chatbot. For any query that you don't know, you will say 'I don't know'. You will answer with the given information:",
+            label="System message",
+        ),
+        gr.Dropdown(choices=["Local", "Global"], label="Select search strategy"),
         gr.Slider(
             minimum=0.1,
             maximum=1.0,
@@ -60,4 +70,4 @@ demo = gr.ChatInterface(
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

rag.py

@@ -0,0 +1,310 @@
+import os
+from neo4j import GraphDatabase, Result
+import pandas as pd
+import numpy as np
+
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from langchain_community.graphs import Neo4jGraph
+from langchain_community.vectorstores import Neo4jVector
+
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+
+from langchain_huggingface import HuggingFaceEndpoint
+
+from typing import Dict, Any
+from tqdm import tqdm
+
+NEO4J_URI = os.getenv("NEO4J_URI")
+NEO4J_USERNAME = os.getenv("NEO4J_USERNAME")
+NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")
+vector_index = os.getenv("VECTOR_INDEX")
+
+chat_llm = HuggingFaceEndpoint(
+    repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
+    task="text-generation",
+    max_new_tokens=100,
+    do_sample=False,
+)
+
+
+def local_retriever(query: str):
+    topChunks = 3
+    topCommunities = 3
+    topOutsideRels = 10
+    topInsideRels = 10
+    topEntities = 10
+
+    driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
+    try:
+        lc_retrieval_query = """
+        WITH collect(node) as nodes
+        // Entity - Text Unit Mapping
+        WITH
+        collect {
+            UNWIND nodes as n
+            MATCH (n)<-[:HAS_ENTITY]->(c:__Chunk__)
+            WITH c, count(distinct n) as freq
+            RETURN c.text AS chunkText
+            ORDER BY freq DESC
+            LIMIT $topChunks
+        } AS text_mapping,
+        // Entity - Report Mapping
+        collect {
+            UNWIND nodes as n
+            MATCH (n)-[:IN_COMMUNITY]->(c:__Community__)
+            WITH c, c.rank as rank, c.weight AS weight
+            RETURN c.summary 
+            ORDER BY rank, weight DESC
+            LIMIT $topCommunities
+        } AS report_mapping,
+        // Outside Relationships 
+        collect {
+            UNWIND nodes as n
+            MATCH (n)-[r:RELATED]-(m) 
+            WHERE NOT m IN nodes
+            RETURN r.description AS descriptionText
+            ORDER BY r.rank, r.weight DESC 
+            LIMIT $topOutsideRels
+        } as outsideRels,
+        // Inside Relationships 
+        collect {
+            UNWIND nodes as n
+            MATCH (n)-[r:RELATED]-(m) 
+            WHERE m IN nodes
+            RETURN r.description AS descriptionText
+            ORDER BY r.rank, r.weight DESC 
+            LIMIT $topInsideRels
+        } as insideRels,
+        // Entities description
+        collect {
+            UNWIND nodes as n
+            RETURN n.description AS descriptionText
+        } as entities
+        // We don't have covariates or claims here
+        RETURN {Chunks: text_mapping, Reports: report_mapping, 
+            Relationships: outsideRels + insideRels, 
+            Entities: entities} AS text, 1.0 AS score, {} AS metadata
+        """
+
+        embedding_model_name = "nomic-ai/nomic-embed-text-v1"
+        embedding_model_kwargs = {"device": "cpu", "trust_remote_code": True}
+        encode_kwargs = {"normalize_embeddings": True}
+        embedding_model = HuggingFaceBgeEmbeddings(
+            model_name=embedding_model_name,
+            model_kwargs=embedding_model_kwargs,
+            encode_kwargs=encode_kwargs,
+        )
+
+        lc_vector = Neo4jVector.from_existing_index(
+            embedding_model,
+            url=NEO4J_URI,
+            username=NEO4J_USERNAME,
+            password=NEO4J_PASSWORD,
+            index_name=vector_index,
+            retrieval_query=lc_retrieval_query,
+        )
+        docs = lc_vector.similarity_search(
+            query,
+            k=topEntities,
+            params={
+                "topChunks": topChunks,
+                "topCommunities": topCommunities,
+                "topOutsideRels": topOutsideRels,
+                "topInsideRels": topInsideRels,
+            },
+        )
+
+        return docs[0]
+    except Exception as err:
+        return f"Error: {err}"
+    finally:
+        try:
+            driver.close()
+        except Exception as e:
+            print(f"Error closing driver: {e}")
+
+
+def global_retriever(query: str, level: int, response_type: str):
+    MAP_SYSTEM_PROMPT = """
+    ---Role---
+
+    You are a helpful assistant responding to questions about data in the tables provided.
+
+    ---Goal---
+
+    Generate a response consisting of a list of key points that responds to the user's question, summarizing all relevant information in the input data tables.
+
+    You should use the data provided in the data tables below as the primary context for generating the response.
+    If you don't know the answer or if the input data tables do not contain sufficient information to provide an answer, just say so. Do not make anything up.
+
+    Each key point in the response should have the following element:
+    - Description: A comprehensive description of the point.
+    - Importance Score: An integer score between 0-100 that indicates how important the point is in answering the user's question. An 'I don't know' type of response should have a score of 0.
+
+    The response shall preserve the original meaning and use of modal verbs such as "shall", "may" or "will".
+
+    Points supported by data should list the relevant reports as references as follows:
+    "This is an example sentence supported by data references [Data: Reports (report ids)]"
+
+    **Do not list more than 5 record ids in a single reference**. Instead, list the top 5 most relevant record ids and add "+more" to indicate that there are more.
+
+    For example:
+    "Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Reports (2, 7, 64, 46, 34, +more)]. He is also CEO of company X [Data: Reports (1, 3)]"
+
+    where 1, 2, 3, 7, 34, 46, and 64 represent the id (not the index) of the relevant data report in the provided tables.
+
+    Do not include information where the supporting evidence for it is not provided. Always start with {{ and end with }}. 
+
+    The response can only be JSON formatted. Do not add any text before or after the JSON-formatted string in the output.
+
+    The response should adhere to the following format:
+    {{
+        "points": [
+            {{"description": "Description of point 1 [Data: Reports (report ids)]", "score": score_value}},
+            {{"description": "Description of point 2 [Data: Reports (report ids)]", "score": score_value}}
+        ]
+    }}
+
+    ---Data tables---
+
+    """
+    map_prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                MAP_SYSTEM_PROMPT,
+            ),
+            ("system", "{context_data}"),
+            (
+                "human",
+                "{question}",
+            ),
+        ]
+    )
+
+    map_chain = map_prompt | chat_llm | StrOutputParser()
+
+    REDUCE_SYSTEM_PROMPT = """
+    ---Role---
+
+    You are a helpful assistant responding to questions about a dataset by synthesizing perspectives from multiple analysts.
+
+
+    ---Goal---
+
+    Generate a response of the target length and format that responds to the user's question, summarize all the reports from multiple analysts who focused on different parts of the dataset.
+
+    Note that the analysts' reports provided below are ranked in the **descending order of importance**.
+
+    If you don't know the answer or if the provided reports do not contain sufficient information to provide an answer, just say so. Do not make anything up.
+
+    The final response should remove all irrelevant information from the analysts' reports and merge the cleaned information into a comprehensive answer that provides explanations of all the key points and implications appropriate for the response length and format.
+
+    Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
+
+    The response shall preserve the original meaning and use of modal verbs such as "shall", "may" or "will".
+
+    The response should also preserve all the data references previously included in the analysts' reports, but do not mention the roles of multiple analysts in the analysis process.
+
+    **Do not list more than 5 record ids in a single reference**. Instead, list the top 5 most relevant record ids and add "+more" to indicate that there are more.
+
+    For example:
+
+    "Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Reports (2, 7, 34, 46, 64, +more)]. He is also CEO of company X [Data: Reports (1, 3)]"
+
+    where 1, 2, 3, 7, 34, 46, and 64 represent the id (not the index) of the relevant data record.
+
+    Do not include information where the supporting evidence for it is not provided.
+
+
+    ---Target response length and format---
+
+    {response_type}
+
+
+    ---Analyst Reports---
+
+    {report_data}
+
+
+    ---Goal---
+
+    Generate a response of the target length and format that responds to the user's question, summarize all the reports from multiple analysts who focused on different parts of the dataset.
+
+    Note that the analysts' reports provided below are ranked in the **descending order of importance**.
+
+    If you don't know the answer or if the provided reports do not contain sufficient information to provide an answer, just say so. Do not make anything up.
+
+    The final response should remove all irrelevant information from the analysts' reports and merge the cleaned information into a comprehensive answer that provides explanations of all the key points and implications appropriate for the response length and format.
+
+    The response shall preserve the original meaning and use of modal verbs such as "shall", "may" or "will".
+
+    The response should also preserve all the data references previously included in the analysts' reports, but do not mention the roles of multiple analysts in the analysis process.
+
+    **Do not list more than 5 record ids in a single reference**. Instead, list the top 5 most relevant record ids and add "+more" to indicate that there are more.
+
+    For example:
+
+    "Person X is the owner of Company Y and subject to many allegations of wrongdoing [Data: Reports (2, 7, 34, 46, 64, +more)]. He is also CEO of company X [Data: Reports (1, 3)]"
+
+    where 1, 2, 3, 7, 34, 46, and 64 represent the id (not the index) of the relevant data record.
+
+    Do not include information where the supporting evidence for it is not provided.
+
+
+    ---Target response length and format---
+
+    {response_type}
+
+    Add sections and commentary to the response as appropriate for the length and format. Style the response in markdown.
+    """
+
+    reduce_prompt = ChatPromptTemplate.from_messages(
+        [
+            (
+                "system",
+                REDUCE_SYSTEM_PROMPT,
+            ),
+            (
+                "human",
+                "{question}",
+            ),
+        ]
+    )
+
+    reduce_chain = reduce_prompt | chat_llm | StrOutputParser()
+
+    graph = Neo4jGraph(
+        url=NEO4J_URI,
+        username=NEO4J_USERNAME,
+        password=NEO4J_PASSWORD,
+        refresh_schema=False,
+    )
+
+    community_data = graph.query(
+        """
+    MATCH (c:__Community__)
+    WHERE c.level = $level
+    RETURN c.full_content AS output
+    """,
+        params={"level": level},
+    )
+    # print(community_data)
+    intermediate_results = []
+    i = 0
+    for community in tqdm(community_data[:10], desc="Processing communities"):
+        intermediate_response = map_chain.invoke(
+            {"question": query, "context_data": community["output"]}
+        )
+        intermediate_results.append(intermediate_response)
+        i += 1
+
+    final_response = reduce_chain.invoke(
+        {
+            "report_data": intermediate_results,
+            "question": query,
+            "response_type": response_type,
+        }
+    )
+    return final_response

requirements.txt

@@ -1 +1,9 @@
-huggingface_hub==0.22.2
+huggingface_hub==0.22.2
+sentence_transformers
+numpy
+pandas
+neo4j
+langchain_community
+langchain_core
+langchain_huggingface
+tqdm