bobine 0.2.0__py3-none-any.whl

bobine/__init__.py

@@ -0,0 +1,84 @@
+"""bobine — PDF / Office / text → Markdown ingestion engine.
+
+Standalone module extracted from the OKFgraph project. Runs on a single
+``onnxruntime`` wheel with no CUDA-version coupling (RapidAI family + pdf_oxide),
+and needs nothing but the stdlib for plain text / markdown documents.
+
+Public API
+----------
+Conversion (PDF/Office)
+    - ``ConverterConfig`` / ``RoutingMode`` — configuration and routing modes.
+    - ``HybridConverter`` — core conversion pipeline (Qt-independent).
+    - ``OnnxRapidEngine`` — lazy ONNX model manager.
+    - ``html_tables_to_gfm`` — HTML table → GFM pipe-table converter.
+
+Images
+    - ``stage_images_as_okf_assets`` — okf-asset:// staging for extracted images.
+    - ``asset_id`` / ``ASSET_STORE_DIRNAME`` — asset naming convention.
+
+Text-type documents
+    - ``Document`` — normalized document model (id, title, body, tags, …).
+    - ``load_markdown_document`` — frontmatter-aware .md loading.
+    - ``wrap_thoughts`` — raw reasoning text → OKF-compliant markdown.
+    - ``lint_markdown`` / ``lint_markdown_file`` — mordant linting (guarded).
+
+Orchestration
+    - ``convert_to_markdown`` — dispatch PDF / Office / text → markdown string.
+    - ``stage_images`` — collect extracted images into an asset store.
+    - ``ingest_document`` — full pipeline: convert → write .md → stage assets
+      → lint, returns a ``ConvertedDocument``.
+
+Versioning
+    - ``check_rapid_versions`` — runtime version check for RapidAI packages.
+"""
+
+from bobine.assets import (
+    ASSET_STORE_DIRNAME,
+    asset_id,
+    stage_images_as_okf_assets,
+)
+from bobine.config import ConverterConfig, RoutingMode
+from bobine.converter import HybridConverter
+from bobine.documents import Document, load_markdown_document, wrap_thoughts
+from bobine.engine import OnnxRapidEngine
+from bobine.markdown import lint_markdown, lint_markdown_file
+from bobine.pipeline import (
+    OFFICE_EXTS,
+    SUPPORTED_EXTENSIONS,
+    TEXT_EXTS,
+    ConvertedDocument,
+    convert_directory,
+    convert_to_markdown,
+    ingest_document,
+    stage_images,
+)
+from bobine.tables import html_tables_to_gfm
+from bobine.versions import check_rapid_versions
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "ASSET_STORE_DIRNAME",
+    "OFFICE_EXTS",
+    "SUPPORTED_EXTENSIONS",
+    "TEXT_EXTS",
+    "ConvertedDocument",
+    "ConverterConfig",
+    "Document",
+    "HybridConverter",
+    "OnnxRapidEngine",
+    "RoutingMode",
+    "__version__",
+    "asset_id",
+    "check_rapid_versions",
+    "convert_directory",
+    "convert_to_markdown",
+    "html_tables_to_gfm",
+    "ingest_document",
+    "lint_markdown",
+    "lint_markdown_file",
+    "load_markdown_document",
+    "stage_images",
+    "stage_images_as_okf_assets",
+    "wrap_thoughts",
+]

bobine/_vendor/rapid_latex_ocr/__init__.py

@@ -0,0 +1,9 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL (vendored into bobine, MIT licensed)
+from .main import LaTeXOCR
+
+# Legacy alias: older builds shipped the class as `LatexOCR`; keep both so
+# bobine's engine import works regardless of which build was vendored.
+LatexOCR = LaTeXOCR
+
+__all__ = ["LaTeXOCR", "LatexOCR"]

bobine/_vendor/rapid_latex_ocr/main.py

