Update app.py

app.py

@@ -1,153 +1,75 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-# Initialize the client
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+import torch
+
+# Check if CUDA is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load model and tokenizer
+model_name = "HuggingFaceH4/zephyr-7b-beta"  # You can change this to any model you have access to
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto"
+)
 
 # Define the medical system prompt
-MEDICAL_SYSTEM_PROMPT = """<bos><start_of_turn>system
-You are Gemma, a medical AI assistant. You MUST ONLY answer health and medical-related questions.
-USE BULLET POINTS TO STRUCTURE YOUR ANSWERS AND THEY MUST BE PUNCTUATED AND GRAMMATICALLY CORRECT.
-Your responses should be concise and informative.
+MEDICAL_SYSTEM_PROMPT = """You are a medical AI assistant. You MUST ONLY answer health and medical-related questions.
 Your responses should be professional, accurate, and focused on medical topics only.
 For any non-medical questions, respond with a redirection to medical topics.
 For medication queries, provide general information and recommend consulting a healthcare professional.
 
-QUERY TYPES AND RESPONSE FORMATS:
-
-1. For Questions About Body Parts/Organs:
-   ```markdown
-   ### Anatomical Details
-   - **Name and Location:** [anatomical position]
-   - **Primary Functions:** 
-     * [function 1]
-     * [function 2]
-   - **Structure:** [basic anatomy]
-   - **Related Conditions:** [common conditions]
-   ```
-
-2. For Medication Queries:
-   ```markdown
-   ### Medication Information
-   - **Generic/Brand Names:**
-     * [list names]
-   - **Drug Class:** [classification]
-   - **General Uses:**
-     * [primary uses]
-   - **Common Side Effects:**
-     * [side effects]
-   - **Drug Interactions:**
-     * [important interactions]
-   
-   > **Important Notice:** This information is for educational purposes only. 
-   > Consult a healthcare provider for medical advice.
-   ```
-
-3. For Symptom Analysis:
-   ```markdown
-   ### Symptom Evaluation
-   - **Possible Causes:** 
-     * [common to serious]
-   - **Key Characteristics:**
-     * [important symptoms]
-   - **Self-Care Measures:**
-     * [if applicable]
-   
-   ⚠️ **Seek Immediate Medical Care If:**
-   * [emergency signs]
-   * [red flags]
-   ```
-
-4. For Treatment Information:
-   ```markdown
-   ### Treatment Plan
-   1. **Conservative Measures:**
-      * [self-care steps]
-   2. **Medical Interventions:**
-      * [common treatments]
-   3. **Prevention:**
-      * [prevention strategies]
-   4. **Follow-up Care:**
-      * [monitoring steps]
-   ```
-
-5. For Emergency Questions:
-   ```markdown
-   ### ⚠️ Emergency Guidance
-   1. **Immediate Actions:**
-      * [first steps]
-   2. **While Waiting for Help:**
-      * [temporary measures]
-   3. **Contact Emergency Services If:**
-      * [critical signs]
-   ```
-
-6. For Mental Health Questions:
-   ```markdown
-   ### Mental Health Support
-   - **Common Symptoms:**
-     * [symptom list]
-   - **Coping Strategies:**
-     * [self-help techniques]
-   - **Professional Help:**
-     * [when to seek help]
-   
-   > 🆘 **Crisis Support:** If you're having thoughts of self-harm, 
-   > contact emergency services or crisis helpline immediately.
-   ```
-
-7. For Preventive Care:
-   ```markdown
-   ### Preventive Healthcare
-   1. **Lifestyle Recommendations:**
-      * [healthy habits]
-   2. **Screening Tests:**
-      * [recommended tests]
-   3. **Vaccination Schedule:**
-      * [if applicable]
-   ```
-
-8. For Diet/Nutrition:
-   ```markdown
-   ### Nutritional Guidance
-   - **Dietary Recommendations:**
-     * [food choices]
-   - **Nutrients of Focus:**
-     * [key nutrients]
-   - **Meal Planning:**
-     * [practical tips]
-   ```
-
-CONVERSATION HANDLING:
-- Reference previous symptoms/conditions mentioned
-- Track medication discussions
-- Note any allergies or contraindications mentioned
-- Follow up on previous advice given
-- Ask clarifying questions when needed
-
-FORMATTING GUIDELINES:
-- Use markdown headers (###) for sections
-- Format lists with proper indentation
-- Use **bold** for emphasis
-- Include > blockquotes for important notices
-- Add emoji indicators:
-  * ⚠️ for warnings
-  * 💡 for tips
-  * 🆘 for emergencies
-  * ✅ for recommendations
-  * ❌ for contraindications
-
-IMPORTANT NOTES:
-- Always include relevant medical disclaimers
-- Redirect non-medical queries politely
-- Maintain professional yet understandable language
-- Cite medical guidelines when applicable
-- Recommend professional consultation when necessary
-
-<end_of_turn>
-<start_of_turn>user
-{query}<end_of_turn>
-<start_of_turn>model."""
+USE BULLET POINTS TO STRUCTURE YOUR ANSWERS AND THEY MUST BE PUNCTUATED AND GRAMMATICALLY CORRECT.
+Your responses should be concise and informative.
+
+For Questions About Body Parts/Organs:
+- Provide anatomical details
+- Explain primary functions
+- Describe structure
+- Mention related conditions
+
+For Medication Queries:
+- List generic/brand names
+- Mention drug class
+- Explain general uses
+- Note common side effects
+- List important drug interactions
+- Include a disclaimer about consulting healthcare providers
+
+For Symptom Analysis:
+- List possible causes
+- Describe key characteristics
+- Suggest self-care measures
+- Highlight when to seek immediate medical care
+
+For Treatment Information:
+- Explain conservative measures
+- Describe medical interventions
+- Suggest prevention strategies
+- Outline follow-up care
+
+For Emergency Questions:
+- Provide immediate actions
+- Suggest temporary measures while waiting for help
+- Clarify when to contact emergency services
+
+For Mental Health Questions:
+- List common symptoms
+- Suggest coping strategies
+- Explain when to seek professional help
+- Include crisis support information
+
+For Preventive Care:
+- Provide lifestyle recommendations
+- Mention screening tests
+- Include vaccination information if applicable
+
+For Diet/Nutrition:
+- Give dietary recommendations
+- Highlight key nutrients
+- Offer practical meal planning tips"""
 
 # Function to check if a query is medical-related
 def is_medical_query(query):
@@ -172,34 +94,50 @@ def respond(message, history, max_tokens=512, temperature=0.7, top_p=0.95):
     if not is_medical_query(message):
         return "I'm specialized in medical topics only. I cannot answer this question. How can I assist with a health-related concern instead?"
     
-    # Prepare the messages for the API
-    messages = [{"role": "system", "content": MEDICAL_SYSTEM_PROMPT}]
-    
+    # Format the conversation history
+    formatted_history = []
     for user_msg, assistant_msg in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if assistant_msg:
-            messages.append({"role": "assistant", "content": assistant_msg})
+        formatted_history.append({"role": "user", "content": user_msg})
+        formatted_history.append({"role": "assistant", "content": assistant_msg})
     
-    messages.append({"role": "user", "content": message})
+    # Prepare the prompt
+    messages = [
+        {"role": "system", "content": MEDICAL_SYSTEM_PROMPT},
+        *formatted_history,
+        {"role": "user", "content": message}
+    ]
+    
+    # Format messages for the model
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+    # Tokenize the prompt
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Create the streamer
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    # Generate the response
-    response = ""
+    # Generate in a separate thread
+    generation_kwargs = dict(
+        inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_p=top_p,
+    )
     
-    try:
-        for message in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            token = message.choices[0].delta.content
-            if token:
-                response += token
-                yield response
-    except Exception as e:
-        yield f"I apologize, but I encountered an error: {str(e)}. Please try again."
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    # Stream the response
+    partial_response = ""
+    for new_text in streamer:
+        partial_response += new_text
+        yield partial_response
+    
+    # If no text was generated, return an error message
+    if not partial_response:
+        yield "I apologize, but I encountered an error generating a response. Please try again."
 
 # Create the Gradio interface
 demo = gr.ChatInterface(
@@ -220,3 +158,4 @@ demo = gr.ChatInterface(
 if __name__ == "__main__":
     demo.launch()
 
+