+1 -1

.gitignore

reviewed

··· 1 1 - src/__pycache__/ 1 1 + __pycache__/ 2 2 data/ 3 3 unsloth_compiled_cache/

+3

pyproject.toml

reviewed

··· 10 10 "torch>=2.3", 11 11 "editdistance", 12 12 "tqdm", 13 13 + "numpy", 14 14 + "torchvision>=0.18", 13 15 ] 14 16 15 17 [project.scripts] ··· 20 22 export = "src.export:main" 21 23 generate-typeset = "src.generate_typeset:main" 22 24 probe = "src.probe:main" 25 25 + eff-mer-evaluate = "src.eff_mer.infer:main" 23 26 24 27 [build-system] 25 28 requires = ["hatchling"]

src/eff_mer/__init__.py

reviewed

This is a binary file and will not be displayed.

+160

src/eff_mer/data.py

reviewed

··· 1 1 + """ 2 2 + Dataset for math expression recognition. 3 3 + 4 4 + Each sample: (image_tensor, token_ids) 5 5 + image_tensor: float32 (1, H, W), normalized to ImageNet grayscale stats 6 6 + token_ids: int64 [<sos>, tok..., <eos>], length <= max_len 7 7 + 8 8 + Images are scaled to a fixed height (IMG_H) preserving aspect ratio, then 9 9 + width is capped at MAX_W and padded to a multiple of STRIDE. 10 10 + 11 11 + collate_fn pads images to the batch-max width and returns a spatial memory 12 12 + mask (B, H_tok * W_tok) -- True at padding token positions -- alongside the 13 13 + usual token padding mask. 14 14 + """ 15 15 + 16 16 + import random 17 17 + from pathlib import Path 18 18 + 19 19 + import torch 20 20 + import torchvision.transforms.functional as TF 21 21 + from PIL import Image, ImageFilter, ImageOps 22 22 + from torch.utils.data import Dataset 23 23 + 24 24 + from .vocab import Vocab 25 25 + 26 26 + # ImageNet RGB mean/std averaged across channels, for grayscale input 27 27 + _MEAN = 0.449 28 28 + _STD = 0.226 29 29 + 30 30 + # Image geometry 31 31 + IMG_H = 256 # fixed height; must be a multiple of STRIDE 32 32 + MAX_W = 1024 # maximum width before scaling down 33 33 + STRIDE = 32 # EfficientNet-B0 spatial stride 34 34 + 35 35 + 36 36 + def scale_and_pad(img: Image.Image, h: int = IMG_H, max_w: int = MAX_W) -> Image.Image: 37 37 + """ 38 38 + Scale image to fixed height h preserving aspect ratio. 39 39 + If resulting width exceeds max_w, scale down further. 40 40 + Pad width to the nearest multiple of STRIDE with white. 41 41 + """ 42 42 + orig_w, orig_h = img.size 43 43 + scale = h / orig_h 44 44 + new_w = round(orig_w * scale) 45 45 + if new_w > max_w: 46 46 + scale = max_w / orig_w 47 47 + new_h = round(orig_h * scale) 48 48 + new_w = max_w 49 49 + else: 50 50 + new_h = h 51 51 + img = img.resize((new_w, new_h), Image.LANCZOS) 52 52 + # Pad height to exactly h (handles rounding), width to multiple of STRIDE 53 53 + pad_w = (-new_w % STRIDE) or 0 54 54 + if new_h != h or pad_w: 55 55 + img = ImageOps.expand(img, (0, 0, pad_w, h - new_h), fill=255) 56 56 + return img 57 57 + 58 58 + 59 59 + def _augment(img: Image.Image) -> Image.Image: 60 60 + """ 61 61 + Online augmentation for training images. 62 62 + Random affine (rotation, scale, translate) + blur + pixel noise. 63 63 + All parameters are small to preserve legibility. 64 64 + """ 65 65 + angle = random.uniform(-5, 5) 66 66 + scale = random.uniform(0.9, 1.1) 67 67 + tx = int(random.uniform(-0.05, 0.05) * img.width) 68 68 + ty = int(random.uniform(-0.05, 0.05) * img.height) 69 69 + img = TF.affine(img, angle=angle, translate=(tx, ty), scale=scale, shear=0, fill=255) 70 70 + 71 71 + radius = random.uniform(0.0, 0.8) 72 72 + if radius > 0.15: 73 73 + img = img.filter(ImageFilter.GaussianBlur(radius)) 74 74 + 75 75 + return img 76 76 + 77 77 + 78 78 + def _spatial_mask(w_toks: list[int], H_tok: int, W_tok_max: int) -> torch.Tensor: 79 79 + """ 80 80 + Build memory key padding mask (B, H_tok * W_tok_max). 81 81 + True at positions corresponding to width-padding columns. 82 82 + """ 83 83 + B = len(w_toks) 84 84 + mask = torch.zeros(B, H_tok, W_tok_max, dtype=torch.bool) 85 85 + for i, wt in enumerate(w_toks): 86 86 + if wt < W_tok_max: 87 87 + mask[i, :, wt:] = True 88 88 + return mask.reshape(B, H_tok * W_tok_max) 89 89 + 90 90 + 91 91 + class MathDataset(Dataset): 92 92 + def __init__( 93 93 + self, 94 94 + manifest_paths: list[Path], 95 95 + vocab: Vocab, 96 96 + max_len: int = 256, 97 97 + augment: bool = False, 98 98 + ): 99 99 + self.vocab = vocab 100 100 + self.max_len = max_len 101 101 + self.augment = augment 102 102 + self.records: list[tuple[Path, str]] = [] 103 103 + 104 104 + for mp in manifest_paths: 105 105 + base = mp.parent 106 106 + for line in mp.open(): 107 107 + r = __import__("json").loads(line) 108 108 + typst = r.get("typst", "") 109 109 + if not typst or typst.startswith("ERROR:"): 110 110 + continue 111 111 + self.records.append((base / r["image"], typst)) 112 112 + 113 113 + def __len__(self) -> int: 114 114 + return len(self.records) 115 115 + 116 116 + def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]: 117 117 + img_path, typst = self.records[idx] 118 118 + 119 119 + img = Image.open(img_path).convert("L") 120 120 + if self.augment: 121 121 + img = _augment(img) 122 122 + img = scale_and_pad(img) 123 123 + img_t = TF.to_tensor(img) # (1, H, W), float32 in [0, 1] 124 124 + if self.augment: 125 125 + img_t = (img_t + torch.randn_like(img_t) * 0.02).clamp(0.0, 1.0) 126 126 + img_t = TF.normalize(img_t, [_MEAN], [_STD]) 127 127 + 128 128 + ids = [self.vocab.sos_id] + self.vocab.encode(typst) + [self.vocab.eos_id] 129 129 + ids = ids[: self.max_len] 130 130 + 131 131 + return img_t, torch.tensor(ids, dtype=torch.long) 132 132 + 133 133 + @staticmethod 134 134 + def collate(batch: list[tuple[torch.Tensor, torch.Tensor]]): 135 135 + imgs, seqs = zip(*batch) 136 136 + 137 137 + # Pad images to batch-max width 138 138 + H = imgs[0].shape[1] 139 139 + W_max = max(t.shape[2] for t in imgs) 140 140 + imgs_t = torch.zeros(len(imgs), 1, H, W_max) 141 141 + for i, t in enumerate(imgs): 142 142 + imgs_t[i, :, :, : t.shape[2]] = t 143 143 + 144 144 + # Spatial memory mask: True at width-padding token positions 145 145 + H_tok = H // STRIDE 146 146 + W_tok_max = W_max // STRIDE 147 147 + w_toks = [t.shape[2] // STRIDE for t in imgs] 148 148 + mem_mask = _spatial_mask(w_toks, H_tok, W_tok_max) # (B, H_tok*W_tok_max) 149 149 + 150 150 + # Token sequence padding 151 151 + max_len = max(s.shape[0] for s in seqs) 152 152 + padded = torch.zeros(len(seqs), max_len, dtype=torch.long) 153 153 + lengths = torch.zeros(len(seqs), dtype=torch.long) 154 154 + for i, s in enumerate(seqs): 155 155 + padded[i, : s.shape[0]] = s 156 156 + lengths[i] = s.shape[0] 157 157 + 158 158 + pad_mask = padded == 0 # (B, T) 159 159 + 160 160 + return imgs_t, padded, lengths, pad_mask, mem_mask

