Fine-tuning Gemma 2B для структурированного извлечения JSON: полный разбор кода и 432 примера

Почему Gemma 2B плюет на ваш JSON и как это исправить

Вы кормите модель регуляторным документом на 20 страниц. Ждете красивый JSON с полями «название компании», «дата регистрации», «штрафы». А получаете... кусок текста, половину ответа на немецком или, что чаще, символ `}` в самом начале. Знакомая история? Особенно с маленькими моделями в 2-3 миллиарда параметров.

Вот и приходится либо использовать огромные и дорогие модели, либо мириться с точностью в 70-80%. Первый вариант бьет по бюджету, второй — по нервам, потому что каждый сломанный JSON надо чинить вручную. Мой путь — fine-tuning. Не на 10 примерах, а на 432. Результат — точность валидного и семантически верного извлечения подскакивает до 94%. И да, это работает на Gemma 2B в 2026 году, когда вокруг уже летают 100B-параметровые монстры, но считать каждый гигабайт видеопамяти все еще актуально.

432 примера: что внутри датасета, который не сломает модель

Секрет не в количестве, а в структуре и разнообразии. Я не просто скопировал 432 PDF-ки в txt. Каждый пример — это пара: исходный текст (промпт) и эталонный JSON (ответ).

• Типы документов: выписки из ЕГРЮЛ, постановления о штрафах, налоговые уведомления, лицензионные соглашения. Все на русском, с типичными бюрократическими оборотами.
• Вариативность сложности: от простых одностраничных справок до многостраничных документов с таблицами и сносками.
• Атака на слабые места: я специально включал документы с пропущенными полями, нестандартными формулировками, цифрами в тексте. Если в документе нет «ИНН», в JSON должно быть `null` или поле отсутствовать — модель должна это понять, а не гадать.

Формат для обучения прост: текст промпта, затем директива `\n###\n`, затем чистый JSON. Никаких лишних слов вроде «Вот JSON:». Модель должна привыкнуть, что после `###` она переключается в режим генерации структуры.

1 Готовим инструменты: что актуально в апреле 2026

Забудьте про туториалы 2024 года. Библиотеки обновились, методы стали эффективнее. Вот стек, который работает сейчас:

# Актуальные версии на 05.04.2026
# transformers >= 4.45.0 (поддержка новых моделей Gemma)
# peft >= 0.11.0 (последние улучшения LoRA)
# torch >= 2.4.0 (оптимизации CUDA 12.6)
# datasets >= 2.20.0
# accelerate >= 0.30.0
# jinja2 >= 3.1.3 (для шаблонов)

# Установка одной командой (для чистого окружения):
# pip install transformers peft torch datasets accelerate jinja2 --upgrade

Модель: `google/gemma-2-2b-it`. Почему Instruct-версия? Она уже немного адаптирована к диалогу, и fine-tuning на структуру ложится лучше. Базовая Gemma 2B тоже сработает, но придется дольше учить.

2 Код: загрузка данных и токенизация без головной боли

Первая ошибка — токенизировать промпт и JSON вместе, как обычный текст. Так вы собьете модель с толку. Нужно явно разделить их и обучить только на генерации JSON-части. Используем чатовый шаблон.

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from datasets import Dataset
import json

model_name = "google/gemma-2-2b-it"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token  # Важно для обучения

# Загружаем наш датасет (предполагаем, что он в JSONL формате)
def load_dataset(path):
    data = []
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            item = json.loads(line)
            # item = {"prompt": "текст документа", "completion": "{\n...}"}
            data.append(item)
    return Dataset.from_list(data)

dataset = load_dataset("regulatory_docs_432.jsonl")

