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