Update app.py

app.py

@@ -1,10 +1,6 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
-import os
-
-# Setup cache directory for models
-os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
 
 class CustomChatDoctor:
     def __init__(self):
@@ -17,31 +13,13 @@ class CustomChatDoctor:
         # Model name
         model_name = "zl111/ChatDoctor"
         
-        # Setup quantization for memory efficiency
-        if self.device == "cuda":
-            # Use 8-bit quantization if GPU is available
-            quantization_config = BitsAndBytesConfig(
-                load_in_8bit=True,
-                llm_int8_threshold=6.0
-            )
-            
-            # Load tokenizer and model
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                quantization_config=quantization_config,
-                device_map="auto",
-                trust_remote_code=True
-            )
-        else:
-            # For CPU, use lighter settings
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-            self.model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                device_map="auto",
-                low_cpu_mem_usage=True,
-                trust_remote_code=True
-            )
+        # Load tokenizer and model
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            device_map="auto",
+            trust_remote_code=True
+        )
         
         print("Model loaded successfully!")
         
@@ -51,11 +29,10 @@ class CustomChatDoctor:
         Your answers should be based on verified medical information.
         If a question doesn't make any sense, or is not factually coherent, explain why instead of answering something incorrect. 
         If you don't know the answer to a question, respond with "I don't have enough information to provide a reliable answer."
-        Always include a disclaimer that you are an AI assistant and not a licensed medical professional.
         """
         
         # Initialize conversation history
-        self.reset_conversation()
+        self.conversation_history = []
     
     def generate_response(self, user_input):
         try:
@@ -67,22 +44,18 @@ class CustomChatDoctor:
             prompt += "\n".join(self.conversation_history[-10:])
             prompt += "\nAI Assistant: "
             
-            # Generate the response with appropriate parameters based on device
+            # Generate the response
             inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)
             
-            with torch.no_grad():  # Disable gradient calculations for inference
-                generation_config = {
-                    "max_new_tokens": 512,
-                    "temperature": 0.7, 
-                    "top_p": 0.9,
-                    "do_sample": True,
-                    "pad_token_id": self.tokenizer.eos_token_id
-                }
-                
+            with torch.no_grad():
                 outputs = self.model.generate(
                     input_ids=inputs.input_ids,
                     attention_mask=inputs.attention_mask,
-                    **generation_config
+                    max_new_tokens=512,
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.eos_token_id
                 )
             
             response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
@@ -93,115 +66,53 @@ class CustomChatDoctor:
             return response
         
         except Exception as e:
-            # Handle any errors during generation
             error_message = f"An error occurred: {str(e)}"
             print(error_message)
-            return "I'm sorry, I encountered an error processing your question. Please try again or ask a different question."
+            return "I'm sorry, I encountered an error processing your question. Please try again."
     
     def reset_conversation(self):
         self.conversation_history = []
         return None
 
-# Create a singleton instance to avoid reloading the model for each user
-chat_doctor = None
-
-def get_chat_doctor():
-    global chat_doctor
-    if chat_doctor is None:
-        chat_doctor = CustomChatDoctor()
-    return chat_doctor
+# Initialize the model
+chat_doctor = CustomChatDoctor()
 
-# Example inputs to help users get started
+# Define example inputs
 examples = [
     "What are the symptoms of diabetes?",
     "How can I manage my hypertension?",
-    "What should I do for a persistent headache?",
-    "Can you explain what asthma is?",
-    "What are the side effects of ibuprofen?"
+    "What should I do for a persistent headache?"
 ]
 
-# Add CSS for styling
-css = """
-.gradio-container {
-    font-family: 'Arial', sans-serif;
-}
-.disclaimer {
-    margin-top: 20px;
-    padding: 10px;
-    background-color: #f8f9fa;
-    border-left: 3px solid #f0ad4e;
-    font-size: 14px;
-}
-"""
-
-# Create welcome message function
-def welcome():
-    return "Welcome to ChatDoctor! I'm an AI assistant trained to provide medical information. How can I help you today?"
-
-# Build the Gradio interface
-with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
+# Create Gradio interface
+with gr.Blocks() as demo:
     gr.Markdown("# Your Custom ChatDoctor")
     gr.Markdown("Ask medical questions and get AI-powered responses.")
     
-    chatbot = gr.Chatbot(height=600, type="messages")
-    msg = gr.Textbox(placeholder="Type your medical question here...", lines=2)
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(placeholder="Type your medical question here...")
     
     with gr.Row():
-        submit_btn = gr.Button("Send", variant="primary")
+        submit_btn = gr.Button("Send")
         clear_btn = gr.Button("Clear Conversation")
     
-    # Display example queries that users can click on
-    gr.Examples(
-        examples=examples,
-        inputs=msg
-    )
+    gr.Examples(examples=examples, inputs=msg)
     
-    with gr.Accordion("About this AI", open=False):
-        gr.Markdown("""
-        **ChatDoctor** is a medical conversation model designed to provide general health information.
-        
-        This AI uses language models to generate responses based on patterns learned from medical texts and conversations.
-        
-        **Important Notes:**
-        - This system is for informational purposes only
-        - Not a substitute for professional medical advice
-        - In emergencies, contact emergency services immediately
-        """)
-    
-    # Add disclaimer at the bottom with custom styling
-    gr.HTML("""
-    <div class="disclaimer">
-        <strong>Disclaimer:</strong> This AI assistant provides information for educational purposes only. 
-        Always consult with a qualified healthcare provider for medical advice, diagnosis, or treatment. 
-        This tool is not intended to replace professional medical consultation.
-    </div>
+    gr.Markdown("""
+    ### Disclaimer
+    This AI assistant provides information for educational purposes only. 
+    Always consult with a qualified healthcare provider for medical advice.
     """)
     
     def respond(message, chat_history):
-        # Lazy-load model on first request
-        doctor = get_chat_doctor()
-        
-        if message.strip() == "":
-            return "", chat_history
-            
-        bot_message = doctor.generate_response(message)
-        chat_history.append({"role": "user", "content": message})
-        chat_history.append({"role": "assistant", "content": bot_message})
+        bot_message = chat_doctor.generate_response(message)
+        chat_history.append((message, bot_message))
         return "", chat_history
     
     def clear_history():
-        doctor = get_chat_doctor()
-        doctor.reset_conversation()
+        chat_doctor.reset_conversation()
         return None
     
-    # Show welcome message when the app starts
-    demo.load(lambda: None, None, chatbot, js="""
-        () => {
-            const welcomeMsg = "Welcome to ChatDoctor! I'm an AI assistant trained to provide medical information. How can I help you today?";
-            return [[null, welcomeMsg]];
-        }
-    """)
-    
     submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
     clear_btn.click(clear_history, None, chatbot)