Update app.py

app.py

@@ -16,28 +16,36 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load the model and tokenizer
-model = AutoModel.from_pretrained(
+# Load the tokenizer first
+tokenizer = AutoTokenizer.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
-    quantization_config=bnb_config,
-    device_map="auto",
-    torch_dtype=torch.float16,
     trust_remote_code=True,
     token=api_token
 )
 
-tokenizer = AutoTokenizer.from_pretrained(
+# Set default tokens if they're missing
+if tokenizer.eos_token is None:
+    tokenizer.eos_token = "</s>"
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    
+# Update the tokenizer's token IDs
+tokenizer.pad_token_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
+tokenizer.eos_token_id = tokenizer.convert_tokens_to_ids(tokenizer.eos_token)
+
+# Load the model with updated tokenizer
+model = AutoModel.from_pretrained(
     "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
+    quantization_config=bnb_config,
+    device_map="auto",
+    torch_dtype=torch.float16,
     trust_remote_code=True,
     token=api_token
 )
 
-# Ensure `eos_token_id` is properly set
-eos_token_id = tokenizer.eos_token_id or tokenizer.pad_token_id
-if eos_token_id is None:
-    raise ValueError(
-        "Neither `eos_token_id` nor `pad_token_id` is defined in the tokenizer. Please specify one explicitly."
-    )
+# Update model's generation config
+model.generation_config.eos_token_id = tokenizer.eos_token_id
+model.generation_config.pad_token_id = tokenizer.pad_token_id
 
 # Preprocess image
 def preprocess_image(image):
@@ -47,37 +55,33 @@ def preprocess_image(image):
 # Handle queries
 def analyze_input(image, question):
     try:
-        # Prepare pixel values for image input
+        # Process the image if provided
         pixel_values = None
         if image is not None:
             image = image.convert('RGB')
             pixel_values = preprocess_image(image)
-
+        
         # Tokenize the question
         tokenized = tokenizer(question, return_tensors="pt")
         input_ids = tokenized.input_ids.to(model.device)
-
-        # Calculate target size for generation
-        tgt_size = input_ids.size(1) + 256  # Input size + max new tokens
-
+        
+        # Calculate target size
+        tgt_size = input_ids.size(1) + 256  # Original input size + max new tokens
+        
         # Construct the model_inputs dictionary
         model_inputs = {
             "input_ids": input_ids,
             "pixel_values": pixel_values,
-            "tgt_sizes": [tgt_size]
+            "tgt_sizes": [tgt_size]  # Add target sizes for generation
         }
-
-        # Generate response
-        outputs = model.generate(
-            model_inputs=model_inputs,  # Explicitly pass the model_inputs
-            max_new_tokens=256,
-            eos_token_id=eos_token_id
-        )
-
+        
+        # Generate the response - Note the changed calling convention
+        outputs = model.generate(model_inputs)
+        
         # Decode the response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return {"status": "success", "response": response}
-
+    
     except Exception as e:
         return {"status": "error", "message": str(e)}
 
@@ -98,4 +102,4 @@ demo.launch(
     share=True,
     server_name="0.0.0.0",
     server_port=7860
-)
+)