| 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 |
|
|
