File size: 2,266 Bytes
82eb32a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | # 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
|