diff --git a/README.md b/README.md index d36bdaf..198852f 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ Model: **Gemma 4 E2B** (`src/train.py`) -- Unsloth QLoRA fine-tuning. ### Overview The training set combines handwritten math datasets with synthetically rendered -Typst documents. Handwriting-font splits (`hw_*`) cover both single-equation +Typst documents. Handwriting-font splits (`typeset_*`) cover both single-equation and structured document content rendered in 6 diverse fonts plus the Typst default, sampled uniformly. @@ -46,8 +46,8 @@ the Typst default (New Computer Modern), sampled uniformly (~14% default). | Split | Samples | Notes | |---|---|---| -| `hw_structured_train` | 15,000 | Whole document uses one uniformly-sampled font. | -| `hw_mixed_train` | 10,000 | Per-paragraph font mixing (~55% of blocks get hw font); requires multi-block bodies. | +| `typeset_uniform_train` | 15,000 | Whole document uses one uniformly-sampled font. | +| `typeset_mixed_train` | 10,000 | Per-paragraph font mixing (~55% of blocks get hw font); requires multi-block bodies. | Body types and generation weights (see `src/generate_mixed.py`): @@ -66,7 +66,7 @@ bold, italic, or underlined. Ink colour sampled from near-black palette. - **Math-only splits** (`mathwriting_*`, `crohme_*`): manifest stores bare math expressions. `data.load_records()` wraps these as `$ ... $` at load time. -- **hw_* splits**: manifest stores complete body content with inline `$...$` +- **typeset_* splits**: manifest stores complete body content with inline `$...$` delimiters already present. No wrapping applied. ### Effective training mix (after caps) @@ -77,8 +77,8 @@ bold, italic, or underlined. Ink colour sampled from near-black palette. | `mathwriting_synthetic` | ~85,879 | **20,000** | 21% | | `crohme_gen_2019` | ~51,855 | **15,000** | 16% | | `crohme_gen_syntactic` | ~69,397 | **15,000** | 16% | -| `hw_structured_train` | 15,000 | 15,000 | 16% | -| `hw_mixed_train` | 10,000 | 10,000 | 11% | +| `typeset_uniform_train` | 15,000 | 15,000 | 16% | +| `typeset_mixed_train` | 10,000 | 10,000 | 11% | | `crohme_gen_2023` | ~3,072 | 3,072 | 3% | | `mathwriting_symbols` | ~6,091 | 6,091 | 6% | | `crohme_real_train` | ~9 | 9 | <1% | @@ -91,8 +91,8 @@ bold, italic, or underlined. Ink colour sampled from near-black palette. | Split | Samples | Notes | |---|---|---| | `mathwriting_val` | 250 | Real handwritten single equations. | -| `hw_structured_val` | 250 | Whole-doc font document fragments. | -| `hw_mixed_val` | 250 | Per-block mixed-font document fragments. | +| `typeset_uniform_val` | 250 | Whole-doc font document fragments. | +| `typeset_mixed_val` | 250 | Per-block mixed-font document fragments. | | **Total** | **750** | | ### Test @@ -100,12 +100,12 @@ bold, italic, or underlined. Ink colour sampled from near-black palette. | Split | Notes | |---|---| | `mathwriting_test` | Held-out real handwritten equations. | -| `hw_structured_test` | Held-out whole-doc font document fragments. | -| `hw_mixed_test` | Held-out mixed-font document fragments. | +| `typeset_uniform_test` | Held-out whole-doc font document fragments. | +| `typeset_mixed_test` | Held-out mixed-font document fragments. | ### Known gaps -- `hw_*` splits are font-based renders, not real handwriting photos. The model +- `typeset_*` splits are font-based renders, not real handwriting photos. The model still lacks real handwritten document fragments at scale. - `crohme_real_train` has only 9 samples. diff --git a/pyproject.toml b/pyproject.toml index 724393d..68c984d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,9 @@ dependencies = [ [project.scripts] train = "src.train:main" +train-deepseek = "src.train_deepseek:main" evaluate = "src.eval:main" +evaluate-deepseek = "src.eval_deepseek:main" train-hnm = "src.train_hnm:main" export = "src.export:main" generate-typeset = "src.generate_typeset:main" diff --git a/src/data.py b/src/data.py index e0b57a7..0a6a2da 100644 --- a/src/data.py +++ b/src/data.py @@ -52,6 +52,9 @@ _MATH_ONLY_SPLITS = { PROMPT = "Transcribe this image to Typst notation. Output only the raw Typst, without explanation. No LaTeX, only Typst." BASE_MODEL = "unsloth/gemma-4-E2B-it" +DEEPSEEK_PROMPT = "\nFree OCR. " +DEEPSEEK_MODEL_DIR = str(Path(__file__).parent.parent / "deepseek_ocr2") + _NUM_RE = re.compile(r"\d+(\.\d+)?") _VAR_RE = re.compile(r"\b[a-zA-Z]\b") @@ -128,15 +131,58 @@ def _augment(img: Image.Image) -> Image.Image: return img +def gemma_format(record: dict, img: Image.Image) -> dict: + return { + "messages": [ + { + "role": "user", + "content": [ + {"type": "image", "image": img}, + {"type": "text", "text": PROMPT}, + ], + }, + { + "role": "assistant", + "content": [{"type": "text", "text": record["typst"]}], + }, + ] + } + + +def deepseek_format(record: dict, img: Image.Image) -> dict: + return { + "messages": [ + { + "role": "<|User|>", + "content": DEEPSEEK_PROMPT, + "images": [img], + }, + { + "role": "<|Assistant|>", + "content": record["typst"], + }, + ] + } + + class MathOCRDataset(Dataset): """ Torch Dataset with lazy image loading. Avoids HuggingFace datasets fingerprinting, which uses dill and breaks on Python 3.14. + + format_fn(record, img) -> dict with a "messages" key in the model's + expected conversation format. Defaults to Gemma format. """ - def __init__(self, records: list[dict], do_augment: bool = False) -> None: + def __init__( + self, + records: list[dict], + do_augment: bool = False, + format_fn=None, + ) -> None: self.records = records self.do_augment = do_augment + self.format_fn = format_fn if format_fn is not None else gemma_format def __len__(self) -> int: return len(self.records) @@ -146,22 +192,12 @@ class MathOCRDataset(Dataset): img = Image.open(r["image_path"]).convert("RGB") if self.do_augment: img = _augment(img) - return { - "messages": [ - { - "role": "user", - "content": [ - {"type": "image", "image": img}, - {"type": "text", "text": PROMPT}, - ], - }, - { - "role": "assistant", - "content": [{"type": "text", "text": r["typst"]}], - }, - ] - } - - -def make_dataset(records: list[dict], do_augment: bool = False) -> MathOCRDataset: - return MathOCRDataset(records, do_augment) + return self.format_fn(r, img) + + +def make_dataset( + records: list[dict], + do_augment: bool = False, + format_fn=None, +) -> MathOCRDataset: + return MathOCRDataset(records, do_augment, format_fn) diff --git a/src/eval_deepseek.py b/src/eval_deepseek.py new file mode 100644 index 0000000..a698489 --- /dev/null +++ b/src/eval_deepseek.py @@ -0,0 +1,115 @@ +""" +ExpRate evaluation of DeepSeek-OCR-2 on the test splits. + +Inference uses model.infer() (the model's native API), which takes a file path, +so images are written to a temp file per sample. + +Usage: uv run evaluate-deepseek [--checkpoint checkpoints/deepseek-ocr2/final] +""" + +import argparse +import os +import re +import tempfile + +import torch +from tqdm import tqdm +from transformers import AutoModel +from unsloth import FastVisionModel + +from .data import DEEPSEEK_MODEL_DIR, TEST_SPLITS, load_records + +os.environ["UNSLOTH_WARN_UNINITIALIZED"] = "0" + + +def normalize(s: str) -> str: + return re.sub(r"\s+", " ", s).strip() + + +def evaluate( + checkpoint: str, + n: int | None = None, + verbose: bool = False, + image_size: int = 768, + base_size: int = 1024, + crop_mode: bool = True, +) -> float: + model, tokenizer = FastVisionModel.from_pretrained( + checkpoint, + load_in_4bit=False, + auto_model=AutoModel, + trust_remote_code=True, + unsloth_force_compile=True, + use_gradient_checkpointing="unsloth", + ) + FastVisionModel.for_inference(model) + model.eval() + + import random + rng = random.Random(42) + records = [] + for split in TEST_SPLITS: + split_recs = load_records([split], dedupe=False) + sample = rng.sample(split_recs, min(n, len(split_recs))) if n is not None else split_recs + records.extend(sample) + + correct = 0 + prompt = "\nFree OCR. " + + with tempfile.TemporaryDirectory() as tmpdir: + tmp_img = os.path.join(tmpdir, "img.png") + for r in tqdm(records, desc="Evaluating"): + # model.infer requires a file path + import shutil + shutil.copy(r["image_path"], tmp_img) + + result = model.infer( + tokenizer, + prompt=prompt, + image_file=tmp_img, + output_path=tmpdir, + image_size=image_size, + base_size=base_size, + crop_mode=crop_mode, + save_results=False, + test_compress=False, + ) + # model.infer returns the transcription string directly + pred = result if isinstance(result, str) else str(result) + + match = normalize(pred) == normalize(r["typst"]) + if match: + correct += 1 + if verbose: + status = "OK" if match else "FAIL" + print(f"[{status}] split={r['split']}") + print(f" GT: {r['typst']}") + print(f" PRED: {pred}") + + exprate = correct / len(records) + split_label = f" ({n}/split)" if n is not None else "" + print(f"ExpRate{split_label}: {exprate:.4f} ({correct}/{len(records)})") + return exprate + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoint", default=f"checkpoints/deepseek-ocr2/final") + parser.add_argument("--n", type=int, default=None, help="Evaluate N random examples per split") + parser.add_argument("--verbose", action="store_true") + parser.add_argument("--image-size", type=int, default=768) + parser.add_argument("--base-size", type=int, default=1024) + parser.add_argument("--no-crop", action="store_true", help="Disable dynamic cropping") + args = parser.parse_args() + evaluate( + args.checkpoint, + n=args.n, + verbose=args.verbose, + image_size=args.image_size, + base_size=args.base_size, + crop_mode=not args.no_crop, + ) + + +if __name__ == "__main__": + main() diff --git a/src/train_deepseek.py b/src/train_deepseek.py new file mode 100644 index 0000000..b72f0cc --- /dev/null +++ b/src/train_deepseek.py @@ -0,0 +1,348 @@ +""" +LoRA fine-tuning of DeepSeek-OCR-2 for Typst math OCR. + +The model uses a fixed "Free OCR" prompt with no way to specify output format, +so finetuning replaces the LaTeX prior with Typst via supervised examples. + +Usage: uv run train-deepseek [--epochs 2] [--cap mathwriting_synthetic 20000] + +Requires the model weights to be downloaded first: + python -c "from huggingface_hub import snapshot_download; snapshot_download('unsloth/DeepSeek-OCR-2', local_dir='deepseek_ocr2')" +""" + +from unsloth import FastVisionModel, is_bf16_supported +from PIL import Image, ImageOps +import torch +from torch.nn.utils.rnn import pad_sequence +from transformers import AutoModel, Trainer, TrainingArguments + +import argparse +import io +import math +import os +import random +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Dict, List, Tuple + + +from .data import ( + DEEPSEEK_MODEL_DIR, + TRAIN_SPLITS, + VAL_SPLITS, + deepseek_format, + load_records, + make_dataset, +) + +os.environ["UNSLOTH_WARN_UNINITIALIZED"] = "0" + +_DEFAULT_CAPS = { + "mathwriting_synthetic": 20_000, + "crohme_gen_2019": 15_000, + "crohme_gen_syntactic": 15_000, + "mathwriting_train": 10_000, +} + + +# --------------------------------------------------------------------------- +# Data collator (DeepSeek-OCR-2 specific) +# --------------------------------------------------------------------------- + +@dataclass +class DeepSeekOCR2DataCollator: + tokenizer: Any + model: Any + image_size: int = 768 + base_size: int = 1024 + crop_mode: bool = True + train_on_responses_only: bool = True + + def __post_init__(self) -> None: + from deepseek_ocr2.modeling_deepseekocr2 import BasicImageTransform, dynamic_preprocess # noqa: F401 + self._dynamic_preprocess = dynamic_preprocess + self.image_token_id = 128815 + self.dtype = self.model.dtype + self.patch_size = 16 + self.downsample_ratio = 4 + self.image_transform = BasicImageTransform( + mean=(0.5, 0.5, 0.5), + std=(0.5, 0.5, 0.5), + normalize=True, + ) + if hasattr(self.tokenizer, "bos_token_id") and self.tokenizer.bos_token_id is not None: + self.bos_id = self.tokenizer.bos_token_id + else: + self.bos_id = 0 + + def _to_pil(self, image_data) -> Image.Image: + if isinstance(image_data, Image.Image): + return image_data.convert("RGB") + if isinstance(image_data, dict) and "bytes" in image_data: + return Image.open(io.BytesIO(image_data["bytes"])).convert("RGB") + raise ValueError(f"Unsupported image format: {type(image_data)}") + + def _process_image(self, image: Image.Image) -> Tuple[List, List, List, List, Tuple]: + from deepseek_ocr2.modeling_deepseekocr2 import dynamic_preprocess + + images_list, images_crop_list, images_spatial_crop = [], [], [] + nq = math.ceil((self.image_size // self.patch_size) / self.downsample_ratio) + nq_base = math.ceil((self.base_size // self.patch_size) / self.downsample_ratio) + + if self.crop_mode: + if image.size[0] <= 768 and image.size[1] <= 768: + crop_ratio = (1, 1) + crops_raw = [] + else: + crops_raw, crop_ratio = dynamic_preprocess( + image, min_num=2, max_num=6, + image_size=self.image_size, use_thumbnail=False, + ) + pad_color = tuple(int(x * 255) for x in self.image_transform.mean) + global_view = ImageOps.pad(image, (self.base_size, self.base_size), color=pad_color) + images_list.append(self.image_transform(global_view).to(self.dtype)) + w, h = crop_ratio + images_spatial_crop.append([w, h]) + for c in crops_raw: + images_crop_list.append(self.image_transform(c).to(self.dtype)) + tok = ([self.image_token_id] * nq_base) * nq_base + [self.image_token_id] + if w > 1 or h > 1: + tok += ([self.image_token_id] * (nq * w)) * (nq * h) + else: + crop_ratio = (1, 1) + images_spatial_crop.append([1, 1]) + if self.base_size <= 768: + resized = image.resize((self.base_size, self.base_size), Image.LANCZOS) + images_list.append(self.image_transform(resized).to(self.dtype)) + else: + pad_color = tuple(int(x * 255) for x in self.image_transform.mean) + global_view = ImageOps.pad(image, (self.base_size, self.base_size), color=pad_color) + images_list.append(self.image_transform(global_view).to(self.dtype)) + tok = ([self.image_token_id] * nq) * nq + [self.image_token_id] + + return images_list, images_crop_list, images_spatial_crop, tok, crop_ratio + + def _process_sample(self, messages: List[Dict]) -> Dict[str, Any]: + from deepseek_ocr2.modeling_deepseekocr2 import text_encode + + images = [ + self._to_pil(img) + for msg in messages + for img in msg.get("images", []) + if img is not None + ] + if not images: + raise ValueError("No images in sample.") + + tokenized_str, images_seq_mask = [self.bos_id], [False] + images_list, images_crop_list, images_spatial_crop = [], [], [] + prompt_token_count = -1 + assistant_started = False + image_idx = 0 + + for msg in messages: + role = msg["role"] + content = msg["content"] + if role == "<|Assistant|>": + if not assistant_started: + prompt_token_count = len(tokenized_str) + assistant_started = True + content = f"{content.strip()} {self.tokenizer.eos_token}" + + for i, part in enumerate(content.split("")): + toks = text_encode(self.tokenizer, part, bos=False, eos=False) + tokenized_str.extend(toks) + images_seq_mask.extend([False] * len(toks)) + if i < len(content.split("")) - 1: + img_list, crop_list, spatial, tok_img, _ = self._process_image(images[image_idx]) + images_list.extend(img_list) + images_crop_list.extend(crop_list) + images_spatial_crop.extend(spatial) + tokenized_str.extend(tok_img) + images_seq_mask.extend([True] * len(tok_img)) + image_idx += 1 + + if not assistant_started: + prompt_token_count = len(tokenized_str) + + images_ori = torch.stack(images_list, dim=0) + images_spatial_crop_t = torch.tensor(images_spatial_crop, dtype=torch.long) + if images_crop_list: + images_crop = torch.stack(images_crop_list, dim=0) + else: + images_crop = torch.zeros((1, 3, self.base_size, self.base_size), dtype=self.dtype) + + return { + "input_ids": torch.tensor(tokenized_str, dtype=torch.long), + "images_seq_mask": torch.tensor(images_seq_mask, dtype=torch.bool), + "images_ori": images_ori, + "images_crop": images_crop, + "images_spatial_crop": images_spatial_crop_t, + "prompt_token_count": prompt_token_count, + } + + def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]: + batch_data = [] + for feat in features: + try: + batch_data.append(self._process_sample(feat["messages"])) + except Exception as e: + print(f"Skipping sample: {e}") + if not batch_data: + raise ValueError("Empty batch after collation.") + + input_ids = pad_sequence( + [d["input_ids"] for d in batch_data], + batch_first=True, + padding_value=self.tokenizer.pad_token_id, + ) + images_seq_mask = pad_sequence( + [d["images_seq_mask"] for d in batch_data], + batch_first=True, + padding_value=False, + ) + labels = input_ids.clone() + labels[labels == self.tokenizer.pad_token_id] = -100 + labels[images_seq_mask] = -100 + if self.train_on_responses_only: + for i, d in enumerate(batch_data): + pc = d["prompt_token_count"] + if pc > 0: + labels[i, :pc] = -100 + + attention_mask = (input_ids != self.tokenizer.pad_token_id).long() + images_batch = [(d["images_crop"], d["images_ori"]) for d in batch_data] + images_spatial_crop = torch.cat([d["images_spatial_crop"] for d in batch_data], dim=0) + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + "images": images_batch, + "images_seq_mask": images_seq_mask, + "images_spatial_crop": images_spatial_crop, + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--epochs", type=int, default=2) + parser.add_argument("--lr", type=float, default=2e-4) + parser.add_argument("--output-dir", default="checkpoints/deepseek-ocr2") + parser.add_argument("--model-dir", default=DEEPSEEK_MODEL_DIR) + parser.add_argument("--cap", nargs=2, action="append", metavar=("SPLIT", "N"), + default=None) + args = parser.parse_args() + + caps = dict(_DEFAULT_CAPS) + if args.cap: + for split, n in args.cap: + caps[split] = int(n) + + if not Path(args.model_dir).resolve().exists(): + print(f"{args.model_dir} not found -- downloading from HuggingFace...") + from huggingface_hub import snapshot_download + snapshot_download("unsloth/DeepSeek-OCR-2", local_dir=args.model_dir) + + model, tokenizer = FastVisionModel.from_pretrained( + args.model_dir, + load_in_4bit=False, + auto_model=AutoModel, + trust_remote_code=True, + unsloth_force_compile=True, + use_gradient_checkpointing="unsloth", + ) + + model = FastVisionModel.get_peft_model( + model, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj"], + r=16, + lora_alpha=16, + lora_dropout=0, + bias="none", + random_state=3407, + use_rslora=False, + loftq_config=None, + ) + FastVisionModel.for_training(model) + + rng = random.Random(42) + train_records: list[dict] = [] + for split in TRAIN_SPLITS: + recs = load_records([split], dedupe=False) + cap = caps.get(split) + if cap and len(recs) > cap: + recs = rng.sample(recs, cap) + train_records.extend(recs) + rng.shuffle(train_records) + + val_rng = random.Random(42) + val_records: list[dict] = [] + for split in VAL_SPLITS: + recs = load_records([split], dedupe=False) + val_records += val_rng.sample(recs, min(250, len(recs))) + val_rng.shuffle(val_records) + + print(f"Train: {len(train_records):,} Val: {len(val_records):,}") + for split in TRAIN_SPLITS: + n = sum(1 for r in train_records if r["split"] == split) + cap = caps.get(split) + print(f" {split}: {n:,}" + (f" (cap {cap:,})" if cap else "")) + + train_ds = make_dataset(train_records, do_augment=True, format_fn=deepseek_format) + val_ds = make_dataset(val_records, do_augment=False, format_fn=deepseek_format) + + data_collator = DeepSeekOCR2DataCollator( + tokenizer=tokenizer, + model=model, + image_size=768, + base_size=1024, + crop_mode=True, + train_on_responses_only=True, + ) + + out_dir = args.output_dir + trainer = Trainer( + model=model, + tokenizer=tokenizer, + data_collator=data_collator, + train_dataset=train_ds, + eval_dataset=val_ds, + args=TrainingArguments( + per_device_train_batch_size=2, + per_device_eval_batch_size=2, + gradient_accumulation_steps=8, + num_train_epochs=args.epochs, + learning_rate=args.lr, + warmup_steps=500, + lr_scheduler_type="cosine", + fp16=not is_bf16_supported(), + bf16=is_bf16_supported(), + dataloader_num_workers=2, + logging_steps=50, + eval_strategy="steps", + eval_steps=500, + save_steps=500, + save_total_limit=7, + load_best_model_at_end=False, + output_dir=out_dir, + run_name="deepseek-ocr2", + report_to=["tensorboard"], + remove_unused_columns=False, + ), + ) + + has_checkpoint = any(Path(out_dir).glob("checkpoint-*")) if Path(out_dir).exists() else False + trainer.train(resume_from_checkpoint=has_checkpoint) + model.save_pretrained(f"{out_dir}/final") + tokenizer.save_pretrained(f"{out_dir}/final") + print(f"Saved to {out_dir}/final") + + +if __name__ == "__main__": + main()