Update test.py

test.py

@@ -1,10 +1,11 @@
-from transformers import AutoModelForVision2Seq, AutoProcessor 
+from transformers import AutoModelForImageTextToText, AutoProcessor
 
-model = AutoModelForVision2Seq.from_pretrained(
+model = AutoModelForImageTextToText.from_pretrained(
     "QuixiAI/Prisma-VL-8B",
-    torch_dtype="auto",
+    dtype="auto",
     device_map="auto"
 )
+
 processor = AutoProcessor.from_pretrained("QuixiAI/Prisma-VL-8B")
 
 messages = [
@@ -15,10 +16,17 @@ messages = [
                 "type": "image",
                 "image": "https://static.wikia.nocookie.net/essentialsdocs/images/7/70/Battle.png/revision/latest?cb=20220523172438",
             },
-            {"type": "text", "text": "Describe your thoughts and your experience of thinking.  The phenomenology is more important than the actual answer."},
+            {
+                "type": "text",
+                "text": (
+                    "Describe your thoughts and your experience of thinking. "
+                    "The phenomenology is more important than the actual answer."
+                ),
+            },
         ],
     }
 ]
+
 inputs = processor.apply_chat_template(
     messages,
     tokenize=True,
@@ -26,12 +34,18 @@ inputs = processor.apply_chat_template(
     return_dict=True,
     return_tensors="pt"
 )
+
 inputs = inputs.to(model.device)
+
 generated_ids = model.generate(**inputs, max_new_tokens=1280)
 generated_ids_trimmed = [
-    out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+    out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
 ]
+
 output_text = processor.batch_decode(
-    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
+    generated_ids_trimmed,
+    skip_special_tokens=True,
+    clean_up_tokenization_spaces=False
 )
+
 print(output_text)