+221

src/eff_mer/infer.py

reviewed

··· 1 1 + """ 2 2 + Inference and evaluation for MathOCR. 3 3 + 4 4 + Evaluate ExpRate on test splits: 5 5 + uv run python scripts/infer.py --checkpoint checkpoints/epoch_009.pt 6 6 + 7 7 + Run on arbitrary image files: 8 8 + uv run python scripts/infer.py --checkpoint checkpoints/epoch_009.pt --images a.png b.png 9 9 + 10 10 + Both modes can be combined. 11 11 + """ 12 12 + 13 13 + import argparse 14 14 + import json 15 15 + from pathlib import Path 16 16 + 17 17 + import numpy as np 18 18 + import torch 19 19 + import torchvision.transforms.functional as TF 20 20 + from PIL import Image, ImageOps 21 21 + from tqdm import tqdm 22 22 + 23 23 + # eff-mer data lives in the sibling eff-mer repo 24 24 + _EFF_MER_ROOT = Path(__file__).parent.parent.parent.parent / "eff-mer" 25 25 + 26 26 + from .data import _MEAN, _STD, scale_and_pad 27 27 + from .model import MathOCR 28 28 + from .vocab import Vocab 29 29 + 30 30 + TEST_SPLITS = {"crohme_test", "mathwriting_Test"} 31 31 + 32 32 + 33 33 + def load_model(checkpoint: Path, vocab: Vocab, device: torch.device) -> MathOCR: 34 34 + ckpt = torch.load(checkpoint, map_location=device, weights_only=True) 35 35 + model = MathOCR(vocab_size=len(vocab)).to(device) 36 36 + model.load_state_dict(ckpt["model"]) 37 37 + model.eval() 38 38 + return model 39 39 + 40 40 + 41 41 + def preprocess_photo(img: Image.Image, margin: float = 0.05) -> Image.Image: 42 42 + """ 43 43 + Crop a phone photo to the bounding box of dark (ink) pixels. 44 44 + 45 45 + Steps: grayscale -> Otsu binarize -> bbox of dark pixels -> crop + margin. 46 46 + Works for clean shots of a single expression on blank paper. 47 47 + Not robust to ruled lines -- use --no-preprocess to skip. 48 48 + """ 49 49 + gray = np.array(img.convert("L")) 50 50 + thresh = int(gray.mean()) # simple global threshold; Otsu would need scipy 51 51 + # Iterative Otsu approximation using numpy 52 52 + for _ in range(10): 53 53 + fg = gray[gray < thresh].mean() if (gray < thresh).any() else 0 54 54 + bg = gray[gray >= thresh].mean() if (gray >= thresh).any() else 255 55 55 + thresh = int((fg + bg) / 2) 56 56 + binary = gray < thresh # True = ink 57 57 + 58 58 + rows = np.where(binary.any(axis=1))[0] 59 59 + cols = np.where(binary.any(axis=0))[0] 60 60 + if rows.size == 0 or cols.size == 0: 61 61 + return img # no ink found, return as-is 62 62 + 63 63 + r0, r1 = int(rows[0]), int(rows[-1]) 64 64 + c0, c1 = int(cols[0]), int(cols[-1]) 65 65 + 66 66 + h, w = gray.shape 67 67 + pad_r = max(1, int((r1 - r0) * margin)) 68 68 + pad_c = max(1, int((c1 - c0) * margin)) 69 69 + r0 = max(0, r0 - pad_r) 70 70 + r1 = min(h, r1 + pad_r) 71 71 + c0 = max(0, c0 - pad_c) 72 72 + c1 = min(w, c1 + pad_c) 73 73 + 74 74 + cropped = img.crop((c0, r0, c1, r1)) 75 75 + # Pad to square to preserve aspect ratio before resize 76 76 + side = max(cropped.width, cropped.height) 77 77 + squared = ImageOps.pad(cropped, (side, side), color=255) 78 78 + return squared 79 79 + 80 80 + 81 81 + def binarize(img: Image.Image) -> Image.Image: 82 82 + """Otsu threshold -> pure black-on-white, color/lighting invariant.""" 83 83 + arr = np.array(img, dtype=np.float32) 84 84 + thresh = arr.mean() 85 85 + for _ in range(10): 86 86 + fg = arr[arr < thresh].mean() if (arr < thresh).any() else 0 87 87 + bg = arr[arr >= thresh].mean() if (arr >= thresh).any() else 255 88 88 + thresh = (fg + bg) / 2 89 89 + ink = arr < thresh 90 90 + # Auto-invert if majority of pixels are dark (light-on-dark input) 91 91 + if ink.mean() > 0.5: 92 92 + ink = ~ink 93 93 + return Image.fromarray(np.where(ink, np.uint8(0), np.uint8(255))) 94 94 + 95 95 + 96 96 + def load_image(path: Path, preprocess: bool = False) -> torch.Tensor: 97 97 + img = Image.open(path).convert("L") 98 98 + img = binarize(img) 99 99 + if preprocess: 100 100 + img = preprocess_photo(img) 101 101 + img = scale_and_pad(img) 102 102 + t = TF.to_tensor(img) 103 103 + return TF.normalize(t, [_MEAN], [_STD]) 104 104 + 105 105 + 106 106 + def pad_images(tensors: list[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor | None]: 107 107 + """ 108 108 + Pad a list of (1, H, W) tensors to the batch-max width. 109 109 + Returns (imgs, mem_mask) where mem_mask is None if all widths are equal. 110 110 + """ 111 111 + from eff_mer.data import STRIDE, _spatial_mask 112 112 + H = tensors[0].shape[1] 113 113 + W_max = max(t.shape[2] for t in tensors) 114 114 + out = torch.zeros(len(tensors), 1, H, W_max) 115 115 + for i, t in enumerate(tensors): 116 116 + out[i, :, :, : t.shape[2]] = t 117 117 + w_toks = [t.shape[2] // STRIDE for t in tensors] 118 118 + H_tok = H // STRIDE 119 119 + W_tok_max = W_max // STRIDE 120 120 + if all(w == W_tok_max for w in w_toks): 121 121 + return out, None 122 122 + return out, _spatial_mask(w_toks, H_tok, W_tok_max) 123 123 + 124 124 + 125 125 + def decode_batch(model: MathOCR, imgs: torch.Tensor, vocab: Vocab, 126 126 + mem_mask: torch.Tensor | None = None) -> list[str]: 127 127 + ids_list = model.generate(imgs, vocab.sos_id, vocab.eos_id, 128 128 + memory_key_padding_mask=mem_mask) 129 129 + return [vocab.decode(ids) for ids in ids_list] 130 130 + 131 131 + 132 132 + def evaluate(model: MathOCR, vocab: Vocab, raster: Path, device: torch.device, batch_size: int): 133 133 + manifests = [mp for mp in sorted(raster.glob("*/manifest.jsonl")) 134 134 + if mp.parent.name in TEST_SPLITS] 135 135 + if not manifests: 136 136 + print(f"No test manifests found under {raster}") 137 137 + return 138 138 + 139 139 + for mp in manifests: 140 140 + split = mp.parent.name 141 141 + records = [] 142 142 + for line in mp.open(): 143 143 + r = json.loads(line) 144 144 + typst = r.get("typst", "") 145 145 + if typst and not typst.startswith("ERROR:"): 146 146 + records.append((mp.parent / r["image"], typst)) 147 147 + 148 148 + exact = 0 149 149 + bar = tqdm(range(0, len(records), batch_size), desc=split, leave=True) 150 150 + examples = [] 151 151 + 152 152 + for start in bar: 153 153 + batch = records[start : start + batch_size] 154 154 + imgs, mem_mask = pad_images([load_image(p) for p, _ in batch]) 155 155 + imgs = imgs.to(device) 156 156 + mem_mask = mem_mask.to(device) if mem_mask is not None else None 157 157 + preds = decode_batch(model, imgs, vocab, mem_mask) 158 158 + 159 159 + for (_, gt), pred in zip(batch, preds): 160 160 + gt_norm = vocab.decode(vocab.encode(gt)) 161 161 + if pred == gt_norm: 162 162 + exact += 1 163 163 + if len(examples) < 5: 164 164 + examples.append((gt_norm, pred)) 165 165 + 166 166 + bar.set_postfix(ExpRate=f"{exact/(start+len(batch)):.3f}") 167 167 + 168 168 + exprate = exact / len(records) 169 169 + print(f"\n{split}: ExpRate={exprate:.4f} ({exact}/{len(records)})") 170 170 + print("Examples (gt | pred):") 171 171 + for gt, pred in examples: 172 172 + match = "OK" if gt == pred else "!!" 173 173 + print(f" [{match}] {gt}") 174 174 + if gt != pred: 175 175 + print(f" {pred}") 176 176 + 177 177 + 178 178 + def infer_files(model: MathOCR, vocab: Vocab, paths: list[Path], device: torch.device, preprocess: bool = False): 179 179 + imgs, mem_mask = pad_images([load_image(p, preprocess=preprocess) for p in paths]) 180 180 + imgs = imgs.to(device) 181 181 + mem_mask = mem_mask.to(device) if mem_mask is not None else None 182 182 + preds = decode_batch(model, imgs, vocab, mem_mask) 183 183 + for path, pred in zip(paths, preds): 184 184 + print(f"{path.name}: {pred}") 185 185 + 186 186 + 187 187 + def main(): 188 188 + parser = argparse.ArgumentParser( 189 189 + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter 190 190 + ) 191 191 + parser.add_argument("--checkpoint", type=Path, required=True) 192 192 + parser.add_argument("--vocab", type=Path, default=_EFF_MER_ROOT / "data" / "vocab.json") 193 193 + parser.add_argument("--raster", type=Path, default=_EFF_MER_ROOT / "data" / "raster") 194 194 + parser.add_argument("--batch", type=int, default=64) 195 195 + parser.add_argument("--images", type=Path, nargs="+", metavar="IMG", 196 196 + help="Image files to run inference on") 197 197 + parser.add_argument("--preprocess", action="store_true", 198 198 + help="Binarize + crop to ink bbox before inference (for phone photos)") 199 199 + parser.add_argument("--eval", action="store_true", default=True, 200 200 + help="Evaluate ExpRate on test splits (default: on unless --images only)") 201 201 + parser.add_argument("--no-eval", dest="eval", action="store_false") 202 202 + args = parser.parse_args() 203 203 + 204 204 + # If --images given without explicit --eval, skip test-set eval 205 205 + if args.images and "--eval" not in sys.argv: 206 206 + args.eval = False 207 207 + 208 208 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 209 209 + vocab = Vocab.load(args.vocab) 210 210 + model = load_model(args.checkpoint, vocab, device) 211 211 + print(f"Loaded {args.checkpoint.name} | vocab={len(vocab)} | device={device}") 212 212 + 213 213 + if args.images: 214 214 + infer_files(model, vocab, args.images, device, preprocess=args.preprocess) 215 215 + 216 216 + if args.eval: 217 217 + evaluate(model, vocab, args.raster, device, args.batch) 218 218 + 219 219 + 220 220 + if __name__ == "__main__": 221 221 + main()