@@ -0,0 +1,214 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import argparse
+import re
+import time
+import traceback
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Tuple, Union
+
+import numpy as np
+import yaml
+from PIL import Image
+
+from .models import EncoderDecoder
+from .utils import DownloadModel, PreProcess, TokenizerCls, get_file_encode
+from .utils_load import InputType, LoadImage, LoadImageError, OrtInferSession
+
+cur_dir = Path(__file__).resolve().parent
+DEFAULT_CONFIG = cur_dir / "config.yaml"
+
+
+@dataclass
+class LaTeXOCRInput:
+    max_width: int = 672
+    max_height: int = 192
+    min_height: int = 32
+    min_width: int = 32
+    bos_token: int = 1
+    max_seq_len: int = 512
+    eos_token: int = 2
+    temperature: float = 0.00001
+
+
+class LaTeXOCR:
+    def __init__(
+        self,
+        config_path: Union[str, Path] = DEFAULT_CONFIG,
+        image_resizer_path: Union[str, Path] = None,
+        encoder_path: Union[str, Path] = None,
+        decoder_path: Union[str, Path] = None,
+        tokenizer_json: Union[str, Path] = None,
+    ):
+        self.image_resizer_path = image_resizer_path
+        self.encoder_path = encoder_path
+        self.decoder_path = decoder_path
+        self.tokenizer_json = tokenizer_json
+
+        self.get_model_path()
+
+        file_encode = get_file_encode(config_path)
+        with open(config_path, "r", encoding=file_encode) as f:
+            args = yaml.load(f, Loader=yaml.FullLoader)
+        input_params = LaTeXOCRInput(**args)
+
+        self.max_dims = [input_params.max_width, input_params.max_height]
+        self.min_dims = [input_params.min_width, input_params.min_height]
+        self.temperature = input_params.temperature
+
+        self.load_img = LoadImage()
+
+        self.pre_pro = PreProcess(max_dims=self.max_dims, min_dims=self.min_dims)
+
+        self.image_resizer = OrtInferSession(self.image_resizer_path)
+
+        self.encoder_decoder = EncoderDecoder(
+            encoder_path=self.encoder_path,
+            decoder_path=self.decoder_path,
+            bos_token=input_params.bos_token,
+            eos_token=input_params.eos_token,
+            max_seq_len=input_params.max_seq_len,
+        )
+        self.tokenizer = TokenizerCls(self.tokenizer_json)
+
+    def get_model_path(
+        self,
+    ) -> Tuple[str]:
+        def try_download(file_name):
+            save_path = default_model_dir / file_name
+            if save_path.exists() or downloader(file_name):
+                return save_path
+            raise FileNotFoundError(f"{file_name} must not be None.")
+
+        downloader = DownloadModel()
+        decoder_name = "decoder.onnx"
+        encoder_name = "encoder.onnx"
+        resizer_name = "image_resizer.onnx"
+        tokenizer_name = "tokenizer.json"
+
+        default_model_dir = cur_dir / "models"
+
+        if self.image_resizer_path is None:
+            self.image_resizer_path = try_download(resizer_name)
+
+        if self.encoder_path is None:
+            self.encoder_path = try_download(encoder_name)
+
+        if self.decoder_path is None:
+            self.decoder_path = try_download(decoder_name)
+
+        if self.tokenizer_json is None:
+            self.tokenizer_json = try_download(tokenizer_name)
+
+    def __call__(self, img: InputType) -> Tuple[str, float]:
+        s = time.perf_counter()
+
+        try:
+            img = self.load_img(img)
+        except LoadImageError as exc:
+            error_info = traceback.format_exc()
+            raise LoadImageError(f"Load the img meets error. Error info is {error_info}") from exc
+
+        try:
+            resizered_img = self.loop_image_resizer(img)
+        except Exception as e:
+            error_info = traceback.format_exc()
+            raise ValueError(f"image resizer meets error. Error info is {error_info}") from e
+
+        try:
+            dec = self.encoder_decoder(resizered_img, temperature=self.temperature)
+        except Exception as e:
+            error_info = traceback.format_exc()
+            raise ValueError(f"EncoderDecoder meets error. Error info is {error_info}") from e
+
+        decode = self.tokenizer.token2str(dec)
+        pred = self.post_process(decode[0])
+
+        elapse = time.perf_counter() - s
+        return pred, elapse
+
+    def loop_image_resizer(self, img: np.ndarray) -> np.ndarray:
+        pillow_img = Image.fromarray(img)
+        pad_img = self.pre_pro.pad(pillow_img)
+        input_image = self.pre_pro.minmax_size(pad_img).convert("RGB")
+        r, w, h = 1, input_image.size[0], input_image.size[1]
+        for _ in range(10):
+            h = int(h * r)
+            final_img, pad_img = self.pre_process(input_image, r, w, h)
+
+            resizer_res = self.image_resizer([final_img.astype(np.float32)])[0]
+
+            argmax_idx = int(np.argmax(resizer_res, axis=-1).item())
+            w = (argmax_idx + 1) * 32
+            if w == pad_img.size[0]:
+                break
+
+            r = w / pad_img.size[0]
+        return final_img
+
+    def pre_process(self, input_image: Image.Image, r, w, h) -> Tuple[np.ndarray, Image.Image]:
+        if r > 1:
+            resize_func = Image.Resampling.BILINEAR
+        else:
+            resize_func = Image.Resampling.LANCZOS
+
+        resize_img = input_image.resize((w, h), resize_func)
+        pad_img = self.pre_pro.pad(self.pre_pro.minmax_size(resize_img))
+        cvt_img = np.array(pad_img.convert("RGB"))
+
+        gray_img = self.pre_pro.to_gray(cvt_img)
+        normal_img = self.pre_pro.normalize(gray_img)
+        final_img = self.pre_pro.transpose_and_four_dim(normal_img)
+        return final_img, pad_img
+
+    @staticmethod
+    def post_process(s: str) -> str:
+        """Remove unnecessary whitespace from LaTeX code.
+
+        Args:
+            s (str): Input string
+
+        Returns:
+            str: Processed image
+        """
+        text_reg = r"(\\(operatorname|mathrm|text|mathbf)\s?\*? {.*?})"
+        letter = "[a-zA-Z]"
+        noletter = r"[\W_^\d]"
+        names = [x[0].replace(" ", "") for x in re.findall(text_reg, s)]
+        s = re.sub(text_reg, lambda match: str(names.pop(0)), s)
+        news = s
+        while True:
+            s = news
+            news = re.sub(r"(?!\\ )(%s)\s+?(%s)" % (noletter, noletter), r"\1\2", s)
+            news = re.sub(r"(?!\\ )(%s)\s+?(%s)" % (noletter, letter), r"\1\2", news)
+            news = re.sub(r"(%s)\s+?(%s)" % (letter, noletter), r"\1\2", news)
+            if news == s:
+                break
+        return s
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-img_resizer", "--image_resizer_path", type=str, default=None)
+    parser.add_argument("-encdoer", "--encoder_path", type=str, default=None)
+    parser.add_argument("-decoder", "--decoder_path", type=str, default=None)
+    parser.add_argument("-tokenizer", "--tokenizer_json", type=str, default=None)
+    parser.add_argument("img_path", type=str, help="Only img path of the formula.")
+    args = parser.parse_args()
+
+    engine = LaTeXOCR(
+        image_resizer_path=args.image_resizer_path,
+        encoder_path=args.encoder_path,
+        decoder_path=args.decoder_path,
+        tokenizer_json=args.tokenizer_json,
+    )
+
+    result, elapse = engine(args.img_path)
+    print(result)
+    print(f"cost: {elapse:.5f}")
+
+
+if __name__ == "__main__":
+    main()

