Jevlike

· 2 мин чтения
machine-learning nlp attention text-classification model-training
📂 Исходный код на GitHub

Train a small model that chooses among a changing list of text options

Jevlike

Jevlike — это исследовательский стартовый репозиторий для обучения компактных моделей, которые за один проход выбирают правильный вариант из динамического списка текстовых опций. В отличие от авторегрессионных декодеров, генерирующих ответ посимвольно, Jevlike вычисляет вероятности для всех вариантов одновременно — это даёт ускорение в ~100 раз при работе с восемью опциями.

Модель вдохновлена коммерческим решением Jev от TypeSafe, однако является полностью независимой реализацией. Автор не воспроизводил дизайн TypeSafe и не претендует на совместимость с их методом.

Архитектура

Каждый вариант текста преобразуется в query-вектор фиксированной длины. Этот вектор вычисляет веса attention по токенам контекста, формируя context-вектор для конкретного варианта. Затем dot product между вариантом и контекстом даёт скалярный скор. Softmax по всем вариантам превращает скоры в вероятности, суммирующиеся в единицу.

option -> query_vector -> attention_weights -> context_vector -> dot_product -> score
    |                                                                              |
scores ---------------------------------------------------------------- softmax -> probabilities

По умолчанию модель использует byte-level encoder, обучающийся с нуля. Альтернативный путь — замороженный pretrained encoder из Hugging Face (например, Qwen2.5-0.5B), который обеспечивает лучшее качество за счёт предобученных весов.

Формат данных

Данные хранятся в JSONL — по одному JSON-объекту на строку:

{"context":"The customer needs a refund.","options":["refund","sales","technical support"],"label":0}

Поле label — нулевой индекс правильного варианта. Количество опций может варьироваться от записи к записи (минимум 2).

Быстрый старт

uv venv
source .venv/bin/activate
uv pip install -e '.[dev]'

jevlike-data synthetic --output data/synthetic
jevlike-train data/synthetic/train.jsonl \
  --validation data/synthetic/validation.jsonl \
  --output runs/synthetic.pt
jevlike-eval runs/synthetic.pt data/synthetic/test.jsonl
jevlike-predict runs/synthetic.pt \
  --context "Choose the exact badge amber badger. Badge: amber badger." \
  --option "azure crane" \
  --option "amber badger" \
  --option "gold heron"

Обучение на своих данных

  1. Экспортируйте train, validation и test в JSONL.
  2. Сохраняйте все опции, которые модель увидит при предсказании, в каждой записи.
  3. Разделяйте связанные записи вместе (например, все записи для одного клиента в один сплит).
  4. Запустите jevlike-train на train и validation.
  5. Оцените на held-out test файле.

Byte-encoder по умолчанию обрезает контекст до 192 байт и каждый вариант до 32 байт. Параметры --context-tokens и --option-tokens позволяют увеличить лимиты.

Использование pretrained encoder

uv pip install -e '.[transformers]'
jevlike-train data/synthetic/train.jsonl \
  --validation data/synthetic/validation.jsonl \
  --output runs/qwen-head.pt \
  --encoder hf \
  --hf-model Qwen/Qwen2.5-0.5B \
  --rank 256 \
  --batch-size 8

Checkpoint сохраняет обученный scorer head и имя encoder. Веса замороженного encoder не копируются — при загрузке требуется доступ к той же модели Hugging Face.

Параметр --rank задаёт ширину scorer head. Более широкий head содержит больше обучаемых параметров и потребляет больше памяти.

Примеры: Doom и шахматы

Jevlike включает демо-примеры с визуальным скринером из jevlike.vision. Один и тот же option-attention head оценивает кнопки управления из image patches.

  • Doom: joint checkpoint набирает 0.60 kills и -97.50 reward за десять эпизодов на deadly_corridor.
  • Шахматы: chess-only checkpoint выигрывает 4, ничьих 46 из 50 игр против случайного противника.

Запись свежего Doom-трейса:

uv pip install -e '.[games]'
python examples/doom/play.py examples/checkpoints/joint-imitation.pt \
  --episodes 10 --game-seconds 35.3 --device cpu \
  --capture-resolution 640x480 --output runs/doom.mp4 --trace runs/doom-trace.json

Результаты

  • ~98% top-1 accuracy на синтетических данных.
  • 26% на Wikispeedia next-click с замороженным Qwen2.5-0.5B (против 8% для shuffled control).
  • 29% при обучении с нуля на 40,000 кликов.
  • При восьми опциях один проход примерно в 100 раз быстрее маленького декодера, генерирующего 400 токенов.

Ограничения

  • Это исследовательский старт, а не копия Jev.
  • Точность зависит от качества данных, сплитов и encoder.
  • Byte-encoder дешёвый, но слабо понимает язык.
  • Pretrained путь может загрузить большую модель и требует больше памяти.
  • Однопроходное скоринг требует полного списка опций до предсказания.

Источник: https://github.com/vinnylarouge/jevlike