Fantominsight/Polina-Lumen-NumDetective

Полина Люмен:: Числовой Инспектор (OCR/digitGPT)

Маленькая GPT-подобная модель (4.8M параметров), которая читает полоску из 1–5 цифр подряд и выписывает их в строку. Реализация — decoder-only трансформер поверх патчей картинки.

Важно про архитектуру репозитория. Это не модель для transformers: нет PreTrainedModel, нет AutoModelForCausalLM, нет tokenizers. Код написан на чистом PyTorch, веса лежат в safetensors, токенизатор — обычный класс DigitTokenizer. Для загрузки нужен код из этого репозитория (файлы model.py, tokenizer.py, checkpoint.py).

Model Details

Архитектура. DigitGPT — decoder-only трансформер с причинной маской внимания. Входная картинка режется на патчи 4×4, каждый линейно проецируется в d_model; затем к патчам приписываются токены цифр, и модель предсказывает следующую цифру авторегрессивно.

параметры	4,804,352
слои	6
d_model	256
головы внимания	4
d_ff	1024
активация / norm	GELU / LayerNorm (без bias)
dropout	0.05
патчи	4×4 → 7×35 = 245
вход	28×140 grayscale, диапазон [0.0, 1.0]
последовательность	245 + 37 = до 282 позиций
точность весов	float32
формат весов	safetensors (53 тензора)

Токенизатор. Собственный, словарь из 12 токенов: 0–9, <pad> = 10, <eos> = 11. Никакого внешнего словаря.

Выход. Строка цифр без разделителей. Длина определяется самой моделью через <eos>, а не задаётся заранее.

Intended use

• Распознавание 1–5 цифр на плотной горизонтальной полоске из MNIST-подобных цифр, шириной ровно 28 px на цифру, выровненных по левому краю.
• Учебный пример GPT-подобной архитектуры, работающей с картинкой как с последовательностью: патчи + токены в одном контексте, причинная маска.
• Позиция baseline: оценка, сколько даёт маленький decoder-only на 3000 синтетических примерах.

Out of scope

Модель не рассчитана на: рукописный ввод, фотографии реальных документов, разное количество цифр вне диапазона 1–5, разделители и другие символы, цветные изображения, шум/размытие/перекос, цифры, не вошедшие в MNIST.

Training data

Синтетически собрано из MNIST (torchvision.datasets.MNIST): 1–5 цифр склеиваются встык на полотно 28×140 — по 28 px на цифру, без зазоров и полей. Модель никогда не видела реальных фотографий.

обучающая выборка	3 000 полосок из MNIST train split
валидационная выборка	500 полосок из MNIST test split
отдельный test split	не выделялся (см. Evaluation results)
шум при генерации	0.0
seed	12345

Выборка мала намеренно (тестовый прогон). Это главное ограничение модели — см. Limitations.

Обучение MNIST распространяется на CC BY-SA 3.0 (Yann LeCun, Corinna Cortes, Christopher J.C. Burges). Атрибуция обязательна при использовании и весов, и обучающих данных.

Training procedure

Обучение ровно то, что описано в model/training_config.yaml:

эпохи	6
batch size	1 (grad accum 4)
оптимизатор	AdamW, lr 3e-4, weight_decay 0.05, betas (0.9, 0.95)
расписание	200 шагов warmup → cosine
grad clip	1.0
AMP	bfloat16 (на CPU отключается)
железо	CPU-only, ~75 мин на 6 эпох
Ctrl+C	сохраняет бандл со stage: interrupted

Кривая обучения — монотонная, переобучения нет:

Эпоха	train loss	val loss	val token acc
1	1.822	1.675	0.403
2	1.445	1.181	0.555
3	1.088	0.858	0.692
4	0.819	0.675	0.763
5	0.714	0.633	0.790
6	0.629	0.574	0.812

Evaluation results

Замер на 200 полосках из MNIST test split — на них модель не обучалась. Это тот же официальный сплит, из которого взят валидационный набор для кривой выше, но другое подмножество из 200 полосок, поэтому цифры ниже немного оптимистичны относительно по-настоящему свежего теста. Метрики считаются жадной генерацией (greedy=True), если не сказано иначе.

Метрика	Значение
точное совпадение всей строки	0.460 (92/200)
совпадение количества цифр	1.000 (200/200)
token accuracy (teacher forcing)	0.812 (658/810)

Как читать эти числа: модель всегда выдаёт правильное число цифр, но одна-две цифры из пяти читает неверно. Разрыв между 0.812 и 0.460 — это именно цена одной ошибки на позицию при строгом сравнении строк. Для использования «распознать все цифры» это существенно: цифра с ошибкой делает весь результат неверным, даже если остальные верны.

How to use

Зависимости: torch, safetensors, numpy, Pillow (+ PyYAML, torchvision — только для обучения).

Вариант 1 — через хелпер репозитория

from checkpoint import load_bundle
from infer import preprocess_image

model, tokenizer, manifest = load_bundle("./model")   # strict=True
tensor = preprocess_image("digits.png", 28, 140)      # сама приведёт к 28x140

ids = model.generate(tensor, max_len=6, greedy=True)
print(tokenizer.decode(ids[0].tolist()))              # например: '93113'

