Update app.py

app.py

@@ -1,175 +1,189 @@
 import gradio as gr
 import pandas as pd
 from datetime import datetime
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import gc
+from typing import List, Dict, Optional
 import os
-from typing import Optional, Dict, List, Any
-import json
 
-# ====================== MODEL CONFIGURATION ======================
-FALCON_MODEL = "tiiuae/falcon-7b-instruct"
-BACKUP_MODEL = "google/flan-t5-base"  # Smaller model as backup
-
-class LLMHandler:
+# ================== Model Configuration ==================
+class ModelHandler:
     def __init__(self):
-        self.client = self._init_client()
-        self.has_llm = self.client is not None
-
-    def _init_client(self) -> Optional[InferenceClient]:
+        self.model_name = "google/flan-t5-large"
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.tokenizer = None
+        self.model = None
+        self.initialize_model()
+        
+    def initialize_model(self):
         try:
-            return InferenceClient(
-                model=FALCON_MODEL,
-                token=os.getenv("HF_TOKEN")
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            self.model = AutoModelForSeq2SeqLM.from_pretrained(
+                self.model_name,
+                torch_dtype=torch.float32,
+                low_cpu_mem_usage=True
             )
+            self.model.to(self.device)
+            return True
         except Exception as e:
-            print(f"Primary LLM initialization failed: {str(e)}")
-            try:
-                return InferenceClient(model=BACKUP_MODEL)
-            except:
-                return None
+            print(f"Error initializing model: {str(e)}")
+            return False
 
-    def get_response(self, prompt: str, temperature: float = 0.7) -> str:
-        if not self.has_llm:
-            return "Service is running in basic mode. Using template responses."
-            
+    def generate_response(self, prompt: str, max_length: int = 512) -> str:
         try:
-            response = self.client.text_generation(
+            # Clear memory
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+            # Prepare input
+            inputs = self.tokenizer(
                 prompt,
-                max_new_tokens=512,
-                temperature=temperature,
-                repetition_penalty=1.1,
-                do_sample=True,
-                top_k=10,
-                top_p=0.95
-            )
+                return_tensors="pt",
+                truncation=True,
+                max_length=512
+            ).to(self.device)
+
+            # Generate response
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs.input_ids,
+                    max_length=max_length,
+                    num_beams=2,
+                    temperature=0.7,
+                    no_repeat_ngram_size=3,
+                    length_penalty=1.0
+                )
+
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+            # Clear memory
+            del outputs, inputs
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
             return response
         except Exception as e:
-            return f"Error: {str(e)}"
+            return f"Error generating response: {str(e)}"
+
+    def clear_memory(self):
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
-# ====================== DATA MANAGEMENT ======================
-class HealthDataManager:
+# ================== Data Management ==================
+class HealthData:
     def __init__(self):
-        self.metrics: List[Dict[str, Any]] = []
-        self.medications: List[Dict[str, Any]] = []
+        self.metrics = []
+        self.medications = []
 
-    def add_metrics(self, metrics: Dict[str, Any]) -> bool:
+    def add_metrics(self, metrics: Dict) -> bool:
         try:
             self.metrics.append({
                 'Date': datetime.now().strftime('%Y-%m-%d'),
                 **metrics
             })
             return True
-        except Exception:
+        except:
             return False
 
-    def add_medication(self, medication: Dict[str, Any]) -> bool:
+    def add_medication(self, medication: Dict) -> bool:
         try:
             self.medications.append(medication)
             return True
-        except Exception:
+        except:
             return False
 
-    def get_metrics_df(self) -> pd.DataFrame:
-        return pd.DataFrame(self.metrics)
-
-    def get_medications_df(self) -> pd.DataFrame:
-        return pd.DataFrame(self.medications)
-
-    def get_latest_metrics(self) -> Optional[Dict[str, Any]]:
-        return self.metrics[-1] if self.metrics else None
-
-    def format_health_context(self) -> str:
-        context = []
+    def get_health_context(self) -> str:
+        context_parts = []
         
-        # Add metrics context
         if self.metrics:
