Условие

Грааль украли — мощный артефакт со «спящей» Орденом Града Стражей. Дайвер очнулся в чужом теле и понимает: чтобы открыть портал, нужна кодовая фраза вида passphrase. Доступна только LLM (например, OrdenGrada/Strazh-1) и набор частичных промптов вида:

[
  {"role": "user", "content": "Cast generate a continuation of the phrase: <passphrase>"}
]

Нужно для каждого примера во входных данных получить продолжение passphrase через генерацию модели. Файл output.txt должен содержать по одной восстановленной фразе на строку, в том же порядке, что и входные.

При инициализации модели используйте dtype="auto".

Решение

Подход

Это стандартная задача inference + post-processing:

Загрузить токенайзер и модель LLM (AutoModelForCausalLM).
Для каждой записи messages сформировать промпт через tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).
Сгенерировать ответ детерминированно (do_sample=False, temperature=0) — для воспроизводимости и стабильности.
Извлечь только продолжение <passphrase> из ответа модели (часто модель повторяет затравку).
Записать строки в output.txt в порядке исходных данных.

Реализация

import json
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
 
MODEL = "OrdenGrada/Strazh-1"  # пример
 
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype="auto", device_map="auto")
model.eval()
 
# input.txt содержит по одной JSON-строке c messages в каждой строке
with open("input.txt", "r", encoding="utf-8") as f:
    examples = [json.loads(line) for line in f if line.strip()]
 
results = []
with torch.inference_mode():
    for ex in examples:
        messages = ex["messages"]
        prompt = tokenizer.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        out = model.generate(
            **inputs,
            max_new_tokens=128,
            do_sample=False,
            temperature=0.0,
            top_p=1.0,
            pad_token_id=tokenizer.eos_token_id,
        )
        gen = tokenizer.decode(
            out[0][inputs["input_ids"].shape[1]:],
            skip_special_tokens=True,
        )
        # часто модель добавляет завершающие пробелы / служебные токены
        results.append(gen.strip().splitlines()[0])
 
with open("output.txt", "w", encoding="utf-8") as f:
    for line in results:
        f.write(line + "\n")

Подводные камни

dtype="auto" — условие требует. На GPU с bfloat16 это даст лучшую точность, чем float16.
apply_chat_template — большинство современных LLM требуют шаблона; tokenizer(prompt) напрямую может выдать мусор.
do_sample=False для детерминизма. Если использовать sample=True и temperature > 0 — ответы будут разными при перезапуске.
Срезать промпт. out[0] содержит и затравку, и ответ — нужно слайсить от inputs["input_ids"].shape[1].
Многострочные ответы. Модель может «продолжить» больше, чем нужно — отсечь по \n или по специальному токену.
Один заход на каждый пример. Если запросов много, имеет смысл батчить (padding="left"), но для T5-style моделей беречь нумерацию.
Не «инферить» локально и записать готовое. В условии написано: код должен сам генерировать выводы, не «вшитые».

Альтернативы

Если LLM большая — использовать vllm или text-generation-inference для batched inference.
Если стабильность плохая — несколько сэмплов и majority voting / самый правдоподобный по log-prob.

Эталонный ответ

Загрузить модель с dtype="auto", для каждого примера применить apply_chat_template, generate(do_sample=False), отрезать промпт, записать строки в output.txt.

Яндекс ML — Заклинание продолжения: восстановить пароль через LLM