+83

src/eff_mer/model/__init__.py

reviewed

··· 1 1 + """ 2 2 + MathOCR: encoder-decoder model for math expression recognition. 3 3 + 4 4 + images -> Encoder -> memory -> Decoder -> logits 5 5 + -> greedy/beam generate 6 6 + """ 7 7 + 8 8 + import torch 9 9 + import torch.nn as nn 10 10 + 11 11 + from .decoder import Decoder 12 12 + from .encoder import Encoder 13 13 + 14 14 + 15 15 + class MathOCR(nn.Module): 16 16 + def __init__( 17 17 + self, 18 18 + vocab_size: int, 19 19 + d_model: int = 256, 20 20 + nhead: int = 8, 21 21 + num_decoder_layers: int = 3, 22 22 + dropout: float = 0.1, 23 23 + max_len: int = 256, 24 24 + ): 25 25 + super().__init__() 26 26 + self.encoder = Encoder(d_model) 27 27 + self.decoder = Decoder(vocab_size, d_model, nhead, num_decoder_layers, dropout, max_len) 28 28 + 29 29 + def forward( 30 30 + self, 31 31 + images: torch.Tensor, 32 32 + tgt: torch.Tensor, 33 33 + tgt_padding_mask: torch.Tensor | None = None, 34 34 + memory_key_padding_mask: torch.Tensor | None = None, 35 35 + ) -> torch.Tensor: 36 36 + """ 37 37 + images: (B, 1, H, W) 38 38 + tgt: (B, T) teacher-forced token ids 39 39 + tgt_padding_mask: (B, T) True at pad positions 40 40 + memory_key_padding_mask: (B, S) True at spatial padding token positions 41 41 + 42 42 + Returns logits (B, T, vocab_size). 43 43 + """ 44 44 + memory = self.encoder(images) 45 45 + return self.decoder( 46 46 + tgt, memory, 47 47 + tgt_padding_mask=tgt_padding_mask, 48 48 + memory_padding_mask=memory_key_padding_mask, 49 49 + ) 50 50 + 51 51 + @torch.no_grad() 52 52 + def generate( 53 53 + self, 54 54 + images: torch.Tensor, 55 55 + sos_id: int, 56 56 + eos_id: int, 57 57 + max_len: int = 256, 58 58 + memory_key_padding_mask: torch.Tensor | None = None, 59 59 + ) -> list[list[int]]: 60 60 + """Greedy decoding. Returns list of token id lists (eos stripped).""" 61 61 + memory = self.encoder(images) 62 62 + B = images.shape[0] 63 63 + 64 64 + seqs = torch.full((B, 1), sos_id, dtype=torch.long, device=images.device) 65 65 + done = torch.zeros(B, dtype=torch.bool, device=images.device) 66 66 + 67 67 + for _ in range(max_len): 68 68 + logits = self.decoder(seqs, memory, 69 69 + memory_padding_mask=memory_key_padding_mask) # (B, T, V) 70 70 + next_tok = logits[:, -1, :].argmax(-1) # (B,) 71 71 + next_tok = torch.where(done, torch.full_like(next_tok, eos_id), next_tok) 72 72 + seqs = torch.cat([seqs, next_tok.unsqueeze(1)], dim=1) 73 73 + done = done | (next_tok == eos_id) 74 74 + if done.all(): 75 75 + break 76 76 + 77 77 + results = [] 78 78 + for i in range(B): 79 79 + toks = seqs[i, 1:].tolist() # strip <sos> 80 80 + if eos_id in toks: 81 81 + toks = toks[: toks.index(eos_id)] 82 82 + results.append(toks) 83 83 + return results

+79

src/eff_mer/model/decoder.py

reviewed

