update chat_template

README.md

@@ -171,18 +171,25 @@ model = AutoModelForCausalLM.from_pretrained(
 ).to(DEVICE).eval()
 processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)
 
-messages = [{"role": "user", "content": PROMPTS[CHOSEN_TASK]}]
-text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-inputs = processor(text=[text], images=[image], return_tensors="pt")
-inputs = {k: (v.to(DEVICE) if isinstance(v, torch.Tensor) else v) for k, v in inputs.items()}
-
-with torch.inference_mode():
-    generated = model.generate(**inputs, max_new_tokens=1024, do_sample=False, use_cache=True)
-
-resp = processor.batch_decode(generated, skip_special_tokens=True)[0]
-answer = resp.split(text)[-1].strip()
-print(answer)
+messages = [
+    {"role": "user",         
+     "content": [
+            {"type": "image", "image": image},
+            {"type": "text", "text": PROMPTS[CHOSEN_TASK]},
+        ]
+    }
+]
+inputs = processor.apply_chat_template(
+    messages, 
+    tokenize=True, 
+    add_generation_prompt=True, 	
+    return_dict=True,
+	return_tensors="pt"
+).to(DEVICE)
+
+outputs = model.generate(**inputs, max_new_tokens=1024)
+outputs = processor.batch_decode(outputs, skip_special_tokens=True)[0]
+print(outputs)
 ```
 
 ## Performance

chat_template.jinja

@@ -7,14 +7,38 @@
 {%- if not sep_token is defined -%}
     {%- set sep_token = "<|end_of_sentence|>" -%}
 {%- endif -%}
+{%- if not image_token is defined -%}
+    {%- set image_token = "<|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>" -%}
+{%- endif -%}
 {{- cls_token -}}
 {%- for message in messages -%}
     {%- if message["role"] == "user" -%}
-        {{- "User: <|IMAGE_START|><|IMAGE_PLACEHOLDER|><|IMAGE_END|>" + message["content"] + "\n" -}}
+        {{- "User: " -}}
+        {%- for content in message["content"] -%}
+            {%- if content["type"] == "image" -%}
+                {{ image_token }}
+            {%- endif -%}
+        {%- endfor -%}
+        {%- for content in message["content"] -%}
+            {%- if content["type"] == "text" -%}
+                {{ content["text"] }}
+            {%- endif -%}
+        {%- endfor -%}
+        {{ "\n" -}}
     {%- elif message["role"] == "assistant" -%}
-        {{- "Assistant: " + message["content"] + sep_token -}}
+        {{- "Assistant: " -}}
+        {%- for content in message["content"] -%}
+            {%- if content["type"] == "text" -%}
+                {{ content["text"] + "\n" }}
+            {%- endif -%}
+        {%- endfor -%}
+        {{ sep_token -}}
     {%- elif message["role"] == "system" -%}
-        {{- message["content"] -}}
+        {%- for content in message["content"] -%}
+            {%- if content["type"] == "text" -%}
+                {{ content["text"] + "\n" }}
+            {%- endif -%}
+        {%- endfor -%}
     {%- endif -%}
 {%- endfor -%}
 {%- if add_generation_prompt -%}