Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig, LlamaTokenizer
 import gradio as gr
 from PIL import Image
 from torchvision.transforms import ToTensor
@@ -16,26 +16,29 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Load the tokenizer first
-tokenizer = AutoTokenizer.from_pretrained(
-    "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
-    trust_remote_code=True,
-    token=api_token
-)
+# Initialize tokenizer using LlamaTokenizer specifically
+model_name = "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1"
+try:
+    tokenizer = LlamaTokenizer.from_pretrained(
+        model_name,
+        trust_remote_code=True,
+        token=api_token
+    )
+except Exception as e:
+    print(f"Failed to load LlamaTokenizer, falling back to AutoTokenizer: {e}")
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        trust_remote_code=True,
+        token=api_token
+    )
 
-# 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)
+# Explicitly set special tokens
+tokenizer.pad_token = tokenizer.eos_token = "</s>"
+tokenizer.pad_token_id = tokenizer.eos_token_id = 2  # Common EOS token ID for Llama models
 
-# Load the model with updated tokenizer
+# Load the model
 model = AutoModel.from_pretrained(
-    "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
+    model_name,
     quantization_config=bnb_config,
     device_map="auto",
     torch_dtype=torch.float16,
@@ -43,9 +46,15 @@ model = AutoModel.from_pretrained(
     token=api_token
 )
 
-# Update model's generation config
+# Ensure the model's generation config is properly set
+if not hasattr(model, 'generation_config'):
+    from transformers import GenerationConfig
+    model.generation_config = GenerationConfig()
+
 model.generation_config.eos_token_id = tokenizer.eos_token_id
 model.generation_config.pad_token_id = tokenizer.pad_token_id
+model.config.eos_token_id = tokenizer.eos_token_id
+model.config.pad_token_id = tokenizer.pad_token_id
 
 # Preprocess image
 def preprocess_image(image):
@@ -66,16 +75,22 @@ def analyze_input(image, question):
         input_ids = tokenized.input_ids.to(model.device)
         
         # Calculate target size
-        tgt_size = input_ids.size(1) + 256  # Original input size + max new tokens
+        tgt_size = input_ids.size(1) + 256
         
         # Construct the model_inputs dictionary
         model_inputs = {
             "input_ids": input_ids,
             "pixel_values": pixel_values,
-            "tgt_sizes": [tgt_size]  # Add target sizes for generation
+            "tgt_sizes": [tgt_size],
+            "pad_token_id": tokenizer.pad_token_id,
+            "eos_token_id": tokenizer.eos_token_id
         }
         
-        # Generate the response - Note the changed calling convention
+        # Print debugging information
+        print(f"Token IDs - EOS: {tokenizer.eos_token_id}, PAD: {tokenizer.pad_token_id}")
+        print(f"Model config - EOS: {model.config.eos_token_id}, PAD: {model.config.pad_token_id}")
+        
+        # Generate the response
         outputs = model.generate(model_inputs)
         
         # Decode the response
@@ -83,6 +98,8 @@ def analyze_input(image, question):
         return {"status": "success", "response": response}
     
     except Exception as e:
+        import traceback
+        print(f"Error details: {traceback.format_exc()}")
         return {"status": "error", "message": str(e)}
 
 # Create Gradio interface
@@ -98,8 +115,9 @@ demo = gr.Interface(
 )
 
 # Launch the Gradio app
-demo.launch(
-    share=True,
-    server_name="0.0.0.0",
-    server_port=7860
-)
+if __name__ == "__main__":
+    demo.launch(
+        share=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )