File size: 1,231 Bytes
7388f4a 9e55c26 7388f4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 | 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
|