Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig, LlamaTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig
 import gradio as gr
 from PIL import Image
 from torchvision.transforms import ToTensor
@@ -16,45 +16,55 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Initialize tokenizer using LlamaTokenizer specifically
+# Model name
 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
-    )
 
-# 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
+# Initialize tokenizer
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    trust_remote_code=True,
+    token=api_token
+)
+
+# Set up tokenizer with default tokens
+default_tokens = {
+    "pad_token": "[PAD]",
+    "eos_token": "</s>",
+    "bos_token": "<s>",
+    "unk_token": "<unk>",
+}
+
+for token_name, token_value in default_tokens.items():
+    if getattr(tokenizer, token_name) is None:
+        setattr(tokenizer, token_name, token_value)
+        token_id_name = f"{token_name}_id"
+        if getattr(tokenizer, token_id_name) is None:
+            token_id = tokenizer.convert_tokens_to_ids(token_value)
+            setattr(tokenizer, token_id_name, token_id)
+
+# Create generation config
+generation_config = GenerationConfig(
+    pad_token_id=tokenizer.pad_token_id,
+    eos_token_id=tokenizer.eos_token_id,
+    bos_token_id=tokenizer.bos_token_id,
+    max_new_tokens=256,
+)
 
 # Load the model
-model = AutoModel.from_pretrained(
+model = AutoModelForCausalLM.from_pretrained(
     model_name,
     quantization_config=bnb_config,
     device_map="auto",
     torch_dtype=torch.float16,
     trust_remote_code=True,
-    token=api_token
+    token=api_token,
+    generation_config=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
+# Ensure model configs are set
 model.config.pad_token_id = tokenizer.pad_token_id
+model.config.eos_token_id = tokenizer.eos_token_id
+model.config.bos_token_id = tokenizer.bos_token_id
 
 # Preprocess image
 def preprocess_image(image):
@@ -64,6 +74,12 @@ def preprocess_image(image):
 # Handle queries
 def analyze_input(image, question):
     try:
+        # Debug print
+        print(f"Tokenizer config:")
+        print(f"EOS token: {tokenizer.eos_token} (id: {tokenizer.eos_token_id})")
+        print(f"PAD token: {tokenizer.pad_token} (id: {tokenizer.pad_token_id})")
+        print(f"BOS token: {tokenizer.bos_token} (id: {tokenizer.bos_token_id})")
+        
         # Process the image if provided
         pixel_values = None
         if image is not None:
@@ -71,35 +87,36 @@ def analyze_input(image, question):
             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
-        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],
-            "pad_token_id": tokenizer.pad_token_id,
-            "eos_token_id": tokenizer.eos_token_id
-        }
+        inputs = tokenizer(
+            question,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=512
+        ).to(model.device)
         
-        # 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}")
+        # Add image if provided
+        if pixel_values is not None:
+            inputs['pixel_values'] = pixel_values
         
-        # Generate the response
-        outputs = model.generate(model_inputs)
+        # Generate response
+        outputs = model.generate(
+            **inputs,
+            generation_config=generation_config,
+            max_new_tokens=256,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+        )
         
-        # Decode the response
+        # Decode response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return {"status": "success", "response": response}
     
     except Exception as e:
         import traceback
-        print(f"Error details: {traceback.format_exc()}")
+        error_trace = traceback.format_exc()
+        print(f"Error occurred: {error_trace}")
         return {"status": "error", "message": str(e)}
 
 # Create Gradio interface