Update app.py

app.py

@@ -17,7 +17,7 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load model with revision pinning - using CausalLM for text generation
+# Load model with revision pinning
 model = AutoModelForCausalLM.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
     quantization_config=bnb_config,
@@ -44,16 +44,25 @@ def analyze_input(image_data, question):
             prompt = f"Medical question: {question}\nAnswer: "
 
         # Tokenize input
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+        
+        # Prepare model inputs
+        model_inputs = {
+            "input_ids": input_ids,
+            "pixel_values": None  # Set to None for text-only queries
+        }
         
         # Generate response
+        generation_config = {
+            "max_new_tokens": 256,
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+        }
+        
         outputs = model.generate(
-            **inputs,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
-            pad_token_id=tokenizer.eos_token_id
+            model_inputs=model_inputs,
+            **generation_config
         )
         
         # Decode and clean up response
@@ -82,7 +91,7 @@ demo = gr.Interface(
     ],
     outputs=gr.JSON(label="Analysis"),
     title="Medical Query Analysis",
-    description="Ask medical questions with or without images. For general medical queries, no image is needed.",
+    description="Ask medical questions. For now, please focus on text-based queries without images.",
     flagging_mode="never"
 )