Условие

Археологи нашли книгу на неизвестном языке. У них есть ограниченный набор переводов фрагментов на русский — это train_data.json (пары source → translation). Осталось много непереведённых фрагментов — это test_data.json (только source).

Задача: дообучить open-source модель (например, T5-small) и сгенерировать submission.json с переводами тестовых фрагментов. Ограничение по метрике BLEU ≥ 20 на тестовом наборе.

Решение

Подход

Классический seq2seq fine-tuning. На малой выборке T5-small достаточно, более крупные модели нужны только если данных много.

Шаги:

Загрузить train_data.json, разбить на train/val (80/20).
Токенизация, формат T5: "translate <unknown> to ru: <source>".
Fine-tune через HuggingFace Trainer или Seq2SeqTrainer.
Генерация на test, beam search с num_beams=4, length penalty.
Сохранение в нужном формате.

Реализация

import json
from datasets import Dataset
from transformers import (
    T5Tokenizer, T5ForConditionalGeneration,
    Seq2SeqTrainer, Seq2SeqTrainingArguments,
    DataCollatorForSeq2Seq,
)
import evaluate
import numpy as np
 
# 1. Данные
with open("train_data.json", encoding="utf-8") as f:
    train_pairs = json.load(f)
with open("test_data.json", encoding="utf-8") as f:
    test_items = json.load(f)
 
ds = Dataset.from_list(train_pairs).train_test_split(test_size=0.1, seed=42)
 
model_name = "google/t5-small"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)
 
 
def preprocess(batch):
    inputs = ["translate to ru: " + s for s in batch["source"]]
    model_inputs = tokenizer(inputs, max_length=256, truncation=True)
    labels = tokenizer(batch["translation"], max_length=256, truncation=True)
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs
 
 
tokenized = ds.map(preprocess, batched=True)
collator = DataCollatorForSeq2Seq(tokenizer, model=model)
 
bleu = evaluate.load("sacrebleu")
 
 
def compute_metrics(eval_preds):
    preds, labels = eval_preds
    if isinstance(preds, tuple):
        preds = preds[0]
    decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
    labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
    return {"bleu": bleu.compute(predictions=decoded_preds, references=[[r] for r in decoded_labels])["score"]}
 
 
args = Seq2SeqTrainingArguments(
    output_dir="t5_lost",
    num_train_epochs=20,
    learning_rate=5e-4,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    eval_strategy="epoch",
    save_strategy="epoch",
    predict_with_generate=True,
    generation_max_length=256,
    generation_num_beams=4,
    weight_decay=0.01,
    logging_steps=50,
    fp16=True,
    load_best_model_at_end=True,
    metric_for_best_model="bleu",
    greater_is_better=True,
)
 
trainer = Seq2SeqTrainer(
    model=model, args=args,
    train_dataset=tokenized["train"], eval_dataset=tokenized["test"],
    tokenizer=tokenizer, data_collator=collator,
    compute_metrics=compute_metrics,
)
trainer.train()
 
# 2. Inference
def translate(text: str) -> str:
    enc = tokenizer("translate to ru: " + text, return_tensors="pt", truncation=True).to(model.device)
    out = model.generate(**enc, max_length=256, num_beams=4, length_penalty=1.0)
    return tokenizer.decode(out[0], skip_special_tokens=True)
 
 
submission = [{"source": x["source"], "translation": translate(x["source"])} for x in test_items]
with open("submission.json", "w", encoding="utf-8") as f:
    json.dump(submission, f, ensure_ascii=False, indent=2)

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

Токенайзер T5 не знает иероглифов нового языка. Если входные символы редкие — T5 начнёт декодировать <unk>. Решение: расширить словарь (tokenizer.add_tokens([...]) + model.resize_token_embeddings) либо использовать байтовый токенайзер (ByT5).
Малый train. Без аугментации (back-translation, добавление шума) на 100–500 парах модель плохо обобщается; лучше использовать пред-обученный multilingual checkpoint (mBART, NLLB-200).
BLEU чувствителен к токенизации. Используйте sacrebleu (стандарт), не «самописный» BLEU.
«Запиши ответы в файл» — условие требует код, генерирующий перевод, а не зашитые строки. Нельзя инферить локально и подсунуть готовый JSON.
dtype="auto" — условие просит, иначе FP16 на CPU упадёт.
Длина последовательности. T5-small имеет лимит 512 токенов; для длинных отрывков — резать на предложения.

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

mBART-50 / NLLB-200 уже знают сотни языков — просто fine-tune.
Few-shot LLM (например, через transformers.pipeline("text-generation")) — без обучения, но менее стабильно.

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

T5-small (или ByT5/mBART для лучшего покрытия символов) + fine-tune на парах train_data.json с BLEU как метрикой выбора модели. Beam search декодирование. На простой задаче перевода BLEU ≥ 20 достижим за 10–20 эпох на T5-small с расширенным словарём.

Яндекс ML — Утерянный язык: машинный перевод с few-shot fine-tuning