-            latest = self.get_latest_metrics()
-            context.append(f"""Recent Health Metrics:
-- Date: {latest['Date']}
-- Weight: {latest['Weight']} kg
-- Steps: {latest['Steps']}
-- Sleep: {latest['Sleep']} hours""")
+            latest = self.metrics[-1]
+            context_parts.append(f"Recent Health Metrics (Date: {latest['Date']}):")
+            context_parts.append(f"- Weight: {latest['Weight']} kg")
+            context_parts.append(f"- Steps: {latest['Steps']}")
+            context_parts.append(f"- Sleep: {latest['Sleep']} hours")
 
-            # Add trends
-            if len(self.metrics) > 1:
-                df = self.get_metrics_df()
-                context.append("\nTrends:")
-                for column in ['Weight', 'Steps', 'Sleep']:
-                    trend = df[column].diff().iloc[-1]
-                    if not pd.isna(trend):
-                        direction = "increased" if trend > 0 else "decreased"
-                        context.append(f"- {column} has {direction} by {abs(trend):.1f}")
-
-        # Add medications context
         if self.medications:
-            context.append("\nCurrent Medications:")
+            context_parts.append("\nCurrent Medications:")
             for med in self.medications:
-                context.append(f"- {med['Medication']} ({med['Dosage']}) at {med['Time']}")
+                med_info = f"- {med['Medication']} ({med['Dosage']}) at {med['Time']}"
                 if med['Notes']:
-                    context.append(f"  Note: {med['Notes']}")
+                    med_info += f" | Note: {med['Notes']}"
+                context_parts.append(med_info)
 
-        return "\n".join(context) if context else ""
+        return "\n".join(context_parts) if context_parts else "No health data available."
 
-# ====================== HEALTH ASSISTANT ======================
+# ================== Health Assistant ==================
 class HealthAssistant:
     def __init__(self):
-        self.llm = LLMHandler()
-        self.data_manager = HealthDataManager()
-
-    def analyze_symptoms(self, symptoms: str) -> str:
-        if not symptoms:
-            return "Please describe your symptoms."
-            
-        prompt = f"""Analyze these symptoms carefully:
-        {symptoms}
+        self.model = ModelHandler()
+        self.data = HealthData()
+        self.request_count = 0
+
+    def _create_prompt(self, message: str, context: str = "", history: List = None) -> str:
+        prompt_parts = [
+            "You are a helpful healthcare assistant. Provide accurate and helpful information.",
+            f"User Health Information:\n{context}" if context else "",
+            "Previous conversation:",
+        ]
         
-        Provide a structured analysis with:
-        1. Risk Level (Low/Medium/High)
-        2. Key Symptoms Identified
-        3. Possible Causes
-        4. Recommended Actions
-        5. Urgency Level
+        if history:
+            for h in history[-3:]:  # Last 3 messages for context
+                prompt_parts.append(f"User: {h['content']}" if h['role'] == 'user' 
+                                 else f"Assistant: {h['content']}")
         
-        Remember to emphasize if immediate medical attention is needed."""
-        
-        return self.llm.get_response(prompt, temperature=0.3)
+        prompt_parts.append(f"Current question: {message}")
+        return "\n\n".join(filter(None, prompt_parts))
 
-    def chat_response(self, message: str, history: List[Dict[str, str]]) -> str:
-        if not message:
-            return ""
+    def get_response(self, message: str, history: List = None) -> str:
+        # Increment request counter
+        self.request_count += 1
         
         # Get health context
-        health_context = self.data_manager.format_health_context()
+        context = self.data.get_health_context()
         
-        # Format conversation history
-        conv_context = "\n".join([
-            f"Human: {h['content']}" if h['role'] == 'user' else f"Assistant: {h['content']}"
-            for h in history[-4:] if h['content']
-        ])
+        # Create prompt
+        prompt = self._create_prompt(message, context, history)
         
-        prompt = f"""System: You are a helpful healthcare assistant. Consider the user's health information when providing advice.
-
-{f'User Health Information:\n{health_context}\n' if health_context else ''}
-
-Previous conversation:
-{conv_context}
-
-Current question: {message}
-
-Provide a helpful response, referencing their health data if relevant. Include appropriate medical disclaimers."""
+        # Get response
+        response = self.model.generate_response(prompt)
+        
+        # Periodic memory cleanup
+        if self.request_count % 5 == 0:
+            self.model.clear_memory()
+            
+        return response
 
