Update nexus.py

nexus.py

@@ -1,12 +1,11 @@
-from llama_index.indices.managed.vectara import VectaraIndex
+from llama_index import VectaraIndex
 from dotenv import load_dotenv
 import os
 from docx import Document
-from llama_index.llms.together import TogetherLLM
-from llama_index.core.llms import ChatMessage, MessageRole
+from llama_index.llms import TogetherLLM
+from llama_index.llms.base import ChatMessage, MessageRole
 from Bio import Entrez 
 import ssl
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import streamlit as st
 from googleapiclient.discovery import build
 from typing import List, Optional
@@ -21,42 +20,20 @@ os.environ["VECTARA_CUSTOMER_ID"] = os.getenv("VECTARA_CUSTOMER_ID", "1452235940
 os.environ["TOGETHER_API"] = os.getenv("TOGETHER_API", "7e6c200b7b36924bc1b4a5973859a20d2efa7180e9b5c977301173a6c099136b")
 os.environ["GOOGLE_SEARCH_API_KEY"] = os.getenv("GOOGLE_SEARCH_API_KEY", "AIzaSyALmmMjvmrmHGtjjuPLEMy6Bp2qgMQJ3Ck")
 
-# Initialize the Vectara index
 index = VectaraIndex()
-
 endpoint = 'https://api.together.xyz/inference'
 
-# Load the hallucination evaluation model
-model_name = "vectara/hallucination_evaluation_model"
-model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-
-def vectara_hallucination_evaluation_model(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
-    outputs = model(**inputs)
-    hallucination_probability = outputs.logits[0][0].item()  
-    return hallucination_probability
-
 def search_pubmed(query: str) -> Optional[List[str]]:
-    """
-    Searches PubMed for a given query and returns a list of formatted results 
-    (or None if no results are found).
-    """
     Entrez.email = "[email protected]"  
-
     try:
         ssl._create_default_https_context = ssl._create_unverified_context
-
         handle = Entrez.esearch(db="pubmed", term=query, retmax=3)
         record = Entrez.read(handle)
         id_list = record["IdList"]
-
         if not id_list:
             return None
-
         handle = Entrez.efetch(db="pubmed", id=id_list, retmode="xml")
         articles = Entrez.read(handle)
-
         results = []
         for article in articles['PubmedArticle']:
             try:
@@ -64,23 +41,17 @@ def search_pubmed(query: str) -> Optional[List[str]]:
                 article_data = medline_citation['Article']
                 title = article_data['ArticleTitle']
                 abstract = article_data.get('Abstract', {}).get('AbstractText', [""])[0]
-
                 result = f"**Title:** {title}\n**Abstract:** {abstract}\n"
-                result += f"**Link:** https://pubmed.ncbi.nlm.gov/{medline_citation['PMID']}\n\n"
+                result += f"**Link:** https://pubmed.ncbi.nlm.nih.gov/{medline_citation['PMID']}\n\n"
                 results.append(result)
             except KeyError as e:
                 print(f"Error parsing article: {article}, Error: {e}")
-
         return results
-
     except Exception as e:
         print(f"Error accessing PubMed: {e}")
         return None
 
 def chat_with_pubmed(article_text, article_link):
-    """
-    Engages in a chat-like interaction with a PubMed article using TogetherLLM.
-    """
     try:
         llm = TogetherLLM(model="QWEN/QWEN1.5-14B-CHAT", api_key=os.environ['TOGETHER_API'])
         messages = [
@@ -94,19 +65,11 @@ def chat_with_pubmed(article_text, article_link):
         return "An error occurred while generating a summary."
 
 def search_web(query: str, num_results: int = 3) -> Optional[List[str]]:
-    """
-    Searches the web using the Google Search API and returns a list of formatted results
-    (or None if no results are found).
-    """
     try:
         service = build("customsearch", "v1", developerKey=os.environ["GOOGLE_SEARCH_API_KEY"])
-
-        # Execute the search request
         res = service.cse().list(q=query, cx="6128965e5bcae442b", num=num_results).execute()
-
         if "items" not in res:
             return None
-
         results = []
         for item in res["items"]:
             title = item["title"]
@@ -114,26 +77,16 @@ def search_web(query: str, num_results: int = 3) -> Optional[List[str]]:
             snippet = item["snippet"]
             result = f"**Title:** {title}\n**Link:** {link}\n**Snippet:** {snippet}\n\n"
             results.append(result)
-
         return results
-
     except Exception as e:
         print(f"Error performing web search: {e}")
         return None
 
 def NEXUS_chatbot(user_input, chat_history=None):
-    """
-    Processes user input, interacts with various resources, and generates a response. 
-    Handles potential errors, maintains chat history, and evaluates hallucination risk.
-    """
-
     if chat_history is None:
         chat_history = []
-
-    response_parts = []  # Collect responses from different sources
-
+    response_parts = []
     try:
-        # Vectara Search
         try:
             query_str = user_input
             response = index.as_query_engine().query(query_str)
@@ -142,7 +95,6 @@ def NEXUS_chatbot(user_input, chat_history=None):
             print(f"Error in Vectara search: {e}")
             response_parts.append("Vectara knowledge base is currently unavailable.")
 
-        # PubMed Search and Chat
         pubmed_results = search_pubmed(user_input)
         if pubmed_results:
             response_parts.append("**PubMed Articles (Chat & Summarize):**")
@@ -153,7 +105,6 @@ def NEXUS_chatbot(user_input, chat_history=None):
         else:
             response_parts.append("No relevant PubMed articles found.")
 
-        # Web Search
         web_results = search_web(user_input)
         if web_results:
             response_parts.append("**Web Search Results:**")
@@ -161,24 +112,10 @@ def NEXUS_chatbot(user_input, chat_history=None):
         else:
             response_parts.append("No relevant web search results found.")
 
-        # Combine response parts into a single string
         response_text = "\n\n".join(response_parts)
-
-        # Hallucination Evaluation
-        def vectara_hallucination_evaluation_model(text):
-            inputs = tokenizer(text, return_tensors="pt")
-            outputs = model(**inputs)
-            hallucination_probability = outputs.logits[0][0].item()  
-            return hallucination_probability
-
-        hallucination_score = vectara_hallucination_evaluation_model(response_text)
-        HIGH_HALLUCINATION_THRESHOLD = 0.9
-        if hallucination_score > HIGH_HALLUCINATION_THRESHOLD:
-            response_text = "I'm still under development and learning. I cannot confidently answer this question yet."
-
     except Exception as e:
         print(f"Error in chatbot: {e}")
-        response_text = "An error occurred. Please try again later."
+        response_text = f"An error occurred: {str(e)}. Please try again later or rephrase your question."
 
     chat_history.append((user_input, response_text))
     return response_text, chat_history
@@ -196,35 +133,26 @@ def show_info_popup():
         *   **Not a substitute for professional medical advice:** NEXUS is not intended to replace professional medical diagnosis and treatment. Always consult a qualified healthcare provider for personalized medical advice.
         *   **General knowledge and educational purposes:** The information provided by NEXUS is for general knowledge and educational purposes only and may not be exhaustive or specific to individual situations.
         *   **Under development:** NEXUS is still under development and may occasionally provide inaccurate or incomplete information. It's important to critically evaluate responses and cross-reference with reliable sources.
-        *   **Hallucination potential:** While NEXUS employs a hallucination evaluation model to minimize the risk of generating fabricated information, there remains a possibility of encountering inaccurate responses, especially for complex or niche queries.
         **How to use:**
         1.  **Type your medical question in the text box.**
         2.  **NEXUS will provide a comprehensive response combining information from various sources.** This may include insights from its knowledge base, summaries of relevant research articles, and safe web search results.
         3.  **You can continue the conversation by asking follow-up questions or providing additional context.** This helps NEXUS refine its search and offer more tailored information.
-        4.  **in case the NEXUS doesn't show the output please check your internet connection or rerun the same command**
-        5.  **user can either chat with the documents or with generate resposne from vectara + pubmed + web search**
-        5.  **chat with document feature is still under development so it would be better to avoid using it for now**
+        4.  **In case NEXUS doesn't show the output, please check your internet connection or rerun the same command.**
         """)
 
-# Initialize session state
 if 'chat_history' not in st.session_state:
     st.session_state.chat_history = []
 
-# Define function to display chat history with highlighted user input and chatbot response
 def display_chat_history():
     for user_msg, bot_msg in st.session_state.chat_history:
         st.info(f"**You:** {user_msg}")
         st.success(f"**NEXUS:** {bot_msg}")
 
-# Define function to clear chat history
 def clear_chat():
     st.session_state.chat_history = []
 
 def main():
-    # Streamlit Page Configuration
     st.set_page_config(page_title="NEXUS Chatbot", layout="wide")
-
-    # Custom Styles
     st.markdown(
         """
         <style>
@@ -249,12 +177,8 @@ def main():
         """,
         unsafe_allow_html=True,
     )
-
-    # Title and Introduction
     st.title("NEXUS Chatbot")
     st.write("Ask your medical questions and get reliable information!")
-
-    # Example Questions (Sidebar)
     example_questions = [
         "What are the symptoms of COVID-19?",
         "How can I manage my diabetes?",
@@ -264,25 +188,18 @@ def main():
     st.sidebar.header("Example Questions")
     for question in example_questions:
         st.sidebar.write(question)
-
-    # Output Container
     output_container = st.container()
-
-    # User Input and Chat History
     input_container = st.container()
     with input_container:
         user_input = st.text_input("You: ", key="input_placeholder", placeholder="Type your medical question here...")
         new_chat_button = st.button("Start New Chat")
         if new_chat_button:
-            st.session_state.chat_history = []  # Clear chat history
-
+            st.session_state.chat_history = []
     if user_input:
         response, st.session_state.chat_history = NEXUS_chatbot(user_input, st.session_state.chat_history)
         with output_container:
             display_chat_history()
-
-    # Information Popup
     show_info_popup()
 
 if __name__ == "__main__":
-    main()
+    main()