UPDATED APP Directly in HF files

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
-from transformers import BertTokenizer, BertForSequenceClassification, TrainingArguments, Trainer
+from transformers import BertTokenizer, BertForSequenceClassification
 import openai
 import os
 import faiss
@@ -8,13 +8,11 @@ import numpy as np
 import requests
 from datasets import load_dataset
 
-ds = load_dataset("epfl-llm/guidelines")
+# Load OpenAI API key and organization ID from environment variables
+openai.api_key = os.getenv("OPENAI_API_KEY")
+openai.Organization = os.getenv("OPENAI_ORG_ID")
 
-# Load OpenAI and Serper API keys from Hugging Face secrets
-openai.api_key = os.getenv("OPENAI_API_KEY")  # Ensure the OpenAI API key is pulled correctly
-serper_api_key = os.getenv("SERPER_API_KEY")  # Ensure the Serper API key is pulled correctly
-
-# Load PubMedBERT tokenizer and model for FDA-related processing
+# Load PubMedBERT tokenizer and model
 tokenizer = BertTokenizer.from_pretrained("microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract")
 model = BertForSequenceClassification.from_pretrained("microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract", num_labels=2)
 
@@ -22,24 +20,25 @@ model = BertForSequenceClassification.from_pretrained("microsoft/BiomedNLP-PubMe
 dimension = 768  # PubMedBERT embedding size
 index = faiss.IndexFlatL2(dimension)
 
+# Embed text using PubMedBERT
 def embed_text(text):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding="max_length", max_length=512)
-    outputs = model(**inputs, output_hidden_states=True)  # Ensure hidden states are returned
-    hidden_state = outputs.hidden_states[-1]  # Get the last hidden state
-    return hidden_state.mean(dim=1).detach().numpy()  # Take the mean across the sequence
+    outputs = model(**inputs, output_hidden_states=True)
+    hidden_state = outputs.hidden_states[-1]
+    return hidden_state.mean(dim=1).detach().numpy()
 
-# Example: Embed past conversation and store in FAISS
+# Add past conversation embedding to FAISS index
 past_conversation = "FDA approval for companion diagnostics requires careful documentation."
 past_embedding = embed_text(past_conversation)
-index.add(past_embedding)
+index.add([past_embedding])
 
-# Embed the incoming query and search for related memory
+# Search past conversations/memory using FAISS
 def search_memory(query):
     query_embedding = embed_text(query)
-    D, I = index.search(query_embedding, k=1)  # Retrieve most similar past conversation
+    D, I = index.search(query_embedding, k=1)
     return I
 
-# Function to handle FDA-specific queries with PubMedBERT
+# Handle FDA-specific queries with PubMedBERT
 def handle_fda_query(query):
     inputs = tokenizer(query, return_tensors="pt", padding="max_length", truncation=True)
     outputs = model(**inputs)
@@ -47,20 +46,23 @@ def handle_fda_query(query):
     response = "Processed FDA-related query via PubMedBERT"
     return response
 
-# Function to handle general queries using GPT-4o
+# Handle general queries using GPT-4O
 def handle_openai_query(prompt):
-    response = openai.Completion.create(
-        engine="gpt-4o",  # Using GPT-4o as per instruction
-        prompt=prompt,
+    response = openai.Chat.create(
+        model="gpt-4-0314-16k-512",
+        messages=[
+            {"role": "user", "content": prompt}
+        ],
+        temperature=0.7,
         max_tokens=100
     )
-    return response.choices[0].text.strip()
+    return response.choices[0].message.content
 
 # Web search with Serper API
 def web_search(query):
     url = f"https://google.serper.dev/search"
     headers = {
-        "X-API-KEY": serper_api_key
+        "X-API-KEY": os.getenv("SERPER_API_KEY")
     }
     params = {
         "q": query
@@ -68,7 +70,45 @@ def web_search(query):
     response = requests.get(url, headers=headers, params=params)
     return response.json()
 
-# Main assistant function that delegates to either OpenAI, PubMedBERT, or Serper (web search)
+# Contextual Short-Term Memory (CSTM)
+cstm = []
+
+# Long-Term Memory (LTM)
+ltm = []  # Load knowledge base articles or FAQs
+
+# Semantic search function
+def semantic_search(query, cstm, ltm):
+    # Generate embeddings for query and CSTM/LTM
+    query_embedding = embed_text(query)
+    cstm_embeddings = [embed_text(text) for text in cstm]
+    ltm_embeddings = [embed_text(text) for text in ltm]
+
+    # Calculate similarity scores
+    cstm_scores = calculate_similarity(query_embedding, cstm_embeddings)
+    ltm_scores = calculate_similarity(query_embedding, ltm_embeddings)
+
+    # Retrieve top relevant results from CSTM and LTM
+    top_cstm = np.argmax(cstm_scores)
+    top_ltm = np.argmax(ltm_scores)
+
+    return top_cstm, top_ltm
+
+# Calculate similarity between embeddings
+def calculate_similarity(query_embedding, embeddings):
+    similarity_scores = []
+    for embedding in embeddings:
+        score = cosine_similarity(query_embedding, embedding)
+        similarity_scores.append(score)
+    return similarity_scores
+
+# Cosine similarity function
+def cosine_similarity(a, b):
+    dot_product = np.dot(a, b)
+    magnitude_a = np.linalg.norm(a)
+    magnitude_b = np.linalg.norm(b)
+    return dot_product / (magnitude_a * magnitude_b)
+
+# Main assistant function
 def respond(
     message,
     history: list[tuple[str, str]],
@@ -77,7 +117,7 @@ def respond(
     temperature,
     top_p,
 ):
-    # Prepare the context for OpenAI and PubMedBERT
+    # Prepare context for OpenAI and PubMedBERT
     messages = [{"role": "system", "content": system_message}]
 
     for val in history:
@@ -88,35 +128,32 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    # Check if the query is related to FDA
+    # Check if query is FDA-related
     openai_response = handle_openai_query(f"Is this query FDA-related: {message}")
     
     if "FDA" in openai_response or "regulatory" in openai_response:
         # Search past conversations/memory using FAISS
         memory_index = search_memory(message)
         if memory_index:
-            return f"Found relevant past memory: {past_conversation}"  # Return past context from memory
+            return f"Found relevant past memory: {past_conversation}"
 
         # If no memory match, proceed with PubMedBERT
         return handle_fda_query(message)
 
-    # If query asks for a web search, perform web search
+    # If query asks for web search, perform web search
     if "search the web" in message.lower():
         return web_search(message)
 
-    # General conversational handling with GPT-4o
-    response = ""
-    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
+    # Perform semantic search on CSTM and LTM
+    top_cstm, top_ltm = semantic_search(message, cstm, ltm)
+    if top_cstm:
+        return f"Found relevant context: {cstm[top_cstm]}"
+    elif top_ltm:
+        return f"Found relevant knowledge: {ltm[top_ltm]}"
 
-        response += token
-        yield response
+    # General conversational handling with GPT-4O
+    response = handle_openai_query(message)
+    return response
 
 
 # Create Gradio ChatInterface for interaction
@@ -130,5 +167,6 @@ demo = gr.ChatInterface(
     ],
 )
 
+
 if __name__ == "__main__":
     demo.launch(share=True)