··· 1 1 + """ 2 2 + Autoregressive transformer decoder. 3 3 + 4 4 + Uses pre-norm (norm_first=True) for training stability, learned positional 5 5 + embeddings, and a causal mask over the target sequence. 6 6 + 7 7 + Coverage attention is a planned addition (see CoMER/WAP) -- the spatial 8 8 + coverage accumulator should be added to cross-attention logits to prevent 9 9 + repeated attendance to the same encoder positions. Deferred to post-baseline. 10 10 + """ 11 11 + 12 12 + import torch 13 13 + import torch.nn as nn 14 14 + 15 15 + 16 16 + class Decoder(nn.Module): 17 17 + def __init__( 18 18 + self, 19 19 + vocab_size: int, 20 20 + d_model: int = 256, 21 21 + nhead: int = 8, 22 22 + num_layers: int = 3, 23 23 + dropout: float = 0.1, 24 24 + max_len: int = 256, 25 25 + ): 26 26 + super().__init__() 27 27 + self.d_model = d_model 28 28 + 29 29 + self.embed = nn.Embedding(vocab_size, d_model, padding_idx=0) 30 30 + self.pos_embed = nn.Embedding(max_len, d_model) 31 31 + 32 32 + layer = nn.TransformerDecoderLayer( 33 33 + d_model=d_model, 34 34 + nhead=nhead, 35 35 + dim_feedforward=d_model * 4, 36 36 + dropout=dropout, 37 37 + batch_first=True, 38 38 + norm_first=True, # pre-norm: more stable training 39 39 + ) 40 40 + self.transformer = nn.TransformerDecoder(layer, num_layers=num_layers) 41 41 + self.out_proj = nn.Linear(d_model, vocab_size) 42 42 + 43 43 + self._init_weights() 44 44 + 45 45 + def _init_weights(self): 46 46 + nn.init.normal_(self.embed.weight, std=0.02) 47 47 + nn.init.normal_(self.pos_embed.weight, std=0.02) 48 48 + nn.init.zeros_(self.out_proj.bias) 49 49 + 50 50 + def forward( 51 51 + self, 52 52 + tgt: torch.Tensor, 53 53 + memory: torch.Tensor, 54 54 + tgt_padding_mask: torch.Tensor | None = None, 55 55 + memory_padding_mask: torch.Tensor | None = None, 56 56 + ) -> torch.Tensor: 57 57 + """ 58 58 + tgt: (B, T) token ids (teacher-forced) 59 59 + memory: (B, S, d) encoder output 60 60 + tgt_padding_mask: (B, T) True at pad positions 61 61 + memory_padding_mask: (B, S) True at pad positions (unused for fixed grids) 62 62 + 63 63 + Returns logits (B, T, vocab_size). 64 64 + """ 65 65 + B, T = tgt.shape 66 66 + causal_mask = nn.Transformer.generate_square_subsequent_mask(T, device=tgt.device, dtype=torch.bool) 67 67 + 68 68 + pos = torch.arange(T, device=tgt.device) 69 69 + x = self.embed(tgt) + self.pos_embed(pos) # (B, T, d_model) 70 70 + 71 71 + x = self.transformer( 72 72 + x, 73 73 + memory, 74 74 + tgt_mask=causal_mask, 75 75 + tgt_key_padding_mask=tgt_padding_mask, 76 76 + memory_key_padding_mask=memory_padding_mask, 77 77 + tgt_is_causal=True, 78 78 + ) 79 79 + return self.out_proj(x) # (B, T, vocab_size)

+59

src/eff_mer/model/encoder.py

reviewed

··· 1 1 + """ 2 2 + Encoder: EfficientNet-B0 backbone + 1x1 projection + 2D sinusoidal positional encoding. 3 3 + 4 4 + Input: (B, 1, H, W) grayscale, normalized to ImageNet grayscale stats 5 5 + Output: (B, h*w, d_model) sequence of spatial tokens 6 6 + 7 7 + At 256x256 input, EfficientNet-B0 produces an 8x8 feature grid (stride 32), 8 8 + giving 64 spatial tokens per image. 9 9 + """ 10 10 + 11 11 + import math 12 12 + 13 13 + import torch 14 14 + import torch.nn as nn 15 15 + import torchvision.models as tvm 16 16 + 17 17 + 18 18 + class Encoder(nn.Module): 19 19 + def __init__(self, d_model: int = 256): 20 20 + super().__init__() 21 21 + assert d_model % 4 == 0, "d_model must be divisible by 4 for 2D sinusoidal PE" 22 22 + 23 23 + backbone = tvm.efficientnet_b0(weights=tvm.EfficientNet_B0_Weights.IMAGENET1K_V1) 24 24 + self.features = backbone.features # (B, 1280, h, w) 25 25 + 26 26 + self.proj = nn.Conv2d(1280, d_model, kernel_size=1, bias=False) 27 27 + self.d_model = d_model 28 28 + 29 29 + def _sinusoidal_pe_2d(self, H: int, W: int, device: torch.device) -> torch.Tensor: 30 30 + """Returns (1, d_model, H, W) positional encoding.""" 31 31 + half_d = self.d_model // 4 # quarter for each of sin/cos × row/col 32 32 + inv_freq = 1.0 / (10000 ** (torch.arange(half_d, device=device).float() / half_d)) 33 33 + 34 34 + row = torch.arange(H, device=device).float() 35 35 + col = torch.arange(W, device=device).float() 36 36 + 37 37 + row_enc = torch.cat([torch.sin(row[:, None] * inv_freq), 38 38 + torch.cos(row[:, None] * inv_freq)], dim=-1) # (H, d/2) 39 39 + col_enc = torch.cat([torch.sin(col[:, None] * inv_freq), 40 40 + torch.cos(col[:, None] * inv_freq)], dim=-1) # (W, d/2) 41 41 + 42 42 + pe = torch.cat([ 43 43 + row_enc.unsqueeze(1).expand(H, W, -1), # (H, W, d/2) 44 44 + col_enc.unsqueeze(0).expand(H, W, -1), # (H, W, d/2) 45 45 + ], dim=-1) # (H, W, d_model) 46 46 + 47 47 + return pe.permute(2, 0, 1).unsqueeze(0) # (1, d_model, H, W) 48 48 + 49 49 + def forward(self, x: torch.Tensor) -> torch.Tensor: 50 50 + # Grayscale -> 3-channel (EfficientNet expects RGB) 51 51 + x3 = x.expand(-1, 3, -1, -1) 52 52 + 53 53 + feat = self.features(x3) # (B, 1280, h, w) 54 54 + feat = self.proj(feat) # (B, d_model, h, w) 55 55 + 56 56 + _, _, H, W = feat.shape 57 57 + feat = feat + self._sinusoidal_pe_2d(H, W, feat.device) 58 58 + 59 59 + return feat.flatten(2).transpose(1, 2) # (B, h*w, d_model)

src/eff_mer/scripts/__init__.py

reviewed

This is a binary file and will not be displayed.

+125

src/eff_mer/scripts/convert_labels.py

reviewed

··· 1 1 + """ 2 2 + Convert LaTeX labels in raster manifests to Typst using latex2typst-batch. 3 3 + 4 4 + Reads each manifest.jsonl, pipes LaTeX strings through the binary, and writes 5 5 + a new manifest.jsonl with an added "typst" field. Records that fail conversion 6 6 + are dropped. Processes manifests in-place (originals backed up as .bak). 7 7 + 8 8 + Usage: 9 9 + uv run convert-labels # all splits under data/raster 10 10 + uv run convert-labels --raster data/raster # explicit raster root 11 11 + uv run convert-labels --binary ../latex2typst-batch 12 12 + """ 13 13 + 14 14 + import argparse 15 15 + import json 16 16 + import re 17 17 + import shlex 18 18 + import shutil 19 19 + import subprocess 20 20 + from pathlib import Path 21 21 + 22 22 + DATA_ROOT = Path(__file__).parent.parent.parent.parent / "data" 23 23 + _PROJECT_ROOT = Path(__file__).parent.parent.parent.parent 24 24 + DEFAULT_CMD = f"node {_PROJECT_ROOT / 'tex2typst-js-batch' / 'batch.mjs'}" 25 25 + 26 26 + _DELIM_RE = re.compile(r"^\$(.+)\$$|^\\\((.+)\\\)$|^\\\[(.+)\\\]$", re.DOTALL) 27 27 + 28 28 + 29 29 + def strip_delimiters(s: str) -> str: 30 30 + """Strip LaTeX/Typst math mode delimiters ($...$, \\(...\\), \\[...\\]).""" 31 31 + s = s.strip() 32 32 + m = _DELIM_RE.match(s) 33 33 + if m: 34 34 + return next(g for g in m.groups() if g is not None).strip() 35 35 + return s 36 36 + 37 37 + 38 38 + def convert_manifest(manifest_path: Path, binary: str) -> tuple[int, int]: 39 39 + """ 40 40 + Convert a single manifest.jsonl in-place. 41 41 + Returns (written, dropped) counts. 42 42 + """ 43 43 + records = [json.loads(line) for line in manifest_path.open()] 44 44 + if not records: 45 45 + return 0, 0 46 46 + 47 47 + # Skip if already converted 48 48 + if "typst" in records[0]: 49 49 + print(f" [skip] already converted: {manifest_path}") 50 50 + return 0, 0 51 51 + 52 52 + latexes = [strip_delimiters(r["latex"]) for r in records] 53 53 + batch_input = "\n".join(latexes) 54 54 + 55 55 + result = subprocess.run( 56 56 + shlex.split(binary), 57 57 + input=batch_input, 58 58 + capture_output=True, 59 59 + text=True, 60 60 + ) 61 61 + typst_lines = result.stdout.splitlines() 62 62 + 63 63 + if len(typst_lines) != len(records): 64 64 + raise RuntimeError( 65 65 + f"{manifest_path}: expected {len(records)} output lines, " 66 66 + f"got {len(typst_lines)}" 67 67 + ) 68 68 + 69 69 + # Backup original 70 70 + bak = manifest_path.with_suffix(".jsonl.bak") 71 71 + shutil.copy2(manifest_path, bak) 72 72 + 73 73 + written = dropped = 0 74 74 + with manifest_path.open("w") as f: 75 75 + for rec, typst in zip(records, typst_lines): 76 76 + if typst.startswith("ERROR:"): 77 77 + dropped += 1 78 78 + continue 79 79 + # Converter re-wraps some outputs in $...$; strip them 80 80 + typst = strip_delimiters(typst) 81 81 + rec["typst"] = typst 82 82 + f.write(json.dumps(rec) + "\n") 83 83 + written += 1 84 84 + 85 85 + return written, dropped 86 86 + 87 87 + 88 88 + def main(): 89 89 + parser = argparse.ArgumentParser( 90 90 + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter 91 91 + ) 92 92 + parser.add_argument( 93 93 + "--raster", 94 94 + type=Path, 95 95 + default=DATA_ROOT / "raster", 96 96 + help="Raster root containing split subdirectories (default: data/raster)", 97 97 + ) 98 98 + parser.add_argument( 99 99 + "--binary", 100 100 + type=str, 101 101 + default=DEFAULT_CMD, 102 102 + help="Converter command (shell-split), default: node tex2typst-js-batch/batch.mjs", 103 103 + ) 104 104 + args = parser.parse_args() 105 105 + 106 106 + manifests = sorted(args.raster.glob("*/manifest.jsonl")) 107 107 + if not manifests: 108 108 + parser.error(f"No manifest.jsonl files found under {args.raster}") 109 109 + 110 110 + total_written = total_dropped = 0 111 111 + for mp in manifests: 112 112 + tag = mp.parent.name 113 113 + print(f"[{tag}]") 114 114 + written, dropped = convert_manifest(mp, args.binary) 115 115 + print(f" {written:,} kept, {dropped:,} dropped") 116 116 + total_written += written 117 117 + total_dropped += dropped 118 118 + 119 119 + total = total_written + total_dropped 120 120 + drop_pct = 100 * total_dropped / total if total else 0 121 121 + print(f"\nTotal: {total_written:,} kept, {total_dropped:,} dropped ({drop_pct:.1f}%)") 122 122 + 123 123 + 124 124 + if __name__ == "__main__": 125 125 + main()

+306

src/eff_mer/scripts/rasterize.py

reviewed

··· 1 1 + #!/usr/bin/env python3 2 2 + """ 3 3 + Rasterize InkML math expression datasets to PNG images + JSONL manifest. 4 4 + 5 5 + Supported datasets: 6 6 + crohme splits: real_train, val, test, gen_LaTeX_data_CROHME_2019, 7 7 + gen_LaTeX_data_CROHME_2023_corpus, gen_syntactic_data 8 8 + mathwriting splits: Train, Val, Test, Symbols, Synthetic 9 9 + typeset [stub] not yet implemented 10 10 + 11 11 + Output structure: 12 12 + <out>/<dataset>_<split>/ 13 13 + images/ *.png 14 14 + manifest.jsonl {"image": "images/<stem>.png", "latex": "...", "mathml": "..."} 15 15 + 16 16 + Usage: 17 17 + python scripts/rasterize.py --dataset crohme --split real_train 18 18 + python scripts/rasterize.py --dataset mathwriting --split Train --augment 19 19 + python scripts/rasterize.py --all 20 20 + python scripts/rasterize.py --all --size 128 --out data/raster_128 21 21 + """ 22 22 + 23 23 + import argparse 24 24 + import json 25 25 + import multiprocessing as mp 26 26 + import random 27 27 + from pathlib import Path 28 28 + from xml.etree import ElementTree as ET 29 29 + 30 30 + import numpy as np 31 31 + from PIL import Image, ImageDraw, ImageFilter 32 32 + 33 33 + # ── Config ──────────────────────────────────────────────────────────────────── 34 34 + 35 35 + DATA_ROOT = Path(__file__).parent.parent.parent.parent / "data" 36 36 + 37 37 + CROHME_SPLITS = [ 38 38 + "real_train", 39 39 + "val", 40 40 + "test", 41 41 + "gen_LaTeX_data_CROHME_2019", 42 42 + "gen_LaTeX_data_CROHME_2023_corpus", 43 43 + "gen_syntactic_data", 44 44 + ] 45 45 + 46 46 + MATHWRITING_SPLITS = ["Train", "Val", "Test", "Symbols", "Synthetic"] 47 47 + 48 48 + INK_NS = "http://www.w3.org/2003/InkML" 49 49 + NS = {"ink": INK_NS} 50 50 + 51 51 + # Fraction of canvas size used as padding on each side 52 52 + MARGIN_FRAC = 0.08 53 53 + 54 54 + # ── InkML parsing ───────────────────────────────────────────────────────────── 55 55 + 56 56 + 57 57 + def _find_annotation(root: ET.Element, *type_keys: str) -> str | None: 58 58 + for key in type_keys: 59 59 + el = root.find(f"{{%s}}annotation[@type='%s']" % (INK_NS, key)) 60 60 + if el is not None and el.text: 61 61 + return el.text.strip() 62 62 + return None 63 63 + 64 64 + 65 65 + def parse_inkml(path: Path) -> dict | None: 66 66 + """ 67 67 + Returns dict with: 68 68 + latex (str) ground-truth LaTeX string 69 69 + mathml (str) serialized MathML annotationXML block, or "" 70 70 + strokes (list of np.ndarray shape (N, 2), float32 xy) 71 71 + Returns None on parse error or missing label. 72 72 + """ 73 73 + try: 74 74 + tree = ET.parse(path) 75 75 + except ET.ParseError: 76 76 + return None 77 77 + 78 78 + root = tree.getroot() 79 79 + 80 80 + # CROHME+ uses "truth"; MathWriting+ uses "label" / "normalizedLabel" 81 81 + latex = _find_annotation(root, "truth", "label", "normalizedLabel") 82 82 + if latex is None: 83 83 + return None 84 84 + 85 85 + # MathML tree: serialize <annotationXML type="truth"> for later use by tree head 86 86 + mathml_el = root.find(f"{{{INK_NS}}}annotationXML[@type='truth']") 87 87 + mathml_str = ET.tostring(mathml_el, encoding="unicode") if mathml_el is not None else "" 88 88 + 89 89 + # Collect strokes: map trace id -> xy array (ignore timestamps / pressure) 90 90 + traces: dict[str, np.ndarray] = {} 91 91 + for trace in root.findall(f"{{{INK_NS}}}trace"): 92 92 + tid = trace.get("id", "") 93 93 + pts = [] 94 94 + if trace.text: 95 95 + for token in trace.text.strip().split(","): 96 96 + coords = token.strip().split() 97 97 + if len(coords) >= 2: 98 98 + try: 99 99 + pts.append((float(coords[0]), float(coords[1]))) 100 100 + except ValueError: 101 101 + continue 102 102 + if pts: 103 103 + traces[tid] = np.array(pts, dtype=np.float32) 104 104 + 105 105 + strokes = list(traces.values()) 106 106 + if not strokes: 107 107 + return None 108 108 + 109 109 + return {"latex": latex, "mathml": mathml_str, "strokes": strokes} 110 110 + 111 111 + 112 112 + # ── Rendering ───────────────────────────────────────────────────────────────── 113 113 + 114 114 + 115 115 + def render_strokes( 116 116 + strokes: list[np.ndarray], 117 117 + size: int = 256, 118 118 + pen_width: float = 2.0, 119 119 + augment: bool = False, 120 120 + ) -> Image.Image: 121 121 + """Render stroke list to a grayscale PIL image (white bg, black ink).""" 122 122 + all_pts = np.concatenate(strokes, axis=0) 123 123 + xmin, ymin = float(all_pts[:, 0].min()), float(all_pts[:, 1].min()) 124 124 + xmax, ymax = float(all_pts[:, 0].max()), float(all_pts[:, 1].max()) 125 125 + 126 126 + span = max(xmax - xmin, ymax - ymin, 1.0) 127 127 + margin_px = int(size * MARGIN_FRAC) 128 128 + canvas_px = size - 2 * margin_px 129 129 + 130 130 + if augment: 131 131 + pen_width = pen_width * random.uniform(0.7, 1.6) 132 132 + 133 133 + def to_px(pts: np.ndarray) -> list[tuple[int, int]]: 134 134 + x = ((pts[:, 0] - xmin) / span * canvas_px + margin_px).round().astype(int) 135 135 + y = ((pts[:, 1] - ymin) / span * canvas_px + margin_px).round().astype(int) 136 136 + return list(zip(x.tolist(), y.tolist())) 137 137 + 138 138 + img = Image.new("L", (size, size), color=255) 139 139 + draw = ImageDraw.Draw(img) 140 140 + pw = max(1, round(pen_width)) 141 141 + 142 142 + for stroke in strokes: 143 143 + px = to_px(stroke) 144 144 + if len(px) == 1: 145 145 + draw.ellipse([px[0][0] - pw, px[0][1] - pw, px[0][0] + pw, px[0][1] + pw], fill=0) 146 146 + else: 147 147 + draw.line(px, fill=0, width=pw) 148 148 + 149 149 + if augment: 150 150 + # Soften jagged strokes 151 151 + img = img.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.0, 0.5))) 152 152 + # Light pixel noise to simulate paper texture 153 153 + arr = np.array(img, dtype=np.int16) 154 154 + arr = np.clip(arr + np.random.normal(0, 4, arr.shape).astype(np.int16), 0, 255).astype(np.uint8) 155 155 + img = Image.fromarray(arr) 156 156 + 157 157 + return img 158 158 + 159 159 + 160 160 + # ── Worker (top-level for multiprocessing pickling) ─────────────────────────── 161 161 + 162 162 + 163 163 + def _worker(args: tuple) -> dict | None: 164 164 + inkml_path, out_img_dir, size, augment = args 165 165 + parsed = parse_inkml(inkml_path) 166 166 + if parsed is None: 167 167 + return None 168 168 + img = render_strokes(parsed["strokes"], size=size, augment=augment) 169 169 + img_name = inkml_path.stem + ".png" 170 170 + img_path = Path(out_img_dir) / img_name 171 171 + img.save(img_path, optimize=False) 172 172 + return { 173 173 + "image": "images/" + img_name, 174 174 + "latex": parsed["latex"], 175 175 + "mathml": parsed["mathml"], 176 176 + } 177 177 + 178 178 + 179 179 + # ── Split processor ─────────────────────────────────────────────────────────── 180 180 + 181 181 + 182 182 + def process_split( 183 183 + inkml_dir: Path, 184 184 + out_dir: Path, 185 185 + size: int, 186 186 + augment: bool, 187 187 + workers: int, 188 188 + ) -> int: 189 189 + inkml_files = sorted(inkml_dir.glob("*.inkml")) 190 190 + if not inkml_files: 191 191 + print(f" [skip] no .inkml files in {inkml_dir}") 192 192 + return 0 193 193 + 194 194 + out_img_dir = out_dir / "images" 195 195 + out_img_dir.mkdir(parents=True, exist_ok=True) 196 196 + 197 197 + job_args = [(f, str(out_img_dir), size, augment) for f in inkml_files] 198 198 + 199 199 + count = 0 200 200 + skipped = 0 201 201 + with open(out_dir / "manifest.jsonl", "w") as mf: 202 202 + with mp.Pool(workers) as pool: 203 203 + for rec in pool.imap_unordered(_worker, job_args, chunksize=128): 204 204 + if rec is not None: 205 205 + mf.write(json.dumps(rec) + "\n") 206 206 + count += 1 207 207 + else: 208 208 + skipped += 1 209 209 + 210 210 + if skipped: 211 211 + print(f" {count:,} written, {skipped} skipped (parse errors / no label)") 212 212 + return count 213 213 + 214 214 + 215 215 + # ── Typeset stub ────────────────────────────────────────────────────────────── 216 216 + 217 217 + 218 218 + def process_typeset(out_dir: Path, size: int, augment: bool) -> int: 219 219 + """ 220 220 + [STUB] Render typeset math expressions to images. 221 221 + 222 222 + Planned sources: 223 223 + - Typst CLI: generate synthetic expressions from a formal grammar, 224 224 + render via `typst compile` to PNG. 225 225 + - im2latex-100k or similar LaTeX-rendered dataset. 226 226 + 227 227 + Expected manifest format is identical to handwriting splits: 228 228 + {"image": "images/<stem>.png", "latex": "...", "mathml": ""} 229 229 + """ 230 230 + raise NotImplementedError( 231 231 + "Typeset rasterization not yet implemented.\n" 232 232 + "Planned: Typst CLI rendering of grammar-generated expressions.\n" 233 233 + "See scripts/rasterize_typeset.py (to be written)." 234 234 + ) 235 235 + 236 236 + 237 237 + # ── CLI ─────────────────────────────────────────────────────────────────────── 238 238 + 239 239 + 240 240 + def build_todo(args) -> list[tuple[str, str, Path]]: 241 241 + """Returns list of (dataset_tag, split, inkml_dir) tuples.""" 242 242 + if args.all: 243 243 + todo = [] 244 244 + for split in CROHME_SPLITS: 245 245 + d = DATA_ROOT / "CROHME+" / split 246 246 + if d.exists(): 247 247 + todo.append(("crohme", split, d)) 248 248 + else: 249 249 + print(f"[warn] CROHME+ split not found: {d}") 250 250 + for split in MATHWRITING_SPLITS: 251 251 + d = DATA_ROOT / "MathWriting+" / split 252 252 + if d.exists(): 253 253 + todo.append(("mathwriting", split, d)) 254 254 + else: 255 255 + print(f"[warn] MathWriting+ split not found: {d}") 256 256 + return todo 257 257 + 258 258 + if args.dataset == "typeset": 259 259 + return [("typeset", "typeset", Path())] 260 260 + 261 261 + dataset_root = { 262 262 + "crohme": DATA_ROOT / "CROHME+", 263 263 + "mathwriting": DATA_ROOT / "MathWriting+", 264 264 + }[args.dataset] 265 265 + 266 266 + inkml_dir = dataset_root / args.split 267 267 + if not inkml_dir.exists(): 268 268 + raise FileNotFoundError(f"Split directory not found: {inkml_dir}") 269 269 + return [(args.dataset, args.split, inkml_dir)] 270 270 + 271 271 + 272 272 + def main(): 273 273 + parser = argparse.ArgumentParser( 274 274 + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter 275 275 + ) 276 276 + parser.add_argument("--dataset", choices=["crohme", "mathwriting", "typeset"]) 277 277 + parser.add_argument("--split", help="Split name (e.g. real_train, Train)") 278 278 + parser.add_argument("--all", action="store_true", help="Process all known handwriting splits") 279 279 + parser.add_argument("--out", type=Path, default=DATA_ROOT / "raster", help="Output root (default: data/raster)") 280 280 + parser.add_argument("--size", type=int, default=256, help="Square canvas size in pixels (default: 256)") 281 281 + parser.add_argument("--augment", action="store_true", help="Apply pen-width jitter + noise at render time") 282 282 + parser.add_argument("--workers", type=int, default=max(1, mp.cpu_count() - 1)) 283 283 + args = parser.parse_args() 284 284 + 285 285 + if not args.all and args.dataset is None: 286 286 + parser.error("Specify --dataset or --all") 287 287 + if not args.all and args.dataset != "typeset" and args.split is None: 288 288 + parser.error("--split required unless --all or --dataset typeset") 289 289 + 290 290 + todo = build_todo(args) 291 291 + 292 292 + for dataset, split, inkml_dir in todo: 293 293 + tag = f"{dataset}_{split}" 294 294 + out_dir = args.out / tag 295 295 + 296 296 + if dataset == "typeset": 297 297 + process_typeset(out_dir, args.size, args.augment) 298 298 + continue 299 299 + 300 300 + print(f"[{tag}] {inkml_dir} -> {out_dir} (workers={args.workers}, size={args.size})") 301 301 + n = process_split(inkml_dir, out_dir, args.size, args.augment, args.workers) 302 302 + print(f" done: {n:,} images") 303 303 + 304 304 + 305 305 + if __name__ == "__main__": 306 306 + main()

+222

src/eff_mer/train.py

reviewed

··· 1 1 + """ 2 2 + Training loop for MathOCR. 3 3 + 4 4 + Usage: 5 5 + uv run train 6 6 + uv run train --epochs 30 --batch 32 --lr 3e-4 --log-every 100 7 7 + 8 8 + Checkpoints saved to checkpoints/epoch_{N}.pt every epoch. 9 9 + Resumes from latest checkpoint if present. 10 10 + """ 11 11 + 12 12 + import argparse 13 13 + import csv 14 14 + import time 15 15 + from pathlib import Path 16 16 + 17 17 + import torch 18 18 + import torch.nn as nn 19 19 + from torch.utils.data import DataLoader 20 20 + from tqdm import tqdm 21 21 + 22 22 + from .data import MathDataset 23 23 + from .model import MathOCR 24 24 + from .vocab import Vocab 25 25 + 26 26 + DATA_ROOT = Path(__file__).parent.parent.parent / "data" 27 27 + CKPT_DIR = Path(__file__).parent.parent.parent / "checkpoints" 28 28 + 29 29 + TRAIN_SPLITS = { 30 30 + "crohme_gen_LaTeX_data_CROHME_2019", 31 31 + "crohme_gen_LaTeX_data_CROHME_2023_corpus", 32 32 + "crohme_gen_syntactic_data", 33 33 + "crohme_real_train", 34 34 + "mathwriting_Train", 35 35 + "mathwriting_Synthetic", 36 36 + "mathwriting_Symbols", 37 37 + } 38 38 + VAL_SPLITS = {"crohme_val", "mathwriting_Val"} 39 39 + 40 40 + 41 41 + def get_manifests(raster: Path, splits: set[str]) -> list[Path]: 42 42 + return sorted(mp for mp in raster.glob("*/manifest.jsonl") if mp.parent.name in splits) 43 43 + 44 44 + 45 45 + def load_latest_checkpoint( 46 46 + ckpt_dir: Path, 47 47 + model: nn.Module, 48 48 + optimizer: torch.optim.Optimizer, 49 49 + scheduler: torch.optim.lr_scheduler.LRScheduler, 50 50 + ): 51 51 + ckpts = sorted(ckpt_dir.glob("epoch_*.pt")) 52 52 + if not ckpts: 53 53 + return 0 54 54 + ckpt = torch.load(ckpts[-1], weights_only=True) 55 55 + model.load_state_dict(ckpt["model"]) 56 56 + optimizer.load_state_dict(ckpt["optimizer"]) 57 57 + if "scheduler" in ckpt: 58 58 + scheduler.load_state_dict(ckpt["scheduler"]) 59 59 + epoch = ckpt["epoch"] 60 60 + print(f"Resumed from {ckpts[-1].name} (epoch {epoch})") 61 61 + return epoch 62 62 + 63 63 + 64 64 + def run_epoch(model, loader, optimizer, criterion, device, train: bool, log_every: int, epoch: int): 65 65 + model.train(train) 66 66 + total_loss = total_tokens = 0 67 67 + step = 0 68 68 + 69 69 + phase = "train" if train else "val" 70 70 + bar = tqdm(loader, desc=f"Epoch {epoch} {phase}", leave=False, dynamic_ncols=True) 71 71 + 72 72 + with torch.set_grad_enabled(train): 73 73 + for imgs, tgt, lengths, pad_mask, mem_mask in bar: 74 74 + imgs = imgs.to(device) 75 75 + tgt = tgt.to(device) 76 76 + pad_mask = pad_mask.to(device) 77 77 + mem_mask = mem_mask.to(device) 78 78 + 79 79 + tgt_in = tgt[:, :-1] 80 80 + tgt_out = tgt[:, 1:] 81 81 + pad_mask_in = pad_mask[:, :-1] 82 82 + 83 83 + logits = model(imgs, tgt_in, tgt_padding_mask=pad_mask_in, 84 84 + memory_key_padding_mask=mem_mask) 85 85 + 86 86 + loss = criterion( 87 87 + logits.reshape(-1, logits.shape[-1]), 88 88 + tgt_out.reshape(-1), 89 89 + ) 90 90 + 91 91 + if train: 92 92 + optimizer.zero_grad() 93 93 + loss.backward() 94 94 + nn.utils.clip_grad_norm_(model.parameters(), 1.0) 95 95 + optimizer.step() 96 96 + 97 97 + n_tokens = (tgt_out != 0).sum().item() 98 98 + total_loss += loss.item() * n_tokens 99 99 + total_tokens += n_tokens 100 100 + step += 1 101 101 + 102 102 + if train and step % log_every == 0: 103 103 + running = total_loss / total_tokens 104 104 + bar.set_postfix(loss=f"{running:.4f}") 105 105 + 106 106 + bar.close() 107 107 + return total_loss / total_tokens if total_tokens else float("inf") 108 108 + 109 109 + 110 110 + def main(): 111 111 + parser = argparse.ArgumentParser( 112 112 + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter 113 113 + ) 114 114 + parser.add_argument("--raster", type=Path, default=DATA_ROOT / "raster") 115 115 + parser.add_argument("--vocab", type=Path, default=DATA_ROOT / "vocab.json") 116 116 + parser.add_argument("--epochs", type=int, default=30) 117 117 + parser.add_argument("--batch", type=int, default=32) 118 118 + parser.add_argument("--lr", type=float, default=3e-4) 119 119 + parser.add_argument("--workers", type=int, default=4) 120 120 + parser.add_argument("--max-len", type=int, default=256) 121 121 + parser.add_argument("--d-model", type=int, default=256) 122 122 + parser.add_argument("--nhead", type=int, default=8) 123 123 + parser.add_argument("--dec-layers", type=int, default=3) 124 124 + parser.add_argument("--dropout", type=float, default=0.1) 125 125 + parser.add_argument("--log-every", type=int, default=100, help="Log loss every N steps") 126 126 + args = parser.parse_args() 127 127 + 128 128 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 129 129 + print(f"Device: {device}") 130 130 + 131 131 + vocab = Vocab.load(args.vocab) 132 132 + print(f"Vocab: {len(vocab)} tokens") 133 133 + 134 134 + train_manifests = get_manifests(args.raster, TRAIN_SPLITS) 135 135 + val_manifests = get_manifests(args.raster, VAL_SPLITS) 136 136 + print(f"Train manifests: {len(train_manifests)}, Val manifests: {len(val_manifests)}") 137 137 + 138 138 + train_ds = MathDataset(train_manifests, vocab, max_len=args.max_len, augment=True) 139 139 + val_ds = MathDataset(val_manifests, vocab, max_len=args.max_len, augment=False) 140 140 + print(f"Train samples: {len(train_ds):,}, Val samples: {len(val_ds):,}") 141 141 + 142 142 + train_loader = DataLoader( 143 143 + train_ds, 144 144 + batch_size=args.batch, 145 145 + shuffle=True, 146 146 + num_workers=args.workers, 147 147 + collate_fn=MathDataset.collate, 148 148 + pin_memory=device.type == "cuda", 149 149 + ) 150 150 + val_loader = DataLoader( 151 151 + val_ds, 152 152 + batch_size=args.batch, 153 153 + shuffle=False, 154 154 + num_workers=args.workers, 155 155 + collate_fn=MathDataset.collate, 156 156 + pin_memory=device.type == "cuda", 157 157 + ) 158 158 + 159 159 + model = MathOCR( 160 160 + vocab_size=len(vocab), 161 161 + d_model=args.d_model, 162 162 + nhead=args.nhead, 163 163 + num_decoder_layers=args.dec_layers, 164 164 + dropout=args.dropout, 165 165 + max_len=args.max_len, 166 166 + ).to(device) 167 167 + 168 168 + n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) 169 169 + print(f"Trainable params: {n_params/1e6:.2f}M") 170 170 + 171 171 + optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) 172 172 + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( 173 173 + optimizer, T_max=args.epochs, eta_min=args.lr * 0.01 174 174 + ) 175 175 + 176 176 + criterion = nn.CrossEntropyLoss(ignore_index=vocab.pad_id, label_smoothing=0.1) 177 177 + 178 178 + CKPT_DIR.mkdir(exist_ok=True) 179 179 + start_epoch = load_latest_checkpoint(CKPT_DIR, model, optimizer, scheduler) 180 180 + 181 181 + log_path = CKPT_DIR / "metrics.csv" 182 182 + log_exists = log_path.exists() 183 183 + log_file = log_path.open("a", newline="") 184 184 + log_writer = csv.writer(log_file) 185 185 + if not log_exists: 186 186 + log_writer.writerow(["epoch", "train_loss", "val_loss", "lr", "elapsed_s"]) 187 187 + 188 188 + for epoch in range(start_epoch + 1, args.epochs + 1): 189 189 + t0 = time.time() 190 190 + train_loss = run_epoch(model, train_loader, optimizer, criterion, device, True, args.log_every, epoch) 191 191 + scheduler.step() 192 192 + val_loss = run_epoch(model, val_loader, optimizer, criterion, device, False, args.log_every, epoch) 193 193 + elapsed = time.time() - t0 194 194 + lr = optimizer.param_groups[0]["lr"] 195 195 + 196 196 + print( 197 197 + f"Epoch {epoch:3d}/{args.epochs} | " 198 198 + f"train={train_loss:.4f} | val={val_loss:.4f} | " 199 199 + f"lr={lr:.2e} | {elapsed:.0f}s" 200 200 + ) 201 201 + 202 202 + ckpt_path = CKPT_DIR / f"epoch_{epoch:03d}.pt" 203 203 + torch.save( 204 204 + { 205 205 + "epoch": epoch, 206 206 + "model": model.state_dict(), 207 207 + "optimizer": optimizer.state_dict(), 208 208 + "scheduler": scheduler.state_dict(), 209 209 + "val_loss": val_loss, 210 210 + }, 211 211 + ckpt_path, 212 212 + ) 213 213 + print(f" checkpoint -> {ckpt_path.name}") 214 214 + 215 215 + log_writer.writerow([epoch, f"{train_loss:.6f}", f"{val_loss:.6f}", f"{lr:.6e}", f"{elapsed:.1f}"]) 216 216 + log_file.flush() 217 217 + 218 218 + log_file.close() 219 219 + 220 220 + 221 221 + if __name__ == "__main__": 222 222 + main()

