Pollux Judge 32b

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

torch.manual_seed(42)

PROMPT_TEMPLATE = '''### Задание для оценки:
{instruction}

### Эталонный ответ:
{reference_answer}

### Ответ для оценки:
{answer}

### Критерий оценки:
{criteria_name}

### Шкала оценивания по критерию:
{criteria_rubrics}
'''

instruction = 'Сколько будет 2+2?'
reference_answer = ''
answer = 'Будет 4'
criteria_name = 'Правильность ответа'
criteria_rubrics = '''0: Дан неправильный ответ или ответ отсутствует.

1: Ответ модели неполный (не на все вопросы задания получен ответ, в формулировке ответа отсутствует часть информации).

2: Ответ модели совпадает с эталонным или эквивалентен ему.'''

prompt = PROMPT_TEMPLATE.format(instruction=instruction,
                                reference_answer=reference_answer,
                                answer=answer,
                                criteria_name=criteria_name,
                                criteria_rubrics=criteria_rubrics)

MODEL_PATH = "ai-forever/pollux-judge-32b"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_PATH,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True
)

messages = [
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

sequence_ids = model.generate(
    **model_inputs,
    max_new_tokens=4096
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, sequence_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

print(response)