Вариант 2 — напрямую, только safetensors + model.py

import json, torch
from safetensors.torch import load_file
from model import DigitGPT, GPTConfig
from tokenizer import DigitTokenizer

manifest = json.load(open("./model/config.json"))
cfg = {k: v for k, v in manifest["model"].items() if k in GPTConfig.__dataclass_fields__}

model = DigitGPT(GPTConfig(**cfg), img_h=cfg["img_h"], img_w=cfg["img_w"])
model.load_state_dict(load_file("./model/model.safetensors"), strict=True)  # missing/unexpected = []
tokenizer = DigitTokenizer.load("./model/tokenizer.json")
model.eval()

with torch.no_grad():
    ids = model.generate(tensor, max_len=6, greedy=True)
print(tokenizer.decode(ids[0].tolist()))

strict=True — при несовпадении весов и config.json будет исключение, а не молчаливо случайная модель.

Обучение с нуля

python train.py                 # создаст ./model
python train_more.py 10 --resume   # продолжить из ./model
python infer.py path/to/image.png

Preprocessing: что модель ожидает

infer.preprocess_image пропорционально вписывает картинку в полотно 28×140 (28 px на цифру, до 5 цифр), выравнивая по левому краю и центруя по вертикали. Увеличение картинки (4x, 10x) и уменьшение работают.

Что не работает и является ограничением по данным, а не багом кода:

• широкие белые поля вокруг цифр;
• инверсия (чёрные цифры на белом фоне);
• сильный шум, размытие, перекос.

Попытка лечить это автокадрированием по чернилам была проверена и отклонена: на бенчмарке из 80 случаев точное совпадение упало с 21/80 до 6/80, потому что цифра MNIST занимает ~20 px из 28, и подгонка по высоте растягивала её до 28 px. Реализовано только LANCZOS при даунскейле — это убирает алиасинг без сдвига масштаба.

Files

файл	что это
model.py	DigitGPT, GPTConfig, PatchEmbed
tokenizer.py	DigitTokenizer (12 токенов)
checkpoint.py	save_bundle / load_bundle
data.py	склейка 1–5 цифр из MNIST
train.py, train_more.py	обучение и дообучение
infer.py	препроцессинг + предсказание
config.yaml	настройки обучения (эталон)
model/	бандл: веса, токенизатор, config.json, metrics.json, runs/
model_leaky_backup/	старые веса с утечкой метки, не использовать (см. ниже)

Limitations

• Слабое качество. 46% точных строк на 5 цифрах — многовато для практики. Причина — 3000 обучающих полосок; это осознанно маленький тестовый бюджет, а не предел архитектуры.
• Узкий домен. Только синтетические склейки MNIST. Любая реальная фотография — вне обучающего распределения.
• Только 1–5 цифр. Другое количество цифр модель не обрабатывает.
• Нет устойчивости к деформациям. Ни поворотов, ни сдвигов, ни размытия при обучении не применялось (noise: 0.0).
• Нет тестового сплита. Все отчёты — на валидации из 500 полосок.
• Только CPU-рецепт. Обучение и инференс на CPU; GPU-путь в конфиге есть, но не проверен и не использовался.

Model history

Первая версия весов обучалась с утечкой метки и была непригодна. Исправлено три бага в model.py:

1. Утечка в forward. logits[:, Np:, :] — позиция Np содержит embedding истинного токена и предсказывала этот же токен, то есть модель училась копировать вход вместо чтения картинки. Правильный сдвиг: logits[:, Np - 1 : T_total - 1, :].
2. generate без причинной маски. Патчи видели друг друга двунаправленно, хотя при обучении маска была причинная.
3. generate прогонял выходы блоков через блоки повторно — torch.cat([cur, ...]), где cur уже результат всех 6 блоков, давал 12 слоёв вместо 6.

Симптом был обманчивым: train_acc показывал 0.98, а сгенерированные картинки читались как мусор. На честной метрике старые веса дают 0.037 token accuracy — уровень случайного угадывания. Старые веса сохранены в model_leaky_backup/ и оставлены только как артефакт; в репозитории лежат только как пример того, как выглядит утечка метки.

Дополнительно: потери в логе были завышены в grad_accum_steps раз; Ctrl+C теперь сохраняет бандл вместо потери прогресса.

Citation

Если пригодится в работе:

@misc{mnist_multidigit_gpt,
  title  = {mnist-multidigit-gpt: patch-based decoder-only digit sequence recognition},
  note   = {Custom PyTorch implementation, 4.8M params, trained on 3000 synthetic MNIST strips},
  year   = {2026},
  license = {GPL-3.0}
}

Обучающие данные:

@inproceedings{lecun1998mnist,
  title   = {Gradient-Based Learning Applied to Document Recognition},
  author  = {LeCun, Yann and Cortes, Corinna and Burges, Christopher J. C.},
  booktitle = {Proceedings of the IEEE},
  year    = {1998}
}

License

GPL-3.0 — на код и веса. Обучающие данные MNIST — CC BY-SA 3.0, атрибуция обязательна. Загружая эти веса, вы принимаете условия GPL-3.0 на производные работы.