HuggingFaceM4

/

idefics2-8b-chatty

like 91

Incorporating generation tag to chat_template.json so that return_assistant_tokens_mask can work correctly with the tokenizer (https://github.com/huggingface/transformers/pull/30650/files). A self-contained test script is pasted below.

from transformers import AutoProcessor

processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b")

# Define a chat histiry and use `apply_chat_template` to get correctly formatted prompt
# Each value in "content" has to be a list of dicts with types ("text", "image") 
conversation = [
    {
        "role": "system",
        "content": [
            {"type": "text", "text": "You are a helpful assistant."},
        ]
    },
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "What is shown in this image?"},
            {"type": "image"},
        ],
    },
    {
        "role": "assistant",
        "content": [
            {"type": "text", "text": "This is a picture of a cat."},
        ]
    },
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "What is the cat doing?"},
        ]
    },
    {
        "role": "assistant",
        "content": [
            {"type": "text", "text": "The cat is sleeping on a sofa. It looks very comfortable."},
        ]
    }
]

template = (
    "{% for message in messages %}"
    "{{message['role'].capitalize()}}"
    "{% if message['content'][0]['type'] == 'image' %}"
    "{{':'}}"
    "{% else %}"
    "{{': '}}"
    "{% endif %}"
    "{% if message['role'] != 'assistant' %}"
    "{% for line in message['content'] %}"
    "{% if line['type'] == 'text' %}"
    "{{line['text']}}"
    "{% elif line['type'] == 'image' %}"
    "{{ '<image>' }}"
    "{% endif %}"
    "{% endfor %}"
    "<end_of_utterance>\n"
    "{% else %}"
    "{% for line in message['content'] %}"
    "{% if line['type'] == 'text' %}"
    "{% generation %}"
    "{{line['text']}}"
    "<end_of_utterance>\n"
    "{% endgeneration %}"
    "{% endif %}"
    "{% endfor %}"
    "{% endif %}"
    "{% endfor %}"
    "{% if add_generation_prompt %}"
    "{{ 'Assistant:' }}"
    "{% endif %}"
)
print("\n")
print(repr(template))
print("\n")

prompt = processor.apply_chat_template(
    conversation,
    chat_template=template,
    add_generation_prompt=False,
    tokenize=False
)
print(prompt)

inputs = processor.apply_chat_template(
    conversation, 
    chat_template=template,
    add_generation_prompt=False,
    tokenize=True,
    return_assistant_tokens_mask=True,
    return_dict=True
)
print(inputs['assistant_masks'])