#!/usr/bin/env python
# coding=utf-8
# Copyright 2024 Statistics and Machine Learning Research Group. All rights reserved.
from .base import StringFormatter, TemplateComponent, ConversationTemplate


# Jinja template format for Ministral (Nemotron-Diffusion-Exp-Ministral-8B)
# Based on chat_template.jinja
MINISTRAL_TEMPLATE = (
    "{{- '<SPECIAL_10>System' }}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'system' %}"
    "{{- '\\n' + message['content'].strip() }}"
    "{%- endif %}"
    "{%- endfor %}"
    "{{- '\\n' }}"
    "{%- for message in messages %}"
    "{%- if message['role'] == 'user' %}"
    "{{- '\\n<SPECIAL_11>User\\n' + message['content'].strip() + '\\n<SPECIAL_11>Assistant\\n' }}"
    "{%- elif message['role'] == 'assistant' %}"
    "{% generation %}"
    "{{- message['content'].strip() }}"
    "{% endgeneration %}"
    "{%- endif %}"
    "{%- endfor %}"
)


# Legacy ConversationTemplate format (kept for backward compatibility)
MINISTRAL_LEGACY_TEMPLATE = ConversationTemplate(
    template_name='ministral_legacy',
    user_formatter=StringFormatter(
        template=[
            TemplateComponent(type='string', content='\n<SPECIAL_11>User\n{{content}}\n<SPECIAL_11>Assistant\n')
        ]
    ),
    assistant_formatter=StringFormatter(
        template=[
            TemplateComponent(type='string', content='{{content}}')
        ]
    ),
    system_formatter=StringFormatter(
        template=[
            TemplateComponent(type='string', content='<SPECIAL_10>System\n{{content}}\n')
        ]
    ),
    force_system=True,
    system_default=''
)
