ksatzke commited on
Commit
1218270
Β·
verified Β·
1 Parent(s): bfdc71d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +356 -0
app.py ADDED
@@ -0,0 +1,356 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
4
+ import logging
5
+ from typing import List, Dict
6
+ import gc
7
+ import os
8
+
9
+ # Setup logging
10
+ logging.basicConfig(
11
+ level=logging.INFO,
12
+ format='%(asctime)s - %(levelname)s - %(message)s'
13
+ )
14
+ logger = logging.getLogger(__name__)
15
+
16
+ # Set environment variables for memory optimization
17
+ os.environ['TRANSFORMERS_CACHE'] = '/home/user/.cache/huggingface/hub'
18
+ os.environ['TOKENIZERS_PARALLELISM'] = 'false'
19
+
20
+ class HealthAssistant:
21
+ def __init__(self):
22
+ self.model_id = "microsoft/Phi-2" # Using smaller Phi-2 model
23
+ self.model = None
24
+ self.tokenizer = None
25
+ self.pipe = None
26
+ self.metrics = []
27
+ self.medications = []
28
+ self.device = "cpu"
29
+ self.is_model_loaded = False
30
+ self.max_history_length = 2
31
+
32
+ def initialize_model(self):
33
+ try:
34
+ if self.is_model_loaded:
35
+ return True
36
+
37
+ logger.info(f"Loading model: {self.model_id}")
38
+
39
+ self.tokenizer = AutoTokenizer.from_pretrained(
40
+ self.model_id,
41
+ trust_remote_code=True,
42
+ model_max_length=256,
43
+ padding_side="left"
44
+ )
45
+ logger.info("Tokenizer loaded")
46
+
47
+ self.model = AutoModelForCausalLM.from_pretrained(
48
+ self.model_id,
49
+ torch_dtype=torch.float32,
50
+ trust_remote_code=True,
51
+ device_map=None,
52
+ low_cpu_mem_usage=True
53
+ ).to(self.device)
54
+
55
+ gc.collect()
56
+
57
+ self.pipe = pipeline(
58
+ "text-generation",
59
+ model=self.model,
60
+ tokenizer=self.tokenizer,
61
+ device=self.device,
62
+ model_kwargs={"low_cpu_mem_usage": True}
63
+ )
64
+
65
+ self.is_model_loaded = True
66
+ logger.info("Model initialized successfully")
67
+ return True
68
+
69
+ except Exception as e:
70
+ logger.error(f"Error in model initialization: {str(e)}")
71
+ raise
72
+
73
+ def unload_model(self):
74
+ if hasattr(self, 'model') and self.model is not None:
75
+ del self.model
76
+ self.model = None
77
+ if hasattr(self, 'pipe') and self.pipe is not None:
78
+ del self.pipe
79
+ self.pipe = None
80
+ if hasattr(self, 'tokenizer') and self.tokenizer is not None:
81
+ del self.tokenizer
82
+ self.tokenizer = None
83
+ self.is_model_loaded = False
84
+ gc.collect()
85
+ logger.info("Model unloaded successfully")
86
+
87
+ def generate_response(self, message: str, history: List = None) -> str:
88
+ try:
89
+ if not self.is_model_loaded:
90
+ self.initialize_model()
91
+
92
+ message = message[:200] # Truncate long messages
93
+
94
+ prompt = self._prepare_prompt(message, history[-self.max_history_length:] if history else None)
95
+
96
+ generation_args = {
97
+ "max_new_tokens": 200,
98
+ "return_full_text": False,
99
+ "temperature": 0.7,
100
+ "do_sample": True,
101
+ "top_k": 50,
102
+ "top_p": 0.9,
103
+ "repetition_penalty": 1.1,
104
+ "num_return_sequences": 1,
105
+ "batch_size": 1
106
+ }
107
+
108
+ output = self.pipe(prompt, **generation_args)
109
+ response = output[0]['generated_text']
110
+
111
+ gc.collect()
112
+
113
+ return response.strip()
114
+
115
+ except Exception as e:
116
+ logger.error(f"Error generating response: {str(e)}")
117
+ return "I apologize, but I encountered an error. Please try again."
118
+
119
+ def _prepare_prompt(self, message: str, history: List = None) -> str:
120
+ prompt_parts = [
121
+ "Medical AI assistant. Be professional, include disclaimers.",
122
+ self._get_health_context()
123
+ ]
124
+
125
+ if history:
126
+ for h in history:
127
+ if isinstance(h, dict): # New message format
128
+ if h['role'] == 'user':
129
+ prompt_parts.append(f"Human: {h['content'][:100]}")
130
+ else:
131
+ prompt_parts.append(f"Assistant: {h['content'][:100]}")
132
+ else: # Old format (tuple)
133
+ prompt_parts.extend([
134
+ f"Human: {h[0][:100]}",
135
+ f"Assistant: {h[1][:100]}"
136
+ ])
137
+
138
+ prompt_parts.extend([
139
+ f"Human: {message}",
140
+ "Assistant:"
141
+ ])
142
+
143
+ return "\n".join(prompt_parts)
144
+
145
+ def _get_health_context(self) -> str:
146
+ if not self.metrics and not self.medications:
147
+ return "No health data"
148
+
149
+ context = []
150
+ if self.metrics:
151
+ latest = self.metrics[-1]
152
+ context.append(f"Metrics: W:{latest['Weight']}kg S:{latest['Steps']} Sl:{latest['Sleep']}h")
153
+
154
+ if self.medications:
155
+ meds = [f"{m['Medication']}({m['Dosage']}@{m['Time']})" for m in self.medications[-2:]]
156
+ context.append("Meds: " + ", ".join(meds))
157
+
158
+ return " | ".join(context)
159
+
160
+ def add_metrics(self, weight: float, steps: int, sleep: float) -> bool:
161
+ try:
162
+ if len(self.metrics) >= 5:
163
+ self.metrics.pop(0)
164
+
165
+ self.metrics.append({
166
+ 'Weight': weight,
167
+ 'Steps': steps,
168
+ 'Sleep': sleep
169
+ })
170
+ return True
171
+ except Exception as e:
172
+ logger.error(f"Error adding metrics: {e}")
173
+ return False
174
+
175
+ def add_medication(self, name: str, dosage: str, time: str, notes: str = "") -> bool:
176
+ try:
177
+ if len(self.medications) >= 5:
178
+ self.medications.pop(0)
179
+
180
+ self.medications.append({
181
+ 'Medication': name,
182
+ 'Dosage': dosage,
183
+ 'Time': time,
184
+ 'Notes': notes
185
+ })
186
+ return True
187
+ except Exception as e:
188
+ logger.error(f"Error adding medication: {e}")
189
+ return False
190
+
191
+ class GradioInterface:
192
+ def __init__(self):
193
+ try:
194
+ logger.info("Initializing Health Assistant...")
195
+ self.assistant = HealthAssistant()
196
+ logger.info("Health Assistant initialized successfully")
197
+ except Exception as e:
198
+ logger.error(f"Failed to initialize Health Assistant: {e}")
199
+ raise
200
+
201
+ def chat_response(self, message: str, history: List) -> tuple:
202
+ if not message.strip():
203
+ return "", history
204
+
205
+ try:
206
+ response = self.assistant.generate_response(message, history)
207
+ # Convert to new message format
208
+ history.append({"role": "user", "content": message})
209
+ history.append({"role": "assistant", "content": response})
210
+
211
+ if len(history) % 3 == 0:
212
+ self.assistant.unload_model()
213
+
214
+ return "", history
215
+ except Exception as e:
216
+ logger.error(f"Error in chat response: {e}")
217
+ return "", history + [
218
+ {"role": "user", "content": message},
219
+ {"role": "assistant", "content": "I apologize, but I encountered an error. Please try again."}
220
+ ]
221
+
222
+ def add_health_metrics(self, weight: float, steps: int, sleep: float) -> str:
223
+ if not all([weight is not None, steps is not None, sleep is not None]):
224
+ return "⚠️ Please fill in all metrics."
225
+
226
+ if weight <= 0 or steps < 0 or sleep < 0:
227
+ return "⚠️ Please enter valid positive numbers."
228
+
229
+ if self.assistant.add_metrics(weight, steps, sleep):
230
+ return f"""βœ… Health metrics saved successfully!
231
+ β€’ Weight: {weight} kg
232
+ β€’ Steps: {steps}
233
+ β€’ Sleep: {sleep} hours"""
234
+ return "❌ Error saving metrics."
235
+
236
+ def add_medication_info(self, name: str, dosage: str, time: str, notes: str) -> str:
237
+ if not all([name, dosage, time]):
238
+ return "⚠️ Please fill in all required fields."
239
+
240
+ if self.assistant.add_medication(name, dosage, time, notes):
241
+ return f"""βœ… Medication added successfully!
242
+ β€’ Medication: {name}
243
+ β€’ Dosage: {dosage}
244
+ β€’ Time: {time}
245
+ β€’ Notes: {notes if notes else 'None'}"""
246
+ return "❌ Error adding medication."
247
+
248
+ def create_interface(self):
249
+ with gr.Blocks(title="Medical Health Assistant") as demo:
250
+ gr.Markdown("""
251
+ # πŸ₯ Medical Health Assistant
252
+ This AI assistant provides general health information and guidance.
253
+ """)
254
+
255
+ with gr.Tabs():
256
+ with gr.Tab("πŸ’¬ Medical Consultation"):
257
+ chatbot = gr.Chatbot(
258
+ value=[],
259
+ height=400,
260
+ label=False,
261
+ type="messages" # Using new message format
262
+ )
263
+ with gr.Row():
264
+ msg = gr.Textbox(
265
+ placeholder="Ask your health question...",
266
+ lines=1,
267
+ label=False,
268
+ scale=9
269
+ )
270
+ send_btn = gr.Button("Send", scale=1)
271
+ clear_btn = gr.Button("Clear Chat")
272
+
273
+ with gr.Tab("πŸ“Š Health Metrics"):
274
+ gr.Markdown("### Track Your Health Metrics")
275
+ with gr.Row():
276
+ weight_input = gr.Number(
277
+ label="Weight (kg)",
278
+ minimum=0,
279
+ maximum=500
280
+ )
281
+ steps_input = gr.Number(
282
+ label="Steps",
283
+ minimum=0,
284
+ maximum=100000
285
+ )
286
+ sleep_input = gr.Number(
287
+ label="Hours Slept",
288
+ minimum=0,
289
+ maximum=24
290
+ )
291
+ metrics_btn = gr.Button("Save Metrics")
292
+ metrics_status = gr.Markdown()
293
+
294
+ with gr.Tab("πŸ’Š Medication Manager"):
295
+ gr.Markdown("### Track Your Medications")
296
+ med_name = gr.Textbox(
297
+ label="Medication Name",
298
+ placeholder="Enter medication name"
299
+ )
300
+ with gr.Row():
301
+ med_dosage = gr.Textbox(
302
+ label="Dosage",
303
+ placeholder="e.g., 500mg"
304
+ )
305
+ med_time = gr.Textbox(
306
+ label="Time",
307
+ placeholder="e.g., 9:00 AM"
308
+ )
309
+ med_notes = gr.Textbox(
310
+ label="Notes (optional)",
311
+ placeholder="Additional instructions or notes"
312
+ )
313
+ med_btn = gr.Button("Add Medication")
314
+ med_status = gr.Markdown()
315
+
316
+ msg.submit(self.chat_response, [msg, chatbot], [msg, chatbot])
317
+ send_btn.click(self.chat_response, [msg, chatbot], [msg, chatbot])
318
+ clear_btn.click(lambda: [], None, chatbot)
319
+
320
+ metrics_btn.click(
321
+ self.add_health_metrics,
322
+ inputs=[weight_input, steps_input, sleep_input],
323
+ outputs=[metrics_status]
324
+ )
325
+
326
+ med_btn.click(
327
+ self.add_medication_info,
328
+ inputs=[med_name, med_dosage, med_time, med_notes],
329
+ outputs=[med_status]
330
+ )
331
+
332
+ gr.Markdown("""
333
+ ### ⚠️ Medical Disclaimer
334
+ This AI assistant provides general health information only. Not a replacement for professional medical advice.
335
+ Always consult healthcare professionals for medical decisions.
336
+ """)
337
+
338
+ demo.queue(max_size=5)
339
+
340
+ return demo
341
+
342
+ def main():
343
+ try:
344
+ interface = GradioInterface()
345
+ demo = interface.create_interface()
346
+ demo.launch(
347
+ server_name="0.0.0.0",
348
+ show_error=True,
349
+ share=True
350
+ )
351
+ except Exception as e:
352
+ logger.error(f"Error starting application: {e}")
353
+ raise
354
+
355
+ if __name__ == "__main__":
356
+ main()