Update app.py

app.py

@@ -1,5 +1,6 @@
-from transformers import RobertaTokenizer, RobertaForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+
 # Define the model path
 MODEL_NAME = "RPW/NIH-1.2_Llama-3.2-11B-Vision-Instruct"
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -7,16 +8,26 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 instruction = "You are an expert radiographer. Describe accurately what you see in this image."
 
 # Load the tokenizer and model
-tokenizer = RobertaTokenizer.from_pretrained(MODEL_NAME)
-model = RobertaForSequenceClassification.from_pretrained(MODEL_NAME).to(device)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME).to(device)
+
+# Example of messages
+messages = [{"role": "system", "content": instruction}]
 
 input_text = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
 inputs = tokenizer(
     image,
     input_text,
     add_special_tokens=False,
-    return_tensors="pt").to("cuda")
+    return_tensors="pt"
+).to(device)
 
 text_streamer = TextStreamer(tokenizer, skip_prompt=True)
-_ = model.generate(**inputs, streamer=text_streamer, max_new_tokens=128,
-                   use_cache=True, temperature=1.5, min_p=0.1)
+_ = model.generate(
+    **inputs,
+    streamer=text_streamer,
+    max_new_tokens=128,
+    use_cache=True,
+    temperature=1.5,
+    min_p=0.1
+)