Update app.py

app.py

@@ -1,8 +1,8 @@
 import os
 import torch
-from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
-from PIL import Image
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
+from PIL import Image
 import base64
 import io
 
@@ -17,8 +17,8 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load model with revision pinning
-model = AutoModel.from_pretrained(
+# Load model for causal language modeling
+model = AutoModelForCausalLM.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
     quantization_config=bnb_config,
     device_map="auto",
@@ -37,38 +37,40 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 def analyze_input(image_data, question):
     try:
-        # Prepare the prompt
+        if not question.strip():
+            return {
+                "status": "error",
+                "message": "Question is required."
+            }
+        
+        # Handle the input image (if any)
         if image_data is not None:
-            prompt = f"Given the medical image and the question: {question}\nPlease provide a detailed analysis."
-        else:
-            prompt = f"Medical question: {question}\nAnswer: "
+            return {
+                "status": "error",
+                "message": "Image support is not implemented yet."
+            }
+        
+        # Prepare prompt for text-only input
+        prompt = f"Medical question: {question}\nAnswer: "
 
         # Tokenize input
-        tokenized = tokenizer(prompt, return_tensors="pt")
-        input_ids = tokenized.input_ids.to(model.device)
-        
-        # Calculate target size (for generation length)
-        tgt_size = input_ids.size(1) + 256  # original length + max new tokens
-        
-        # Prepare model inputs
-        model_inputs = {
-            "input_ids": input_ids,
-            "pixel_values": None,  # Set to None for text-only queries
-            "tgt_sizes": [tgt_size]  # Add target size for generation
-        }
-        
+        inputs = tokenizer(prompt, return_tensors="pt")
+        input_ids = inputs.input_ids.to(model.device)
+
         # Generate response
         outputs = model.generate(
-            model_inputs=model_inputs,
+            input_ids=input_ids,
+            max_length=256,  # Limit the length of the generated text
+            eos_token_id=tokenizer.eos_token_id,  # Ensure generation stops correctly
+            pad_token_id=tokenizer.pad_token_id,
+            temperature=0.7,  # Control randomness
+            top_p=0.9,  # Nucleus sampling
+            top_k=50  # Top-k sampling
         )
-        
-        # Decode and clean up response
+
+        # Decode response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Remove the prompt from the response
-        if prompt in response:
-            response = response[len(prompt):].strip()
-        
+
         return {
             "status": "success",
             "response": response
@@ -88,7 +90,7 @@ demo = gr.Interface(
     ],
     outputs=gr.JSON(label="Analysis"),
     title="Medical Query Analysis",
-    description="Ask medical questions. For now, please focus on text-based queries without images.",
+    description="Ask medical questions. For now, please focus on text-based queries.",
     flagging_mode="never"
 )
 
@@ -97,4 +99,4 @@ demo.launch(
     share=True,
     server_name="0.0.0.0",
     server_port=7860
-)
+)