import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
)


trained_model_name = "culturalheritagenus/Jawi-OCR-Qwen-v2"

model = AutoModelForCausalLM.from_pretrained(
    trained_model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="eager",
)

tokenizer = AutoTokenizer.from_pretrained(trained_model_name)

chat_template = """{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}
{% if message['role'] == 'assistant' %}<start_of_turn>{{ role }}
{% generation %}{{ message['content'] | trim }}{% endgeneration %}
<end_of_turn>{% else %}
{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>' }}
{% endif %}
{% endfor %}
{% if add_generation_prompt %}<start_of_turn>model
{% generation %}{% endgeneration %}
{% endif %}
"""

tokenizer.chat_template = chat_template