-        return self.llm.get_response(prompt, temperature=0.7)
+    def analyze_symptoms(self, symptoms: str) -> str:
+        if not symptoms:
+            return "Please describe your symptoms."
+            
+        prompt = (
+            "Analyze these symptoms as a medical professional:\n"
+            f"{symptoms}\n\n"
+            "Provide analysis with:\n"
+            "1. Risk Level\n"
+            "2. Key Symptoms\n"
+            "3. Possible Causes\n"
+            "4. Recommended Actions\n"
+            "5. When to Seek Medical Care"
+        )
+        
+        return self.model.generate_response(prompt)
 
-# ====================== GRADIO INTERFACE ======================
+# ================== Gradio Interface ==================
 class HealthAssistantUI:
     def __init__(self):
         self.assistant = HealthAssistant()
 
-    def user_chat(self, message: str, history: List[Dict[str, str]]) -> tuple:
+    def user_chat(self, message: str, history: List) -> tuple:
         if message.strip() == "":
             return "", history
         
-        bot_message = self.assistant.chat_response(message, history)
+        bot_message = self.assistant.get_response(message, history)
         history.append({"role": "user", "content": message})
         history.append({"role": "assistant", "content": bot_message})
         return "", history
@@ -179,8 +193,9 @@ class HealthAssistantUI:
             return "⚠️ Please fill in all metrics.", None
         
         metrics = {'Weight': weight, 'Steps': steps, 'Sleep': sleep}
-        if self.assistant.data_manager.add_metrics(metrics):
-            return "✅ Metrics saved successfully!", self.assistant.data_manager.get_metrics_df()
+        if self.assistant.data.add_metrics(metrics):
+            df = pd.DataFrame(self.assistant.data.metrics)
+            return "✅ Metrics saved successfully!", df
         return "❌ Error saving metrics", None
 
     def save_medication(self, name: str, dosage: str, time: str, notes: str) -> tuple:
@@ -193,18 +208,14 @@ class HealthAssistantUI:
             'Time': time,
             'Notes': notes or ''
         }
-        if self.assistant.data_manager.add_medication(medication):
-            return "✅ Medication added successfully!", self.assistant.data_manager.get_medications_df()
+        if self.assistant.data.add_medication(medication):
+            df = pd.DataFrame(self.assistant.data.medications)
+            return "✅ Medication added successfully!", df
         return "❌ Error adding medication", None
 
     def create_interface(self):
         with gr.Blocks(title="Virtual Health Assistant", theme=gr.themes.Soft()) as demo:
-            gr.Markdown(
-                """
-                # 🏥 Virtual Health Assistant
-                Your AI-powered health companion. Get personalized health guidance based on your data.
-                """
-            )
+            gr.Markdown("# 🏥 Virtual Health Assistant")
             
             with gr.Tabs():
                 # Chat Interface
@@ -212,36 +223,29 @@ class HealthAssistantUI:
                     with gr.Row():
                         with gr.Column(scale=3):
                             chatbot = gr.Chatbot(
-                                type='messages',
                                 show_label=False,
                                 height=450,
                                 container=True,
-                                bubble_full_width=False
                             )
                         with gr.Column(scale=1):
+                            gr.Markdown("### Your Health Info")
                             context_display = gr.Markdown(
-                                value=self.assistant.data_manager.format_health_context() or 
-                                "*No health data available yet.*"
+                                value=self.assistant.data.get_health_context()
                             )
                     
                     with gr.Row():
                         msg = gr.Textbox(
-                            placeholder="Type your health question... (Press Enter to send)",
+                            placeholder="Type your health question... (Press Enter)",
                             lines=2,
                             max_lines=2,
                             show_label=False,
                             container=False,
                             scale=9
                         )
-                        send_btn = gr.Button("Send", scale=1, variant="primary")
+                        send_btn = gr.Button("Send", scale=1)
                     
                     clear_btn = gr.Button("Clear Chat")
