Upload folder using huggingface_hub

image_processing_latex_ocr.py

@@ -1,11 +1,31 @@
 import torch
 import numpy as np
-from PIL import Image
+from PIL import Image, ImageOps, ImageEnhance
 from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
 from transformers.utils import logging
 
 logger = logging.get_logger(__name__)
 
+
+def _prepare_for_inference(img: Image.Image) -> Image.Image:
+    """
+    Normalize real-world inputs (screenshots, camera, PDF crops) to the
+    clean white-background style the model was trained on.
+
+    Steps applied in order:
+      1. Convert to grayscale luminance to check background tone
+      2. If dark background (mean < 0.45), invert — handles dark mode / night mode
+      3. Auto-contrast to stretch histogram — fixes low-contrast scans/photos
+      4. Mild sharpening to counter screenshot JPEG blur
+    """
+    arr = np.array(img.convert("L"), dtype=np.float32) / 255.0
+    if arr.mean() < 0.45:
+        img = ImageOps.invert(img.convert("RGB"))
+    img = ImageOps.autocontrast(img, cutoff=1)
+    img = ImageEnhance.Sharpness(img).enhance(1.4)
+    return img.convert("RGB")
+
+
 class LaTeXOCRImageProcessor(BaseImageProcessor):
     model_type = "latex_ocr"
 
@@ -21,24 +41,28 @@ class LaTeXOCRImageProcessor(BaseImageProcessor):
         self.max_image_width = max_image_width
         self.patch_size = patch_size
 
-    def preprocess(self, images, **kwargs) -> BatchFeature:
+    def preprocess(self, images, do_prepare=True, **kwargs) -> BatchFeature:
         if not isinstance(images, list):
             images = [images]
-        
+
         processed_images = []
         for img in images:
             if img.mode != "RGB":
                 img = img.convert("RGB")
 
