Fantominsight/Polina-Lumen-OCR

Полина Люмен (OCR/DigitGPT)

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

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

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 при обучении на CPU.

Out of scope

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

Training data

Датасет собирается на лету из MNIST (torchvision.datasets.MNIST), а не заранее в файл:

1. выбирается случайное число цифр k из [1, 5];
2. k цифр берётся случайно из MNIST и склеивается встык на полотно 28×140 — по 28 px на цифру, без зазоров и полей;
3. train_samples / val_samples задают количество полосок в каждой эпохе и в каждой валидации, а не размер фиксированного датасета.

Каждая эпоха поэтому показывает модель новые изображения, а не повторяет одни и те же. Модель никогда не видела реальных фотографий.

полосок в эпоху	10 000 (из MNIST train split, 60 000 цифр)
полосок в валидации	500 (из MNIST test split, 10 000 цифр)
отдельный test split	не выделялся (см. Evaluation results)
шум при генерации	0.0
seed	12345 (random, numpy, torch)

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

Training procedure

Параметры взяты из training_config.yaml, сохранённого в бандле, — это авторитетная запись запуска:

эпохи	6
batch size	32, grad accum 4 → эффективный батч 128
оптимизатор	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, ~20 мин на эпоху
Ctrl+C	сохраняет бандл со stage: interrupted

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

Эпоха	train loss	train acc	val loss	val token acc
1	2.025	0.326	1.509	0.498
2	1.615	0.435	1.361	0.518
3	1.399	0.507	1.134	0.593
4	0.984	0.658	0.667	0.783
5	0.633	0.788	0.487	0.837
6	0.518	0.827	0.437	0.850

Evaluation results

85% — это teacher-forced token accuracy на валидации, то есть доля правильно названных токенов, когда модель «посмотрела» на истинные предыдущие токены. Это мягкая метрика. Для настоящей задачи «прочитать цифры» важнее другая — точное совпадение всей строки.

Замер независимой проверкой на 200 полосках из MNIST test split (генерация с нуля, greedy=True, модель истинные токены не видит):

Метрика	Значение
val teacher-forced token acc (500 полосок, из metrics.json)	0.850
token accuracy, независимая проверка (200 полосок)	0.828
точное совпадение всей строки	0.480 (96/200)
совпадение количества цифр	1.000 (200/200)

Разница между 0.850 и 0.828 — шум выборки, а не расхождение: при n=500 стандартная ошибка при p=0.85 равна ±1.6 пп, при n=200 — ±2.5 пп.

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

How to use

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

pip install -r scripts/requirements.txt

Команда запуска — из корня репозитория:

python scripts/infer.py path/to/image.png -m .

Флаг -m . обязателен: веса и config.json лежат в корне репозитория, а по умолчанию код ищет бандл в ./model.

Опции:

python scripts/infer.py img.png -m . --sampling    # сэмплирование вместо greedy
python scripts/infer.py img.png -m . --show-input  # сохранить препроцесс 28x140
python scripts/infer.py img.png -m . --device cuda

Из Python:

import sys; sys.path.insert(0, "scripts")
from checkpoint import load_bundle
from infer import preprocess_image

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

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

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

Обучение с нуля (внимание: train.py читает config.yaml из текущего каталога, поэтому запускать из папки scripts):

cd scripts
python train.py
python train_more.py 10 --resume -m ..   # продолжить из корня

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

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

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

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

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

Files

README.md                    эта карточка
model.safetensors            веса (float32)
config.json                  архитектура + спецификация картинки (источник истины)
tokenizer.json               словарь: 0-9, <pad>=10, <eos>=11
training_config.yaml         полный конфиг запуска, из которого получен бандл
metrics.json                 метрики по эпохам
scripts/                     весь исходный код

В scripts/:

файл	что это
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	настройки обучения
requirements.txt	зависимости
README.md	короткая инструкция по раскладке репозитория

Limitations

• Слабое качество. 48% точных строк на 5 цифр — многовато для практики. Причина — 10000 полосок в эпоху; это маленький бюджет, а не предел архитектуры. Увеличение числа примеров даёт малый прирост: переход с 3000 на 10000 дал +1.6 пп token accuracy и +2.0 пп exact match.
• Узкий домен. Только синтетические склейки MNIST. Любая реальная фотография — вне обучающего распределения.
• Только 1–5 цифр. Другое количество цифр модель не обрабатывает.
• Нет устойчивости к деформациям. Ни поворотов, ни сдвигов, ни размытия при обучении не применялось (noise: 0.0).
• Валидация мала. 500 полосок дают стандартную ошибку ±1.6 пп при p=0.85, так что числа на 1–3 пп варьируются от запуска к запуску.
• Нет выделенного тестового сплита. Отчётные цифры — на валидации из MNIST test split.
• Только CPU-рецепт. Обучение и инференс на CPU; GPU-путь в конфиге есть, но не проверен и не использовался.

Model history

Текущий бандл — 10000 полосок/эпоху, 6 эпох, val_token_acc = 0.850.

Первая версия весов (3000 полосок) обучалась с утечкой метки и была непригодна. Исправлено три бага в 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 — уровень случайного угадывания.

Все бандлы после этих исправлений — только исправленный код. Валидация после фиксов:

бандл	полосок/эпоху	token acc (200 полосок)	exact match
первый исправленный	3 000	0.812	0.460
текущий	10 000	0.828	0.480

Дополнительно было исправлено: потери в логе были завышены в 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 10000 synthetic MNIST strips per epoch},
  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 на производные работы.