Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, GenerationConfig
+from transformers import AutoModel, AutoTokenizer, BitsAndBytesConfig
 import gradio as gr
 from PIL import Image
 from torchvision.transforms import ToTensor
@@ -16,102 +16,49 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Model name
-model_name = "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1"
-
-# 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 = AutoModelForCausalLM.from_pretrained(
-    model_name,
+# Initialize model and 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,
-    generation_config=generation_config
+    attn_implementation="flash_attention_2",
+    token=api_token
 )
 
-# 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):
-    transform = ToTensor()
-    return transform(image).unsqueeze(0).to(model.device)
+tokenizer = AutoTokenizer.from_pretrained(
+    "ContactDoctor/Bio-Medical-MultiModal-Llama-3-8B-V1",
+    trust_remote_code=True,
+    token=api_token
+)
 
-# 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:
+            # Convert to RGB if image is provided
             image = image.convert('RGB')
-            pixel_values = preprocess_image(image)
-        
-        # Tokenize the question
-        inputs = tokenizer(
-            question,
-            return_tensors="pt",
-            padding=True,
-            truncation=True,
-            max_length=512
-        ).to(model.device)
         
-        # Add image if provided
-        if pixel_values is not None:
-            inputs['pixel_values'] = pixel_values
+        # Prepare messages in the format expected by the model
+        msgs = [{'role': 'user', 'content': [image, question]}]
         
-        # Generate response
-        outputs = model.generate(
-            **inputs,
-            generation_config=generation_config,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
+        # Generate response using the chat method
+        response_stream = model.chat(
+            image=image,
+            msgs=msgs,
+            tokenizer=tokenizer,
+            sampling=True,
+            temperature=0.95,
+            stream=True
         )
         
-        # Decode response
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return {"status": "success", "response": response}
+        # Collect the streamed response
+        generated_text = ""
+        for new_text in response_stream:
+            generated_text += new_text
+            print(new_text, flush=True, end='')
+        
+        return {"status": "success", "response": generated_text}
     
     except Exception as e:
         import traceback
@@ -123,12 +70,16 @@ def analyze_input(image, question):
 demo = gr.Interface(
     fn=analyze_input,
     inputs=[
-        gr.Image(type="pil", label="Upload Medical Image (Optional)"),
-        gr.Textbox(label="Medical Question")
+        gr.Image(type="pil", label="Upload Medical Image"),
+        gr.Textbox(
+            label="Medical Question",
+            placeholder="Give the modality, organ, analysis, abnormalities (if any), treatment (if abnormalities are present)?",
+            value="Give the modality, organ, analysis, abnormalities (if any), treatment (if abnormalities are present)?"
+        )
     ],
     outputs=gr.JSON(label="Analysis"),
-    title="ContactDoctor Medical Assistant",
-    description="Upload a medical image and/or enter a question to receive detailed AI-powered responses."
+    title="Medical Image Analysis Assistant",
+    description="Upload a medical image and ask questions about it. The AI will analyze the image and provide detailed responses."
 )
 
 # Launch the Gradio app