+            if do_prepare:
+                img = _prepare_for_inference(img)
+
             w, h = img.size
             new_w = int(round(w * self.image_height / max(h, 1)))
             new_w = min(new_w, self.max_image_width)
             new_w = max((new_w // self.patch_size) * self.patch_size, self.patch_size)
-            
+
             if (w, h) != (new_w, self.image_height):
                 img = img.resize((new_w, self.image_height), Image.BILINEAR)
 
             img_array = np.array(img).astype(np.float32) / 255.0
+            img_array = (img_array - 0.5) / 0.5
             img_array = np.transpose(img_array, (2, 0, 1))
             processed_images.append(img_array)
 

modeling_latex_decoder.py

@@ -11,7 +11,7 @@ from transformers.modeling_outputs import CausalLMOutput
 try:
     from .configuration_latex_decoder import LaTeXDecoderConfig
 except ImportError:
-    from configuration_latex_decoder import LaTeXDecoderConfig
+    from latex_ocr.configuration_latex_decoder import LaTeXDecoderConfig
 
 
 class RMSNorm(nn.Module):

modeling_latex_ocr.py

@@ -12,9 +12,9 @@ try:
     from .configuration_latex_ocr import LaTeXOCRConfig
     from .modeling_latex_decoder import LaTeXDecoderForCausalLM
 except ImportError:
-    from configuration_latex_decoder import LaTeXDecoderConfig
-    from configuration_latex_ocr import LaTeXOCRConfig
-    from modeling_latex_decoder import LaTeXDecoderForCausalLM
+    from latex_ocr.configuration_latex_decoder import LaTeXDecoderConfig
+    from latex_ocr.configuration_latex_ocr import LaTeXOCRConfig
+    from latex_ocr.modeling_latex_decoder import LaTeXDecoderForCausalLM
 
 try:
     from flash_attn import flash_attn_func, flash_attn_varlen_func
@@ -319,14 +319,67 @@ class CustomDecoder(nn.Module):
 
     @torch.no_grad()
     def generate(self, inputs_embeds, attention_mask, max_new_tokens, num_beams=1):
-        eos_id = self.eos_token_id
         device = inputs_embeds.device
-        batch = inputs_embeds.shape[0]
-        assert batch == 1, "beam search only supports batch_size=1"
-        vis_emb = inputs_embeds[0]
-        vis_len = vis_emb.shape[0]
+        batch  = inputs_embeds.shape[0]
+
+        if num_beams > 1:
+            # beam search: only supports batch_size=1
+            assert batch == 1, "beam search only supports batch_size=1"
+            return self._beam_search(inputs_embeds, attention_mask, max_new_tokens, num_beams)
+
+        return self._greedy_batch(inputs_embeds, attention_mask, max_new_tokens)
+
+    @torch.no_grad()
+    def _greedy_batch(self, inputs_embeds, attention_mask, max_new_tokens):
+        """Greedy decoding with true batch support."""
+        eos_id  = self.eos_token_id
+        pad_id  = self._pad_id
+        device  = inputs_embeds.device
+        batch   = inputs_embeds.shape[0]
+        d_model = inputs_embeds.shape[-1]
+
+        # generated token ids per sample, and finished flags
+        gen_ids  = [[] for _ in range(batch)]
+        finished = torch.zeros(batch, dtype=torch.bool, device=device)
+
+        cur_embeds = inputs_embeds                          # (B, vis_len, D)
+        cur_mask   = attention_mask                         # (B, vis_len)
+
+        for _ in range(max_new_tokens):
+            logits   = self._forward_embeds(cur_embeds, cur_mask)  # (B, seq, vocab)
+            next_tok = logits[:, -1, :].argmax(dim=-1)             # (B,)
+
+            for i in range(batch):
+                if not finished[i]:
+                    gen_ids[i].append(next_tok[i].item())
+            finished |= (next_tok == eos_id)
+            if finished.all():
+                break
+
+            tok_emb  = self._model.embed_tokens(next_tok.unsqueeze(1))   # (B, 1, D)
+            tok_mask = cur_mask.new_ones(batch, 1)
+            cur_embeds = torch.cat([cur_embeds, tok_emb], dim=1)
+            cur_mask   = torch.cat([cur_mask,   tok_mask], dim=1)
+
+        # pad to same length and return (B, max_len)
+        max_len = max((len(ids) for ids in gen_ids), default=0)
+        if max_len == 0:
+            return torch.zeros(batch, 0, dtype=torch.long, device=device)
+        out = torch.full((batch, max_len), pad_id, dtype=torch.long, device=device)
+        for i, ids in enumerate(gen_ids):
+            if ids:
+                out[i, :len(ids)] = torch.tensor(ids, dtype=torch.long, device=device)
+        return out
+
+    @torch.no_grad()
+    def _beam_search(self, inputs_embeds, attention_mask, max_new_tokens, num_beams):
+        """Original beam search (batch_size=1 only)."""
+        eos_id   = self.eos_token_id
+        device   = inputs_embeds.device
+        vis_emb  = inputs_embeds[0]
+        vis_len  = vis_emb.shape[0]
         vis_mask = attention_mask[0] if attention_mask is not None else None
-        beams = [(0.0, [], False) for _ in range(num_beams)]
+        beams    = [(0.0, [], False) for _ in range(num_beams)]
 
         for _ in range(max_new_tokens):
             all_embeds, all_masks = [], []
@@ -341,7 +394,7 @@ class CustomDecoder(nn.Module):
             max_len = max(e.shape[0] for e in all_embeds)
             d_model = all_embeds[0].shape[-1]
             padded_embeds = vis_emb.new_zeros(num_beams, max_len, d_model)
-            padded_mask = vis_mask.new_zeros(num_beams, max_len) if vis_mask is not None else None
+            padded_mask   = vis_mask.new_zeros(num_beams, max_len) if vis_mask is not None else None
             for idx, emb in enumerate(all_embeds):
                 padded_embeds[idx, :emb.shape[0]] = emb
                 if padded_mask is not None:
@@ -366,7 +419,7 @@ class CustomDecoder(nn.Module):
 
         best_ids = max(beams, key=lambda x: x[0])[1]
         if not best_ids:
-            return torch.zeros(batch, 0, dtype=torch.long, device=device)
+            return torch.zeros(1, 0, dtype=torch.long, device=device)
         return torch.tensor(best_ids, dtype=torch.long, device=device).unsqueeze(0)
 
 
@@ -400,6 +453,12 @@ class LaTeXOCRModel(PreTrainedModel):
         self.decoder = CustomDecoder(config)
         self.post_init()
 
+    def tie_weights(self):
+        if self.config.decoder_weights_tied:
+            self.decoder.tie_weights()
+        else:
+            self.decoder.untie_weights()
+
     def _init_weights(self, module):
         return
 

pipeline_latex_ocr.py

@@ -22,11 +22,11 @@ class LaTeXOCRPipeline:
 
         sys.path.insert(0, str(path))
 
-        from tokenization_latex_ocr import LaTeXTokenizer
-        from image_processing_latex_ocr import LaTeXOCRImageProcessor
-        from processing_latex_ocr import LaTeXOCRProcessor
-        from modeling_latex_ocr import LaTeXOCRModel
-        from configuration_latex_ocr import LaTeXOCRConfig
+        from latex_ocr.tokenization_latex_ocr import LaTeXTokenizer
+        from latex_ocr.image_processing_latex_ocr import LaTeXOCRImageProcessor
+        from latex_ocr.processing_latex_ocr import LaTeXOCRProcessor
+        from latex_ocr.modeling_latex_ocr import LaTeXOCRModel
+        from latex_ocr.configuration_latex_ocr import LaTeXOCRConfig
 
         config = LaTeXOCRConfig.from_pretrained(str(path))
         image_processor = LaTeXOCRImageProcessor.from_pretrained(str(path))
@@ -36,16 +36,19 @@ class LaTeXOCRPipeline:
 
         return cls(model=model, processor=processor, device=device)
 
-    def __call__(self, image, max_new_tokens: int = None, num_beams: int = None) -> str:
-        if isinstance(image, (str, Path)):
-            image = Image.open(image).convert("RGB")
-        elif isinstance(image, Image.Image):
-            image = image.convert("RGB")
-        else:
-            raise TypeError(f"Unsupported image type: {type(image)}")
+    def __call__(self, image, max_new_tokens: int = None, num_beams: int = None):
+        single = not isinstance(image, list)
+        images = [image] if single else image
 
-        inputs = self.processor(images=image, return_tensors="pt")
-        pixel_values = inputs["pixel_values"].to(self.device)
+        loaded = []
+        for img in images:
+            if isinstance(img, (str, Path)):
+                img = Image.open(img).convert("RGB")
+            elif isinstance(img, Image.Image):
+                img = img.convert("RGB")
+            else:
+                raise TypeError(f"Unsupported image type: {type(img)}")
+            loaded.append(img)
 
         kwargs = {}
         if max_new_tokens is not None:
@@ -53,7 +56,15 @@ class LaTeXOCRPipeline:
         if num_beams is not None:
             kwargs["num_beams"] = num_beams
 
+        # image processor handles variable-width images one at a time;
+        # collect pixel_values as a list for NaViT's batched_images path
+        pixel_values = [
+            self.processor(images=img, return_tensors="pt")["pixel_values"].to(self.device)
+            for img in loaded
+        ]
+
         with torch.no_grad():
             generated_ids = self.model.generate(pixel_values, **kwargs)
 
-        return self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        results = self.processor.batch_decode(generated_ids, skip_special_tokens=True)
+        return results[0] if single else results

processing_latex_ocr.py

@@ -1,6 +1,6 @@
 from transformers import ProcessorMixin
-from image_processing_latex_ocr import LaTeXOCRImageProcessor
-from tokenization_latex_ocr import LaTeXTokenizer
+from latex_ocr.image_processing_latex_ocr import LaTeXOCRImageProcessor
+from latex_ocr.tokenization_latex_ocr import LaTeXTokenizer
 
 class LaTeXOCRProcessor(ProcessorMixin):
     attributes = ["image_processor", "tokenizer"]

tokenization_latex_ocr.py

@@ -23,19 +23,14 @@ class LaTeXTokenizer(PreTrainedTokenizer):
         if "model" in data:
             self.token2id: Dict[str, int] = data["model"]["vocab"]
             self.id2token: Dict[int, str] = {int(v): k for k, v in self.token2id.items()}
-            self.merges: List[Tuple[str, str]] = []
+            self.merges = []
             cfg = {}
         else:
             self.token2id = data["token2id"]
             self.id2token = {int(k): v for k, v in data["id2token"].items()}
-            self.merges = [tuple(m) for m in data.get("merges", [])]
+            self.merges = data.get("merges", [])
             cfg = data.get("config", {})
 
-        self.bpe_ranks: Dict[Tuple[str, str], int] = {
-            pair: idx for idx, pair in enumerate(self.merges)
-        }
-        self._bpe_cache: Dict[str, str] = {}
-
         kwargs.setdefault("model_max_length", cfg.get("model_max_length", 256))
         kwargs.setdefault("padding_side", cfg.get("padding_side", "right"))
         kwargs.setdefault("truncation_side", cfg.get("truncation_side", "right"))
@@ -48,42 +43,6 @@ class LaTeXTokenizer(PreTrainedTokenizer):
             **kwargs,
         )
 
-    def _get_pairs(self, word: Tuple[str, ...]):
-        return {(word[i], word[i + 1]) for i in range(len(word) - 1)}
-
-    def _bpe(self, token: str) -> str:
-        if token in self._bpe_cache:
-            return self._bpe_cache[token]
-
-        word = tuple(token)
-        pairs = self._get_pairs(word)
-
-        if not pairs:
-            return token
-
-        while True:
-            bigram = min(pairs, key=lambda p: self.bpe_ranks.get(p, float("inf")))
-            if bigram not in self.bpe_ranks:
-                break
-            first, second = bigram
-            new_word = []
-            i = 0
-            while i < len(word):
-                if i < len(word) - 1 and word[i] == first and word[i + 1] == second:
-                    new_word.append(first + second)
-                    i += 2
-                else:
-                    new_word.append(word[i])
-                    i += 1
-            word = tuple(new_word)
-            pairs = self._get_pairs(word)
-            if not pairs:
-                break
-
-        result = " ".join(word)
-        self._bpe_cache[token] = result
-        return result
-
     @property
     def vocab_size(self) -> int:
         return len(self.token2id)
@@ -99,7 +58,7 @@ class LaTeXTokenizer(PreTrainedTokenizer):
             for length in range(min(20, len(text) - i), 0, -1):
                 substr = text[i:i + length]
                 if substr in self.token2id:
-                    tokens.extend(self._bpe(substr).split())
+                    tokens.append(substr)
                     i += length
                     matched = True
                     break