Updated files for enhanced PubMedBERT and GPT-4o-mini integration

.gitignore

@@ -1 +1,2 @@
 venv/
+.env

app.py

@@ -1,177 +1,88 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-from transformers import BertTokenizer, BertForSequenceClassification
 import openai
 import os
+from dotenv import load_dotenv
+import requests
+from transformers import BertTokenizer, BertForSequenceClassification
+import torch
 import faiss
 import numpy as np
-import requests
-from datasets import load_dataset
 
-# Load OpenAI API key and organization ID from environment variables
+# Load .env
+load_dotenv()
+
+# API Keys and Org ID
 openai.api_key = os.getenv("OPENAI_API_KEY")
-openai.Organization = os.getenv("OPENAI_ORG_ID")
+openai.organization = os.getenv("OPENAI_ORG_ID")
+serper_api_key = os.getenv("SERPER_API_KEY")
 
 # 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)
 
-# FAISS setup for vector search (embedding-based memory)
-dimension = 768  # PubMedBERT embedding size
+# FAISS setup for vector search
+dimension = 768
 index = faiss.IndexFlatL2(dimension)
 
-# Embed text using PubMedBERT
+# Function to embed text (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)
     hidden_state = outputs.hidden_states[-1]
     return hidden_state.mean(dim=1).detach().numpy()
 
-# Add past conversation embedding to FAISS index
-past_conversation = "FDA approval for companion diagnostics requires careful documentation."
-past_embedding = embed_text(past_conversation)
-past_embedding = np.array(past_embedding)  # Convert to numpy array
-
-# Reshape if necessary (e.g., (1, 768) for PubMedBERT)
-past_embedding = past_embedding.reshape(1, -1)
-
-index.add(past_embedding)
-
-# Search past conversations/memory using FAISS
-def search_memory(query):
-    query_embedding = embed_text(query)
-    D, I = index.search(query_embedding, k=1)
-    return I
-
-# Handle FDA-specific queries with PubMedBERT
+# Function to retrieve info from PubMedBERT
 def handle_fda_query(query):
-    inputs = tokenizer(query, return_tensors="pt", padding="max_length", truncation=True)
+    inputs = tokenizer(query, return_tensors="pt", padding="max_length", truncation=True, max_length=512)
     outputs = model(**inputs)
     logits = outputs.logits
-    response = "Processed FDA-related query via PubMedBERT"
-    return response
-
-# Handle general queries using GPT-4O
-def handle_openai_query(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].message.content
-
-# Web search with Serper API
-def web_search(query):
-    url = f"https://google.serper.dev/search"
-    headers = {
-        "X-API-KEY": os.getenv("SERPER_API_KEY")
-    }
-    params = {
-        "q": query
-    }
-    response = requests.get(url, headers=headers, params=params)
-    return response.json()
-
-# 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]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    # Prepare context for OpenAI and PubMedBERT
-    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]})
-
-    messages.append({"role": "user", "content": message})
-
-    # 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}"
-
-        # If no memory match, proceed with PubMedBERT
-        return handle_fda_query(message)
-
-    # If query asks for web search, perform web search
-    if "search the web" in message.lower():
-        return web_search(message)
-
-    # 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]}"
-
-    # General conversational handling with GPT-4O
-    response = handle_openai_query(message)
-    return response
-
-
-# Create Gradio ChatInterface for interaction
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are Ferris2.0, an FDA expert.", label="System message"),
+    prediction = torch.argmax(logits, dim=1).item()
+
+    # Simulate a meaningful FDA-related response
+    if prediction == 1:
+        return f"FDA Query Processed: '{query}' contains important regulatory information."
+    else:
+        return f"FDA Query Processed: '{query}' seems to be general and not regulatory-heavy."
+
+# Function to enhance info via GPT-4o-mini
+def enhance_with_gpt4o(fda_response):
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-4o-mini",  # Correct model
+            messages=[{"role": "system", "content": "You are an expert FDA assistant."}, {"role": "user", "content": f"Enhance this FDA info: {fda_response}"}],
+            max_tokens=150
+        )
+        return response['choices'][0]['message']['content']
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Main function that gets PubMedBERT output and enhances it using GPT-4o-mini
+def respond(message, system_message, max_tokens, temperature, top_p):
+    try:
+        # First retrieve info via PubMedBERT
+        fda_response = handle_fda_query(message)
+        
+        # Then enhance this info via GPT-4o-mini
+        enhanced_response = enhance_with_gpt4o(fda_response)
+        
+        # Return both the PubMedBERT result and the enhanced version
+        return f"Original Info from PubMedBERT: {fda_response}\n\nEnhanced Info via GPT-4o-mini: {enhanced_response}"
+
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Gradio Interface
+demo = gr.Interface(
+    fn=respond,
+    inputs=[
+        gr.Textbox(label="Enter your FDA query", placeholder="Ask Ferris2.0 anything FDA-related."),
+        gr.Textbox(value="You are Ferris2.0, the most advanced FDA Regulatory Assistant.", 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.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
     ],
+    outputs="text",
 )
 
-
 if __name__ == "__main__":
     demo.launch()