TRAINING.md

Training / fine-tuning

Baberu OCR is a from-scratch 115M model: a frozen DINOv2 vision encoder, an MLP projector, and a custom 6-layer character-level GQA decoder. The training code is included so you can fine-tune the released checkpoint on your own bubbles, or reproduce the full recipe. All scripts are flat — run them from the repo root.

Fine-tune the released model on your own crops (Step 3)

Continues image->text training from the released weights with a fresh optimizer:

python train_ocr.py \
  --crops-dir     /path/to/your/crops \
  --index         /path/to/ocr_text_index.parquet \
  --tokenizer-dir ./tokenizer \
  --finetune-from .            # load full weights from this repo, fresh schedule
  --out-dir       ./ft-out \
  --epochs 1 --batch-size 96 --num-workers 8
  # add  --unfreeze-vision --vision-lr 1e-5  to also adapt the encoder

• --finetune-from <dir> loads config.json + model.safetensors from <dir> (point it at this repo) and starts a fresh optimizer/schedule.
• --resume-from <dir> instead continues an interrupted run of this trainer (restores optimizer/scheduler/RNG from training_state.pt).

Data format

--index is a parquet (built by build_ocr_index.py) pairing each crop id with its text and language; --crops-dir holds the crop images the index refers to. ocr_pairs.py / data_ocr.py show the exact loader — adapt them to your store.

Full recipe (from scratch)

1. train_text.py — pretrain the decoder as a character LM (FineWeb2 + text).
2. train_text.py again — adapt to manga-style text (Step 2).
3. train_ocr.py --init-decoder-from <step2> — connect vision and train on (crop, text) pairs; then re-run with --unfreeze-vision to adapt DINOv2.

See each script's module docstring for the exact flags.

Files

• train_ocr.py — Step 3 image->text training / fine-tuning entry
• train_text.py — Step 1/2 character-LM pretraining
• data_ocr.py, ocr_pairs.py — OCR (crop, text) data loading
• data_baberu.py, data_fineweb.py — text data loading for the LM stages
• build_ocr_index.py — build the crop->text index parquet
• modeling_baberu.py, configuration_baberu.py, tokenization_baberu.py — the model