Update app.py

app.py

@@ -1,31 +1,53 @@
-import streamlit as st
-import time
-import requests
 import os
+import streamlit as st
 from huggingface_hub import InferenceClient
+from langchain_community.vectorstores import Neo4jVector
+from transformers import AutoTokenizer, AutoModel
+import torch
+
+# Custom Embedding Class
+class CustomHuggingFaceEmbeddings:
+    def __init__(self, model_name="sentence-transformers/all-MiniLM-L6-v2"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModel.from_pretrained(model_name)
+
+    def embed_text(self, text):
+        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        return outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
+
+    def embed_query(self, text):
+        return self.embed_text(text)
+    
+    def embed_documents(self, text):
+        return self.embed_text(text)
+
+# Function to set up the Neo4j Vector Index
+@st.cache_resource 
+def setup_vector_index():
+    return Neo4jVector.from_existing_graph(
+        CustomHuggingFaceEmbeddings(),
+        url=os.environ['NEO4J_URI'],
+        username=os.environ['NEO4J_USERNAME'],
+        password=os.environ['NEO4J_PASSWORD'],
+        index_name='articles',
+        node_label="Article",
+        text_node_properties=['topic', 'title', 'abstract'],
+        embedding_node_property='embedding',
+    )
 
 # Hugging Face API Setup
 API_TOKEN = os.environ.get("HUGGINGFACE_API_TOKEN")
-GPT2XL_API_URL = "https://api-inference.huggingface.co/models/openai-community/gpt2-xl"
 MISTRAL_MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
 client = InferenceClient(api_key=API_TOKEN)
 
-# Query GPT-2 XL
-def query_from_gpt2xl(text: str):
-    headers = {"Authorization": f"Bearer {API_TOKEN}"}
-    while True:
-        response = requests.post(GPT2XL_API_URL, headers=headers, json={"inputs": text})
-        response_data = response.json()
-        if "error" in response_data and "loading" in response_data["error"]:
-            wait_time = response_data.get("estimated_time", 10)
-            st.info(f"Model is loading. Waiting for {wait_time:.2f} seconds...")
-            time.sleep(wait_time)
-        else:
-            return response_data[0]["generated_text"]
-
 # Query Mistral
-def query_from_mistral(text: str):
-    messages = [{"role": "user", "content": text}]
+def query_from_mistral(context: str, user_input: str):
+    messages = [
+        {"role": "system", "content": f"Use the following context to answer the query:\n{context}"},
+        {"role": "user", "content": user_input},
+    ]
     completion = client.chat.completions.create(
         model=MISTRAL_MODEL_NAME,
         messages=messages,
@@ -33,32 +55,65 @@ def query_from_mistral(text: str):
     )
     return completion.choices[0].message["content"]
 
+# extract data from retriever response
+def extract_data(documents):
+    result = []
+
+    for doc in documents:
+        # Extract metadata
+        publication_date = doc.metadata.get('publication_date')
+        if publication_date:
+            publication_date = publication_date.isoformat()
+        
+        # Extract page content
+        page_content = doc.page_content.strip().split("\n")
+        topic = page_content[1].strip() if len(page_content) > 1 else "N/A"
+        title = page_content[2].strip() if len(page_content) > 2 else "N/A"
+        abstract = page_content[3].strip() if len(page_content) > 3 else "N/A"
+        
+        # Format the extracted data as a string
+        doc_data = (
+            f"Publication Date: {publication_date}\n"
+            f"Topic: {topic}\n"
+            f"Title: {title}\n"
+            f"Abstract: {abstract}\n"
+        )
+        result.append(doc_data)
+    
+    return result
+
+# Main Streamlit Application
 def main():
-    st.set_page_config(page_title="Multi-Model Chat", layout="centered")
-    st.title("🤖 Multi-Model Chat")
-    st.markdown("Chat with either **GPT-2 XL** or **Mistral-7B-Instruct** via Hugging Face API.")
+    st.set_page_config(page_title="Vector Chat with Mistral", layout="centered")
+    st.title("🤖 Vector Chat with Mistral")
+    st.markdown("Chat with **Mistral-7B-Instruct** using context retrieved from a Neo4j vector index.")
+
+    # Initialize the vector index
+    vector_index = setup_vector_index()
 
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
-    model_choice = st.selectbox("Select a model:", ["GPT-2 XL", "Mistral-7B-Instruct"])
-
     with st.form(key="chat_form", clear_on_submit=True):
         user_input = st.text_input("You:", "")
         submit = st.form_submit_button("Send")
 
     if submit and user_input:
         st.session_state.messages.append({"role": "user", "content": user_input})
+
         with st.spinner("Fetching response..."):
             try:
-                if model_choice == "GPT-2 XL":
-                    response = query_from_gpt2xl(user_input)
-                elif model_choice == "Mistral-7B-Instruct":
-                    response = query_from_mistral(user_input)
+                # Retrieve context from the vector index
+                context_results = vector_index.similarity_search(user_input, top_k=3)
+                context = extract_data(context_results)[0]
+
+                # Get response from Mistral
+                response = query_from_mistral(context, user_input)
                 st.session_state.messages.append({"role": "bot", "content": response})
             except Exception as e:
                 st.error(f"Error: {e}")
 
+    # Display chat history
     for message in st.session_state.messages:
         if message["role"] == "user":
             st.markdown(f"**You:** {message['content']}")