Haseeb-001 commited on
Commit
385bd53
·
verified ·
1 Parent(s): 9d63b8a

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +179 -0
app.py ADDED
@@ -0,0 +1,179 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import streamlit as st
3
+ import numpy as np
4
+ import faiss
5
+ from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModel
6
+ from groq import Groq
7
+
8
+ # Load API Key from Environment
9
+ groq_api_key = os.environ.get("GROQ_API_KEY")
10
+ if groq_api_key is None:
11
+ st.error("GROQ_API_KEY environment variable not set.")
12
+ st.stop()
13
+
14
+ # Initialize Groq Client
15
+ try:
16
+ client = Groq(api_key=groq_api_key)
17
+ except Exception as e:
18
+ st.error(f"Error initializing Groq client: {e}")
19
+ st.stop()
20
+
21
+ # Load PubMedBERT Model (Try Groq API first, then Hugging Face)
22
+ try:
23
+ pubmedbert_tokenizer = AutoTokenizer.from_pretrained("NeuML/pubmedbert-base-embeddings")
24
+ pubmedbert_model = AutoModel.from_pretrained("NeuML/pubmedbert-base-embeddings")
25
+ pubmedbert_pipeline = pipeline('feature-extraction', model=pubmedbert_model, tokenizer=pubmedbert_tokenizer, device=-1)
26
+ except Exception:
27
+ st.warning("Error loading PubMedBERT from Groq API. Using Hugging Face model.")
28
+ pubmedbert_tokenizer = AutoTokenizer.from_pretrained("microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext")
29
+ pubmedbert_model = AutoModelForSequenceClassification.from_pretrained("microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext")
30
+ pubmedbert_pipeline = pipeline('feature-extraction', model=pubmedbert_model, tokenizer=pubmedbert_tokenizer, device=-1)
31
+
32
+ # Initialize FAISS Index
33
+ embedding_dim = 768
34
+ index = faiss.IndexFlatL2(embedding_dim)
35
+
36
+ # Function to Check if Query is Related to Epilepsy
37
+ def preprocess_query(query):
38
+ tokens = query.lower().split()
39
+ epilepsy_keywords = ["seizure", "epilepsy", "convulsion", "neurology", "brain activity"]
40
+
41
+ is_epilepsy_related = any(k in tokens for k in epilepsy_keywords)
42
+
43
+ return tokens, is_epilepsy_related
44
+
45
+ # Function to Generate Response with Chat History
46
+ def generate_response(user_query, chat_history):
47
+ # Grammatical Correction using LLaMA (Hidden from User)
48
+ try:
49
+ correction_prompt = f"""
50
+ Correct the following user query for grammar and spelling errors, but keep the original intent intact.
51
+ Do not add or remove any information, just fix the grammar.
52
+ User Query: {user_query}
53
+ Corrected Query:
54
+ """
55
+ grammar_completion = client.chat.completions.create(
56
+ messages=[{"role": "user", "content": correction_prompt}],
57
+ model="llama-3.3-70b-versatile",
58
+ stream=False,
59
+ )
60
+ corrected_query = grammar_completion.choices[0].message.content.strip()
61
+ # If correction fails or returns empty, use original query
62
+ if not corrected_query:
63
+ corrected_query = user_query
64
+ except Exception as e:
65
+ corrected_query = user_query # Fallback to original query if correction fails
66
+ print(f"⚠️ Grammar correction error: {e}") # Optional: Log the error for debugging
67
+
68
+ tokens, is_epilepsy_related = preprocess_query(corrected_query) # Use corrected query for processing
69
+
70
+ # Greeting Responses
71
+ greetings = ["hello", "hi", "hey"]
72
+ if any(word in tokens for word in greetings):
73
+ return "👋 Hello! How can I assist you today?"
74
+
75
+ # If Epilepsy Related - Use Epilepsy Focused Response
76
+ if is_epilepsy_related:
77
+ # Try Getting Medical Insights from PubMedBERT
78
+ try:
79
+ pubmedbert_embeddings = pubmedbert_pipeline(corrected_query) # Use corrected query for PubMedBERT
80
+ embedding_mean = np.mean(pubmedbert_embeddings[0], axis=0)
81
+ index.add(np.array([embedding_mean]))
82
+ pubmedbert_insights = "**PubMedBERT Analysis:** ✅ Query is relevant to epilepsy research."
83
+ except Exception as e:
84
+ pubmedbert_insights = f"⚠️ Error during PubMedBERT analysis: {e}"
85
+
86
+ # Use LLaMA for Final Response Generation with Chat History Context (Epilepsy Focus)
87
+ try:
88
+ prompt_history = ""
89
+ if chat_history:
90
+ prompt_history += "**Chat History:**\n"
91
+ for message in chat_history:
92
+ prompt_history += f"{message['role'].capitalize()}: {message['content']}\n"
93
+ prompt_history += "\n"
94
+
95
+ epilepsy_prompt = f"""
96
+ {prompt_history}
97
+ **User Query:** {corrected_query} # Use corrected query for final response generation
98
+ **Instructions:** Provide a concise, structured, and human-friendly response specifically about epilepsy or seizures, considering the conversation history if available.
99
+ """
100
+
101
+ chat_completion = client.chat.completions.create(
102
+ messages=[{"role": "user", "content": epilepsy_prompt}],
103
+ model="llama-3.3-70b-versatile",
104
+ stream=False,
105
+ )
106
+ model_response = chat_completion.choices[0].message.content.strip()
107
+ except Exception as e:
108
+ model_response = f"⚠️ Error generating response with LLaMA: {e}"
109
+
110
+ return f"**NeuroGuard:** ✅ **Analysis:**\n{pubmedbert_insights}\n\n**Response:**\n{model_response}"
111
+
112
+
113
+ # If Not Epilepsy Related - Try to Answer as General Health Query
114
+ else:
115
+ # Try Getting Medical Insights from PubMedBERT (even for general health)
116
+ try:
117
+ pubmedbert_embeddings = pubmedbert_pipeline(corrected_query)
118
+ embedding_mean = np.mean(pubmedbert_embeddings[0], axis=0)
119
+ index.add(np.array([embedding_mean]))
120
+ pubmedbert_insights = "**PubMedBERT Analysis:** PubMed analysis performed for health-related context." # General analysis message
121
+ except Exception as e:
122
+ pubmedbert_insights = f"⚠️ Error during PubMedBERT analysis: {e}"
123
+
124
+ # Use LLaMA for General Health Response Generation with Chat History Context
125
+ try:
126
+ prompt_history = ""
127
+ if chat_history:
128
+ prompt_history += "**Chat History:**\n"
129
+ for message in chat_history:
130
+ prompt_history += f"{message['role'].capitalize()}: {message['content']}\n"
131
+ prompt_history += "\n"
132
+
133
+ general_health_prompt = f"""
134
+ {prompt_history}
135
+ **User Query:** {corrected_query}
136
+ **Instructions:** Provide a concise, structured, and human-friendly response to the general health query, considering the conversation history if available. If the query is clearly not health-related, respond generally.
137
+ """
138
+
139
+ chat_completion = client.chat.completions.create(
140
+ messages=[{"role": "user", "content": general_health_prompt}],
141
+ model="llama-3.3-70b-versatile",
142
+ stream=False,
143
+ )
144
+ model_response = chat_completion.choices[0].message.content.strip()
145
+ except Exception as e:
146
+ model_response = f"⚠️ Error generating response with LLaMA: {e}"
147
+
148
+ return f"**NeuroGuard:** ✅ **Analysis:**\n{pubmedbert_insights}\n\n**Response:**\n{model_response}"
149
+
150
+
151
+ # Streamlit UI Setup
152
+ st.set_page_config(page_title="NeuroGuard: Epilepsy & Health Chatbot", layout="wide") # Updated title
153
+ st.title("🧠 NeuroGuard: Epilepsy & Health Chatbot") # Updated title
154
+ st.write("💬 Ask me anything about epilepsy, seizures, and general health. I remember our conversation!") # Updated description
155
+
156
+ # Initialize Chat History in Session State
157
+ if "chat_history" not in st.session_state:
158
+ st.session_state.chat_history = []
159
+
160
+ # Display Chat History
161
+ for message in st.session_state.chat_history:
162
+ with st.chat_message(message["role"]):
163
+ st.markdown(message["content"])
164
+
165
+ # User Input
166
+ if prompt := st.chat_input("Type your question here..."):
167
+ st.session_state.chat_history.append({"role": "user", "content": prompt})
168
+ with st.chat_message("user"):
169
+ st.markdown(prompt)
170
+
171
+ # Generate Bot Response
172
+ with st.chat_message("bot"):
173
+ with st.spinner("🤖 Thinking..."):
174
+ try:
175
+ response = generate_response(prompt, st.session_state.chat_history) # Pass chat history here
176
+ st.markdown(response)
177
+ st.session_state.chat_history.append({"role": "bot", "content": response})
178
+ except Exception as e:
179
+ st.error(f"⚠️ Error processing query: {e}")