От учителя к студенту: путь к изотропным эмбеддингам в 29M модели. Часть 3/10

Важное предупреждение:
Весь код по моей методологии написан нейросетью DeepSeek, сам бы не осилил (я не программист).
За мной лишь отладка, исправление неточностей и ошибок.

3.1. Выбор датасета для contrastive

deepvk/ru-HNP — русский датасет hard negatives, сгенерированный с помощью ChatGPT (см. карточку deepvk/ru-HNP). Отлично подходит для нашей задачи, потому что даёт готовую разметку без необходимости что-то придумывать самостоятельно.

Структура данных для одного запроса:

• 1 query — исходное предложение.

• 5 positives — перефразировки.

• 5 hard negatives — семантически близкие, но не парафразы.

Пример (из датасета):

query:      С 1952 года после окончания Куйбышевского индустриального
            института им. В. В. Куйбышева работал в системе атомной
            промышленности СССР.

positives:  После окончания Куйбышевского индустриального института
            им. В. В. Куйбышева он начал работать в системе атомной
            промышленности СССР с 1952 года.
            С 1952 года после выпуска из Куйбышевского индустриального
            института им. В. В. Куйбышева он трудился в системе атомной
            промышленности СССР.
            ...

negatives:  До 1952 года перед окончанием Куйбышевского
            индустриального института им. В. В. Куйбышева он не работал
            в системе атомной промышленности СССР.
            После окончания Куйбышевского индустриального института
            им. В. В. Куйбышева он не устроился на работу в систему
            атомной промышленности СССР с 1952 года.
            ...

Почему это подходит:

• Positives — качественные парафразы с разной степенью перефразирования.

• Negatives — осмысленные: отрицания, временные сдвиги («до/после»), замены глаголов на противоположные.

• Объём: 100K запросов → 500K позитивов + 500K негативов в развёрнутом виде.

Размеры сплитов:

• train: 100,000

• validation: 590

• test: 2,000

3.2. Формат данных для обучения

Динамический сэмплинг (выборка) вместо фиксированных комбинаций. У нас 5 positives × 5 negatives = 25 возможных троек (1 pos + 1 neg). Использование 1 pos и 2 neg даёт 50 возможных наборов на каждый запрос (5 способов выбрать pos × 10 способов выбрать 2 neg из 5); если хранить все развёрнуто, получится 5M четвёрок (query, pos, neg1, neg2) вместо 2.5M троек (query, pos, neg).

Лучше — хранить полные списки (query, pos[], neg[]) и на каждой эпохе случайно выбирать:

• 1 positive из 5;

• 2 hard negatives из 5.

Так модель видит разные комбинации и не переобучается на конкретные пары.

Внутренний формат jsonl:

{"query": "...", "pos": ["...", "...", "...", "...", "..."], "neg": ["...", "...", "...", "...", "..."]}

Размер файлов: 250,6 MB на train.

3.3. Скрипт подготовки датасета

Файл: scripts/2.4_prepare_hnp.py

"""Подготовка deepvk/ru-HNP: сохранение в jsonl + meta.

Сохраняем полные списки pos и neg — динамический sampling на уровне Dataset.
Читает датасет с HuggingFace, фильтрует пустые, пишет hnp_{train,val,test}.jsonl.
"""
import json
import time
from pathlib import Path

import datasets

ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"
DATA.mkdir(parents=True, exist_ok=True)
META_DIR = DATA / "meta"
META_DIR.mkdir(exist_ok=True)
META = META_DIR / "2.4_prepare_hnp_meta.json"

OUT_TRAIN = DATA / "hnp_train.jsonl"
OUT_VAL = DATA / "hnp_val.jsonl"
OUT_TEST = DATA / "hnp_test.jsonl"


def dump_split(ds, path: Path):
    """Фильтрует пустые и пишет jsonl. Возвращает статистику."""
    n_raw = len(ds)
    n_written = 0
    n_skip_empty_query = 0
    n_skip_empty_pos = 0
    n_skip_empty_neg = 0
    pos_counts, neg_counts = [], []

    t0 = time.time()
    with path.open("w", encoding="utf-8") as f:
        for row in ds:
            query = (row["query"] or "").strip()
            if not query:
                n_skip_empty_query += 1
                continue
            pos = [p for p in row["pos"] if p and p.strip()]
            neg = [n for n in row["neg"] if n and n.strip()]
            if not pos:
                n_skip_empty_pos += 1
                continue
            if not neg:
                n_skip_empty_neg += 1
                continue
            f.write(json.dumps(
                {"query": query, "pos": pos, "neg": neg},
                ensure_ascii=False,
            ) + "\n")
            n_written += 1
            pos_counts.append(len(pos))
            neg_counts.append(len(neg))

    size_mb = path.stat().st_size / 1e6
    print(f"  {path.name}: {n_written:,}/{n_raw:,} → {size_mb:.1f} MB "
          f"({time.time() - t0:.0f}s)", flush=True)

    return {
        "raw": n_raw,
        "written": n_written,
        "skipped": {
            "empty_query": n_skip_empty_query,
            "empty_pos": n_skip_empty_pos,
            "empty_neg": n_skip_empty_neg,
        },
        "pos": {
            "min": min(pos_counts) if pos_counts else 0,
            "max": max(pos_counts) if pos_counts else 0,
            "mean": round(sum(pos_counts) / len(pos_counts), 2) if pos_counts else 0,
        },
        "neg": {
            "min": min(neg_counts) if neg_counts else 0,
            "max": max(neg_counts) if neg_counts else 0,
            "mean": round(sum(neg_counts) / len(neg_counts), 2) if neg_counts else 0,
        },
        "file_mb": round(size_mb, 1),
    }


def main():
    t_start = time.time()

    print("Loading deepvk/ru-HNP from HuggingFace...", flush=True)
    ds = datasets.load_dataset("deepvk/ru-HNP")
    print(ds, flush=True)

    print("\nDumping splits:", flush=True)
    stats = {
        "train": dump_split(ds["train"], OUT_TRAIN),
        "val":   dump_split(ds["validation"], OUT_VAL),
        "test":  dump_split(ds["test"], OUT_TEST),
    }

    meta = {
        "script": "2.4_prepare_hnp",
        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "source": "deepvk/ru-HNP",
        "splits": stats,
        "elapsed_sec": round(time.time() - t_start, 1),
    }
    META.write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")

    print(f"\nFiles in {DATA}:")
    for p in sorted(DATA.glob("hnp_*.jsonl")):
        print(f"  {p.name:20s} {p.stat().st_size / 1e6:6.1f} MB")
    print(f"meta: {META}")


if __name__ == "__main__":
    main()

Ожидаемый вывод:

Dumping splits:
  hnp_train.jsonl: 99,999/100,000 → 250.6 MB (4s)
  hnp_val.jsonl: 590/590 → 1.5 MB (0s)
  hnp_test.jsonl: 2,000/2,000 → 5.1 MB (0s)

Files in /..../rubert-tiny2-distill/data:
  hnp_test.jsonl          5.1 MB
  hnp_train.jsonl       250.6 MB
  hnp_val.jsonl           1.5 MB
meta: /..../rubert-tiny2-distill/data/meta/2.4_prepare_hnp_meta.json