def tokenize_function(example):
    # Формируем промпт в формате чата
    messages = [
        {"role": "user", "content": example["prompt"] + "\n###\n"},
        {"role": "assistant", "content": example["completion"]}
    ]
    # Применяем чатовый шаблон токенизатора Gemma
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
    # Токенизируем всю последовательность
    tokenized = tokenizer(text, truncation=True, max_length=2048, padding="max_length")
    # Создаем маски labels: -100 для промпта, реальные id для ответа (JSON)
    # Найдем, где начинается ответ ассистента
    tokens = tokenizer.convert_ids_to_tokens(tokenized["input_ids"])
    # Простой поиск по токену начала ответа (зависит от шаблона)
    # Для Gemma обычно это `model`
    assistant_start = tokens.index("model") if "model" in tokens else -1
    labels = [-100] * len(tokenized["input_ids"])
    if assistant_start != -1:
        labels[assistant_start:] = tokenized["input_ids"][assistant_start:]
    tokenized["labels"] = labels
    return tokenized

tokenized_dataset = dataset.map(tokenize_function, batched=False)

Если кажется сложным — это потому, что так и есть. Но альтернатива — модель будет учиться пересказывать промпт, а не генерировать JSON. Кстати, если ваш токенизатор криво находит границу, посмотрите туториал по настройке сэмплеров для JSON, там похожие принципы.

3 Fine-tuning с LoRA: настраиваем только то, что нужно

Полный fine-tuning 2B-модели — избыточен. LoRA (Low-Rank Adaptation) в 2026 году — стандарт. Но не всякая LoRA одинакова. Параметры по умолчанию могут дать плохой результат. Вот конфигурация, которая работает для структурированного вывода:

from peft import LoraConfig, get_peft_model, TaskType

lora_config = LoraConfig(
    r=32,                  # Ранг. Не 8, не 64. 32 — золотая середина для задачи структуры.
    lora_alpha=64,         # Коэффициент масштабирования.
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],  # Задействуем все линейные слои.
    lora_dropout=0.1,
    bias="none",
    task_type=TaskType.CAUSAL_LM,
)

model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()  # Должно показать ~0.5-1% обучаемых параметров

4 Цикл обучения: 3 эпохи, теплый старт и взвешенная потеря

Больше 3 эпох на 432 примерах — почти гарантированное переобучение. Используем оптимизатор AdamW с cosine расписанием и warmup.

from transformers import TrainingArguments, Trainer

training_args = TrainingArguments(
    output_dir="./gemma-2b-json-lora",
    num_train_epochs=3,
    per_device_train_batch_size=2,  # Зависит от памяти. Для 16 ГБ хватит.
    gradient_accumulation_steps=4,
    warmup_steps=50,
    logging_steps=10,
    save_strategy="epoch",
    learning_rate=3e-4,             # Не 5e-5! Для LoRA и JSON нужен более агрессивный LR.
    fp16=False,                     # Используем bfloat16, если карта поддерживает.
    bf16=True,                      # Предпочтительнее для Gemma.
    gradient_checkpointing=True,
    remove_unused_columns=False,    # Важно, чтобы не потерять labels.
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=lambda data: {"input_ids": torch.stack([d["input_ids"] for d in data]),
                               "attention_mask": torch.stack([d["attention_mask"] for d in data]),
                               "labels": torch.stack([d["labels"] for d in data])},
)

trainer.train()
model.save_pretrained("./gemma-2b-json-lora-final")
tokenizer.save_pretrained("./gemma-2b-json-lora-final")

Обучение на одной RTX 5070 Ti (16 ГБ) займет около 1.5 часа. Если хочется быстрее, можно попробовать методы ускорения из фиинтюнов для Gemma 3, но для нашей задачи это оверкилл.

Что сломалось и как починить: 4 самые частые ошибки

Даже с идеальным кодом что-то пойдет не так. Вот что я видел чаще всего:

Вопросы, которые вы хотели задать, но боялись

И последнее. Не надейтесь, что fine-tuning раз и навсегда решит проблему с JSON. Форматы документов меняются, появляются новые поля. Модель нужно дообучать. Но теперь у вас есть план и 432 примера, чтобы начать. Дальше — только вперед, собирать свой датасет и калибровать под свою задачу. Удачи.

Подписаться на канал