Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from PIL import Image
 import gradio as gr
 import base64
@@ -17,8 +17,8 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load model with revision pinning
-model = AutoModel.from_pretrained(
+# Load model with revision pinning - using CausalLM for text generation
+model = AutoModelForCausalLM.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
     quantization_config=bnb_config,
     device_map="auto",
@@ -37,45 +37,32 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 def analyze_input(image_data, question):
     try:
-        # Handle base64 image if provided
-        if isinstance(image_data, str) and image_data.startswith('data:image'):
-            base64_data = image_data.split(',')[1]
-            image_bytes = base64.b64decode(base64_data)
-            image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
-        # Handle direct image input
-        elif image_data is not None:
-            image = Image.fromarray(image_data).convert('RGB')
+        # Prepare the prompt
+        if image_data is not None:
+            prompt = f"Given the medical image and the question: {question}\nPlease provide a detailed analysis."
         else:
-            image = None
+            prompt = f"Medical question: {question}\nAnswer: "
 
-        # Process with or without image
-        if image is not None:
-            # Prepare inputs for multimodal generation
-            model_inputs = {
-                "input_ids": tokenizer(question, return_tensors="pt").input_ids.to(model.device),
-                "images": [image]
-            }
-        else:
-            # Prepare inputs for text-only generation
-            model_inputs = {
-                "input_ids": tokenizer(question, return_tensors="pt").input_ids.to(model.device)
-            }
-
-        # Generate response with proper inputs
-        generation_config = {
-            "max_new_tokens": 256,
-            "do_sample": True,
-            "temperature": 0.7,
-            "top_p": 0.9,
-        }
+        # Tokenize input
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         
+        # Generate response
         outputs = model.generate(
-            model_inputs=model_inputs,
-            **generation_config
+            **inputs,
+            max_new_tokens=256,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
         )
         
+        # Decode and clean up 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
@@ -90,12 +77,12 @@ def analyze_input(image_data, question):
 demo = gr.Interface(
     fn=analyze_input,
     inputs=[
-        gr.Image(type="numpy", label="Medical Image"),
+        gr.Image(type="numpy", label="Medical Image (Optional)"),
         gr.Textbox(label="Question", placeholder="Enter your medical query...")
     ],
     outputs=gr.JSON(label="Analysis"),
-    title="Bio-Medical MultiModal Analysis",
-    description="Ask questions with or without an image",
+    title="Medical Query Analysis",
+    description="Ask medical questions with or without images. For general medical queries, no image is needed.",
     flagging_mode="never"
 )