Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 from PIL import Image
 import gradio as gr
 import base64
@@ -18,7 +18,7 @@ bnb_config = BitsAndBytesConfig(
 )
 
 # Load model with revision pinning
-model = AutoModelForCausalLM.from_pretrained(
+model = AutoModel.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
     quantization_config=bnb_config,
     device_map="auto",
@@ -44,25 +44,22 @@ def analyze_input(image_data, question):
             prompt = f"Medical question: {question}\nAnswer: "
 
         # Tokenize input
-        input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+        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
+            "pixel_values": None,  # Set to None for text-only queries
+            "tgt_sizes": [tgt_size]  # Add target size for generation
         }
         
         # Generate response
-        generation_config = {
-            "max_new_tokens": 256,
-            "do_sample": True,
-            "temperature": 0.7,
-            "top_p": 0.9,
-        }
-        
         outputs = model.generate(
             model_inputs=model_inputs,
-            **generation_config
         )
         
         # Decode and clean up response