Update idefics user assistant tags

kushal-10 · Apr 29, 2024 · 71257c1 · 71257c1
1 parent 9ddeb69
commit 71257c1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/backends/huggingface_multimodal_api.py b/backends/huggingface_multimodal_api.py
@@ -121,18 +121,18 @@ def generate_idefics_output(messages: list[Dict], model: IdeficsForVisionText2Te
     idefics_input = [] 
     for m in messages:
         if m['role'] == 'user':
-            idefics_input.append('\nUSER: ' + m['content'])
+            idefics_input.append('\nUser: ' + m['content'])
             if 'image' in m.keys():
                 idefics_input.append(m['image'])
             idefics_input.append('<end_of_utterance>')
         elif m['role'] == 'assistant':
-            idefics_input.append('\nASSISTANT: ' + m['content'])        
+            idefics_input.append('\nAssistant: ' + m['content'])        
             idefics_input.append('<end_of_utterance>')    
-    idefics_input.append('\nASSISTANT:')  
+    idefics_input.append('\nAssistant:')  
     idefics_input = [idefics_input]
 
     inputs = processor(idefics_input, add_end_of_utterance_token=False, return_tensors="pt").to(device)
-    
+
     # Generation args for Idefics
     exit_condition = processor.tokenizer("<end_of_utterance>", add_special_tokens=False).input_ids
     bad_words_ids = processor.tokenizer(["<image>", "<fake_token_around_image>"], add_special_tokens=False).input_ids

diff --git a/backends/model_registry.json b/backends/model_registry.json
@@ -435,7 +435,7 @@
     "backend": "huggingface_multimodal",
     "huggingface_id": "HuggingFaceM4/idefics-80b-instruct",
     "model_type": "Idefics",
-    "eos_to_cull": "ASSISTANT:",
+    "eos_to_cull": "Assistant:",
     "padding": false,
     "premade_chat_template": false,
     "custom_chat_template": "{%- for message in messages -%}{% if message['role'] == 'user' %}USER: {{ message['content'] }} <end_of_utterance>\n{% else %}ASSISTANT: {{ message['content'] }} <end_of_utterance>\n{% endif %}{% endfor %}ASSISTANT:"