bobine/_vendor/rapid_latex_ocr/models.py

@@ -0,0 +1,154 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+from pathlib import Path
+from typing import Tuple, Union
+
+import numpy as np
+
+from .utils_load import OrtInferSession
+
+
+class EncoderDecoder:
+    def __init__(
+        self,
+        encoder_path: Union[Path, str],
+        decoder_path: Union[Path, str],
+        bos_token: int,
+        eos_token: int,
+        max_seq_len: int,
+    ):
+        self.bos_token = bos_token
+        self.eos_token = eos_token
+        self.max_seq_len = max_seq_len
+
+        self.encoder = OrtInferSession(encoder_path)
+        self.decoder = Decoder(decoder_path)
+
+    def __call__(self, x: np.ndarray, temperature: float = 0.25):
+        ort_input_data = np.array([self.bos_token] * len(x))[:, None]
+        context = self.encoder([x])[0]
+        output = self.decoder(
+            ort_input_data,
+            self.max_seq_len,
+            eos_token=self.eos_token,
+            context=context,
+            temperature=temperature,
+        )
+        return output
+
+
+class Decoder:
+    def __init__(self, decoder_path: Union[Path, str]):
+        self.max_seq_len = 512
+        self.session = OrtInferSession(decoder_path)
+
+    def __call__(
+        self,
+        start_tokens,
+        seq_len=256,
+        eos_token=None,
+        temperature=1.0,
+        filter_thres=0.9,
+        context=None,
+    ):
+        num_dims = len(start_tokens.shape)
+
+        b, t = start_tokens.shape
+
+        out = start_tokens
+        mask = np.full_like(start_tokens, True, dtype=bool)
+
+        for _ in range(seq_len):
+            x = out[:, -self.max_seq_len :]
+            mask = mask[:, -self.max_seq_len :]
+
+            ort_outs = self.session([x.astype(np.int64), mask, context])[0]
+            np_preds = ort_outs
+            np_logits = np_preds[:, -1, :]
+
+            np_filtered_logits = self.npp_top_k(np_logits, thres=filter_thres)
+            np_probs = self.softmax(np_filtered_logits / temperature, axis=-1)
+
+            sample = self.multinomial(np_probs.squeeze(), 1)[None, ...]
+
+            out = np.concatenate([out, sample], axis=-1)
+            mask = np.pad(mask, [(0, 0), (0, 1)], "constant", constant_values=True)
+
+            if eos_token is not None and (np.cumsum(out == eos_token, axis=1)[:, -1] >= 1).all():
+                break
+
+        out = out[:, t:]
+        if num_dims == 1:
+            out = out.squeeze(0)
+        return out
+
+    @staticmethod
+    def softmax(x, axis=None) -> float:
+        def logsumexp(a, axis=None, b=None, keepdims=False):
+            a_max = np.amax(a, axis=axis, keepdims=True)
+
+            if a_max.ndim > 0:
+                a_max[~np.isfinite(a_max)] = 0
+            elif not np.isfinite(a_max):
+                a_max = 0
+
+            tmp = np.exp(a - a_max)
+
+            # suppress warnings about log of zero
+            with np.errstate(divide="ignore"):
+                s = np.sum(tmp, axis=axis, keepdims=keepdims)
+                out = np.log(s)
+
+            if not keepdims:
+                a_max = np.squeeze(a_max, axis=axis)
+            out += a_max
+            return out
+
+        return np.exp(x - logsumexp(x, axis=axis, keepdims=True))
+
+    def npp_top_k(self, logits, thres=0.9):
+        k = int((1 - thres) * logits.shape[-1])
+        val, ind = self.np_top_k(logits, k)
+        probs = np.full_like(logits, float("-inf"))
+        np.put_along_axis(probs, ind, val, axis=1)
+        return probs
+
+    @staticmethod
+    def np_top_k(
+        a: np.ndarray, k: int, axis=-1, largest=True, sorted=True
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        if axis is None:
+            axis_size = a.size
+        else:
+            axis_size = a.shape[axis]
+
+        assert 1 <= k <= axis_size
+
+        a = np.asanyarray(a)
+        if largest:
+            index_array = np.argpartition(a, axis_size - k, axis=axis)
+            topk_indices = np.take(index_array, -np.arange(k) - 1, axis=axis)
+        else:
+            index_array = np.argpartition(a, k - 1, axis=axis)
+            topk_indices = np.take(index_array, np.arange(k), axis=axis)
+
+        topk_values = np.take_along_axis(a, topk_indices, axis=axis)
+        if sorted:
+            sorted_indices_in_topk = np.argsort(topk_values, axis=axis)
+            if largest:
+                sorted_indices_in_topk = np.flip(sorted_indices_in_topk, axis=axis)
+            sorted_topk_values = np.take_along_axis(topk_values, sorted_indices_in_topk, axis=axis)
+            sorted_topk_indices = np.take_along_axis(
+                topk_indices, sorted_indices_in_topk, axis=axis
+            )
+            return sorted_topk_values, sorted_topk_indices
+        return topk_values, topk_indices
+
+    @staticmethod
+    def multinomial(weights, num_samples, replacement=True):
+        weights = np.asarray(weights)
+        weights /= np.sum(weights)  # 确保权重之和为1
+        indices = np.arange(len(weights))
+        samples = np.random.choice(indices, size=num_samples, replace=replacement, p=weights)
+        return samples

bobine/_vendor/rapid_latex_ocr/utils.py

@@ -0,0 +1,185 @@
+# -*- encoding: utf-8 -*-
+# @Author: SWHL
+# @Contact: liekkaskono@163.com
+import io
+from pathlib import Path
+from typing import List, Optional, Union
+
+import chardet
+import cv2
+import numpy as np
+import requests
+import tqdm
+from PIL import Image
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+
+
+class PreProcess:
+    def __init__(self, max_dims: List[int], min_dims: List[int]):
+        self.max_dims, self.min_dims = max_dims, min_dims
+        self.mean = np.array([0.7931, 0.7931, 0.7931]).astype(np.float32)
+        self.std = np.array([0.1738, 0.1738, 0.1738]).astype(np.float32)
+
+    @staticmethod
+    def pad(img: Image.Image, divable: int = 32) -> Image.Image:
+        """Pad an Image to the next full divisible value of `divable`. Also normalizes the image and invert if needed.
+
+        Args:
+            img (PIL.Image): input image
+            divable (int, optional): . Defaults to 32.
+
+        Returns:
+            PIL.Image
+        """
+        threshold = 128
+        data = np.array(img.convert("LA"))
+        if data[..., -1].var() == 0:
+            data = (data[..., 0]).astype(np.uint8)
+        else:
+            data = (255 - data[..., -1]).astype(np.uint8)
+
+        data = (data - data.min()) / (data.max() - data.min()) * 255
+        if data.mean() > threshold:
+            # To invert the text to white
+            gray = 255 * (data < threshold).astype(np.uint8)
+        else:
+            gray = 255 * (data > threshold).astype(np.uint8)
+            data = 255 - data
+
+        coords = cv2.findNonZero(gray)  # Find all non-zero points (text)
+        a, b, w, h = cv2.boundingRect(coords)  # Find minimum spanning bounding box
+        rect = data[b : b + h, a : a + w]
+        im = Image.fromarray(rect).convert("L")
+        dims: List[Union[int, int]] = []
+        for x in [w, h]:
+            div, mod = divmod(x, divable)
+            dims.append(divable * (div + (1 if mod > 0 else 0)))
+
+        padded = Image.new("L", tuple(dims), 255)
+        padded.paste(im, (0, 0, im.size[0], im.size[1]))
+        return padded
+
+    def minmax_size(
+        self,
+        img: Image.Image,
+    ) -> Image.Image:
+        """Resize or pad an image to fit into given dimensions
+
+        Args:
+            img (Image): Image to scale up/down.
+
+        Returns:
+            Image: Image with correct dimensionality
+        """
+        if self.max_dims is not None:
+            ratios = [a / b for a, b in zip(img.size, self.max_dims)]
+            if any([r > 1 for r in ratios]):
+                size = np.array(img.size) // max(ratios)
+                size = np.maximum(size, 1)
+                img = img.resize(tuple(size.astype(int)), Image.BILINEAR)
+
+        if self.min_dims is not None:
+            padded_size: List[Union[int, int]] = [
+                max(img_dim, min_dim) for img_dim, min_dim in zip(img.size, self.min_dims)
+            ]
+
+            new_pad_size = tuple(padded_size)
+            if new_pad_size != img.size:  # assert hypothesis
+                padded_im = Image.new("L", new_pad_size, 255)
+                padded_im.paste(img, img.getbbox())
+                img = padded_im
+        return img
+
+    def normalize(self, img: np.ndarray, max_pixel_value=255.0) -> np.ndarray:
+        mean = self.mean * max_pixel_value
+        std = self.std * max_pixel_value
+        denominator = np.reciprocal(std, dtype=np.float32)
+        img = img.astype(np.float32)
+        img -= mean
+        img *= denominator
+        return img
+
+    @staticmethod
+    def to_gray(img) -> np.ndarray:
+        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+        return cv2.cvtColor(gray, cv2.COLOR_GRAY2RGB)
+
+    @staticmethod
+    def transpose_and_four_dim(img: np.ndarray) -> np.ndarray:
+        return img.transpose(2, 0, 1)[:1][None, ...]
+
+
+class TokenizerCls:
+    def __init__(self, json_file: Union[Path, str]):
+        self.tokenizer = Tokenizer(BPE()).from_file(str(json_file))
+
+    def token2str(self, tokens) -> List[str]:
+        if len(tokens.shape) == 1:
+            tokens = tokens[None, :]
+
+        dec = [self.tokenizer.decode(tok.tolist()) for tok in tokens]
+        return [
+            "".join(detok.split(" "))
+            .replace("Ġ", " ")
+            .replace("[EOS]", "")
+            .replace("[BOS]", "")
+            .replace("[PAD]", "")
+            .strip()
+            for detok in dec
+        ]
+
+
+class DownloadModel:
+    """Modified from https://github.com/lukas-blecher/LaTeX-OCR/blob/1781514fb8c92ea9f94057295fdae0e683f4648e/pix2tex/model/checkpoints/get_latest_checkpoint.py"""
+
+    def __init__(self) -> None:
+        self.url = "https://github.com/RapidAI/RapidLaTeXOCR/releases/download/v0.0.0"
+        self.cur_dir = Path(__file__).resolve().parent
+
+    def __call__(self, file_name: str) -> bool:
+        save_dir = self.cur_dir / "models"
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+        full_path = f"{self.url}/{file_name}"
+        print(f"Download {full_path} to {self.cur_dir}/models")
+
+        try:
+            file = self.download_as_bytes_with_progress(full_path, file_name)
+            save_file_path = save_dir / file_name
+            self.save_file(save_file_path, file)
+        except Exception:
+            return False
+        return True
+
+    @staticmethod
+    def download_as_bytes_with_progress(url: str, name: Optional[str] = None) -> bytes:
+        resp = requests.get(url, stream=True, allow_redirects=True)
+        total = int(resp.headers.get("content-length", 0))
+        bio = io.BytesIO()
+        with tqdm.tqdm(desc=name, total=total, unit="b", unit_scale=True, unit_divisor=1024) as bar:
+            for chunk in resp.iter_content(chunk_size=65536):
+                bar.update(len(chunk))
+                bio.write(chunk)
+        return bio.getvalue()
+
+    @staticmethod
+    def save_file(save_path: Union[str, Path], file: bytes):
+        with open(save_path, "wb") as f:
+            f.write(file)
+
+
+def get_file_encode(file_path: Union[str, Path]) -> str:
+    try:
+        with open(file_path, "rb") as f:
+            raw_data = f.read(100)
+            result = chardet.detect(raw_data)
+            encoding = result["encoding"]
+        return encoding
+    except Exception:
+        return "utf-8"
+
+
+if __name__ == "__main__":
+    downloader = DownloadModel()
+    downloader("decoder.onnx")