How to do a special operation for the first element in a loop with jinja grammar?

Question

I'm trying to convert the following python code into a jinja template, so that I can use my own chat tempolate in llama tokenizer. But I'm new to this field and don't know what to do.

My python code looks like this:

for i in range(len(messages)):
    if i == 0:
        do_something()
    else:
        do_other_things()

I tried marking the first element with a variable, the value of which should be changed inside the loop. But in practice the variable stays constant, so my idea didin't work.

The chat template I actually tried:

{% set is_first = "false" %}
{% for message in messages %}
    {% if message['role'] == 'system' %}
        {{ '<<SYS>>\\n' + message['content'].strip() + '\\n<</SYS>>\\n\\n'}}
        {% set is_first = "true" %}
    {% elif message['role'] == 'user' and is_first == "true" %}
        {% set is_first = "false" %}
        {{ message['content'].strip() + ' [/INST]' }}
    {% elif message['role'] == 'user' %}
        {{ '[INST] ' + message['content'].strip() + ' [/INST]' }}
    {% elif message['role'] == 'assistant' %}
        {{ message['content'].strip() }}
    {% endif %}
{% endfor %}"""

Is there a solution to make my code work? By the way, is there any method to tell jinja raise an error when the order of dialog isn't in "system -> user -> assistant -> user -> ..." fashion?

Answer 1

Found a solution in this link. Use loop.index0 to access the index number in the loop.