#!/usr/bin/env python3
"""Xattot formula moduli (2026-10-31): formula kesmasi → LaTeX, PP-FormulaNet_plus-L (Apache-2.0).

Rasmiy ONNX reposi bo'sh, shuning uchun torch va transformers kerak: faqat `--formula` bilan yuklanadi, oddiy quvur torch'siz qoladi.
Sinov (939 kesma, `sinov/jf_1031`): 0.62 s/kesma GPU, 11.7 s CPU (lokal ~7.4 s), ya'ni ommaviy o'qish GPU'da. 48 kesmada 60% to'g'ri, 19% kichik xato, 21% xato.
"""
import re

import torch
from transformers import AutoModelForImageTextToText, AutoProcessor

MODEL = 'PaddlePaddle/PP-FormulaNet_plus-L_safetensors'
REV = '6cc004a8e91ea3f3819b05317f40b290cbe99c2b'                           # pinned (Apache-2.0)
YAKKA = re.compile(r'(?<![\\A-Za-z])[A-Za-z](?: [A-Za-z](?![A-Za-z])){2,}')      # `S c a l e d`: matematik rejimda bo'shliq ahamiyatsiz


class Formula:
    def __init__(self, qurilma=None, batch=None, max_token=600):
        self.dev = qurilma or ('cuda' if torch.cuda.is_available() else 'cpu'); self.max_token = max_token
        self.batch = batch or (8 if self.dev == 'cuda' else 1)   # CPU'da 8 ta matritsali kesma: 8.2 GB, OOM (10-31, p08)
        self.p = AutoProcessor.from_pretrained(MODEL, revision=REV); self.m = AutoModelForImageTextToText.from_pretrained(MODEL, revision=REV).to(self.dev).eval()
        self.tur = next(self.m.parameters()).dtype

    @torch.inference_mode()
    def oqi(self, rasmlar):
        # rasmlar: PIL RGB; 600 token chegarasi: axlat kesmada generatsiya cho'zilmasin (939 dan 17 tasi yetgan, ular uzun tizimlar)
        out = []
        for i in range(0, len(rasmlar), self.batch):
            x = self.p(images=rasmlar[i:i + self.batch], return_tensors='pt').to(self.dev); x['pixel_values'] = x['pixel_values'].to(self.tur)
            out += self.p.post_process(self.m.generate(**x, max_new_tokens=self.max_token))
        return [YAKKA.sub(lambda m: m.group().replace(' ', ''), t) for t in out]