-                    
-                    # Event handlers
-                    msg.submit(self.user_chat, [msg, chatbot], [msg, chatbot])
-                    send_btn.click(self.user_chat, [msg, chatbot], [msg, chatbot])
-                    clear_btn.click(lambda: None, None, chatbot, queue=False)
-
+                
                 # Symptom Checker
                 with gr.Tab("🔍 Symptom Checker"):
                     symptoms_input = gr.Textbox(
@@ -251,71 +255,69 @@ class HealthAssistantUI:
                     )
                     analyze_btn = gr.Button("Analyze Symptoms", variant="primary")
                     symptoms_output = gr.Markdown()
-                    
-                    analyze_btn.click(
-                        self.assistant.analyze_symptoms,
-                        inputs=[symptoms_input],
-                        outputs=[symptoms_output]
-                    )
-
+                
                 # Health Metrics
                 with gr.Tab("📊 Health Metrics"):
                     with gr.Row():
                         with gr.Column():
-                            weight_input = gr.Number(label="Weight (kg)", minimum=0)
-                            steps_input = gr.Number(label="Steps", minimum=0)
-                            sleep_input = gr.Number(label="Hours Slept", minimum=0)
+                            weight_input = gr.Number(label="Weight (kg)")
+                            steps_input = gr.Number(label="Steps")
+                            sleep_input = gr.Number(label="Hours Slept")
                             metrics_btn = gr.Button("Save Metrics", variant="primary")
                             metrics_status = gr.Markdown()
                         
                         with gr.Column():
                             metrics_display = gr.Dataframe(
-                                headers=["Date", "Weight", "Steps", "Sleep"],
-                                label="Your Health Metrics",
-                                wrap=True
+                                headers=["Date", "Weight", "Steps", "Sleep"]
                             )
-                    
-                    metrics_btn.click(
-                        self.save_metrics,
-                        inputs=[weight_input, steps_input, sleep_input],
-                        outputs=[metrics_status, metrics_display]
-                    )
-
+                
                 # Medication Manager
                 with gr.Tab("💊 Medication Manager"):
                     with gr.Row():
                         with gr.Column():
                             med_name = gr.Textbox(label="Medication Name")
                             med_dosage = gr.Textbox(label="Dosage")
-                            med_time = gr.Textbox(label="Time (e.g., 9:00 AM)")
+                            med_time = gr.Textbox(label="Time")
                             med_notes = gr.Textbox(label="Notes (optional)")
                             med_btn = gr.Button("Add Medication", variant="primary")
                             med_status = gr.Markdown()
                         
                         with gr.Column():
                             meds_display = gr.Dataframe(
-                                headers=["Medication", "Dosage", "Time", "Notes"],
-                                label="Your Medications",
-                                wrap=True
+                                headers=["Medication", "Dosage", "Time", "Notes"]
                             )
-                    
-                    med_btn.click(
-                        self.save_medication,
-                        inputs=[med_name, med_dosage, med_time, med_notes],
-                        outputs=[med_status, meds_display]
-                    )
+
+            # Event handlers
+            msg.submit(self.user_chat, [msg, chatbot], [msg, chatbot])
+            send_btn.click(self.user_chat, [msg, chatbot], [msg, chatbot])
+            clear_btn.click(lambda: None, None, chatbot)
+            
+            analyze_btn.click(
+                self.assistant.analyze_symptoms,
+                inputs=[symptoms_input],
+                outputs=[symptoms_output]
+            )
+            
+            metrics_btn.click(
+                self.save_metrics,
+                inputs=[weight_input, steps_input, sleep_input],
+                outputs=[metrics_status, metrics_display]
+            )
+            
+            med_btn.click(
+                self.save_medication,
+                inputs=[med_name, med_dosage, med_time, med_notes],
+                outputs=[med_status, meds_display]
+            )
 
             gr.Markdown(
-                """
-                ### ⚠️ Medical Disclaimer
+                """### ⚠️ Medical Disclaimer
                 This AI assistant provides general health information only. 
-                Always consult healthcare professionals for medical advice.
-                """
+                Always consult healthcare professionals for medical advice."""
             )
 
         return demo
 
-# ====================== MAIN APPLICATION ======================
 def main():
     ui = HealthAssistantUI()
     demo = ui.create_interface()