+67

src/eff_mer/vocab.py

reviewed

··· 1 1 + """ 2 2 + Fixed vocabulary over Typst math tokens. 3 3 + 4 4 + Tokenization: regex split into word-character runs and individual symbols, 5 5 + matching how Typst math is written (multi-char identifiers like `frac`, `sqrt` 6 6 + stay atomic; operators like `+`, `(`, `^` are single tokens). 7 7 + """ 8 8 + 9 9 + import json 10 10 + import re 11 11 + from collections import Counter 12 12 + from pathlib import Path 13 13 + 14 14 + SPECIAL = ["<pad>", "<sos>", "<eos>", "<unk>"] 15 15 + 16 16 + _TOKEN_RE = re.compile(r"[a-zA-Z0-9]+|[^\w\s]|_") 17 17 + 18 18 + 19 19 + def tokenize(typst: str) -> list[str]: 20 20 + return _TOKEN_RE.findall(typst) 21 21 + 22 22 + 23 23 + class Vocab: 24 24 + def __init__(self, tokens: list[str]): 25 25 + self._id2tok = tokens 26 26 + self._tok2id = {t: i for i, t in enumerate(tokens)} 27 27 + 28 28 + def __len__(self) -> int: 29 29 + return len(self._id2tok) 30 30 + 31 31 + @property 32 32 + def pad_id(self) -> int: 33 33 + return 0 # always first by construction 34 34 + 35 35 + @property 36 36 + def sos_id(self) -> int: 37 37 + return 1 38 38 + 39 39 + @property 40 40 + def eos_id(self) -> int: 41 41 + return 2 42 42 + 43 43 + @property 44 44 + def unk_id(self) -> int: 45 45 + return 3 46 46 + 47 47 + def encode(self, typst: str) -> list[int]: 48 48 + return [self._tok2id.get(t, self.unk_id) for t in tokenize(typst)] 49 49 + 50 50 + def decode(self, ids: list[int], skip_special: bool = True) -> str: 51 51 + skip = {self.pad_id, self.sos_id, self.eos_id} if skip_special else set() 52 52 + return " ".join(self._id2tok[i] for i in ids if i not in skip) 53 53 + 54 54 + def save(self, path: Path) -> None: 55 55 + path.write_text(json.dumps(self._id2tok, ensure_ascii=False)) 56 56 + 57 57 + @classmethod 58 58 + def load(cls, path: Path) -> "Vocab": 59 59 + return cls(json.loads(path.read_text())) 60 60 + 61 61 + @classmethod 62 62 + def build(cls, typst_strings: list[str], min_freq: int = 2) -> "Vocab": 63 63 + counter: Counter = Counter() 64 64 + for s in typst_strings: 65 65 + counter.update(tokenize(s)) 66 66 + tokens = SPECIAL + [t for t, c in counter.most_common() if c >= min_freq] 67 67 + return cls(tokens)