assets/sdk/src/controlmt/client.py

"""High-level client for ControlMT v2.3."""

from __future__ import annotations
import re
import time
import warnings
from concurrent.futures import ThreadPoolExecutor
from typing import Iterable, Sequence

from controlmt.device import ResolvedConfig, detect_libraries, resolve
from controlmt.batching import auto_batch_size


DEFAULT_MODEL_ID = "anandkaman/controlmt-v2.3"

# Heuristic: 'kn2en' if input is mostly Kannada chars, 'en2kn' otherwise.
_KN_RE = re.compile(r"[ಀ-೿]")


class ControlMT:
    """High-level wrapper around the HuggingFace `model.translate()` API.

    Usage:
        from controlmt import ControlMT
        model = ControlMT.from_hf()                              # auto
        model.translate("ನಾನು ಕನ್ನಡ ಮಾತನಾಡುತ್ತೇನೆ.")
        model.batch_translate(texts, batch_size=8)
        model.batch_translate(texts, auto_batch=True)            # GPU only
    """

    # ──────────────────────────────────────────────────────────────
    #  Construction
    # ──────────────────────────────────────────────────────────────
    def __init__(self, hf_model, tokenizer, config: ResolvedConfig, model_id: str):
        self._model = hf_model
        self._tokenizer = tokenizer
        self.config = config
        self.model_id = model_id

    @classmethod
    def from_hf(
        cls,
        model_id: str = DEFAULT_MODEL_ID,
        *,
        device: str = "auto",                    # "auto" | "gpu" | "cuda" | "cpu"
        dtype: str | None = None,                # "float32" | "bfloat16" | "float16" | None
        quant: str = "none",                     # "none" | "int8" (CPU-only)
        revision: str | None = None,             # HF revision
        verbose: bool = False,
    ) -> "ControlMT":
        """Load ControlMT from HuggingFace + auto-resolve config.

        Defaults: GPU if available, else CPU. fp16 on GPU, bf16 on CPU. No quantization.
        """
        import torch
        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

        cfg = resolve(device=device, dtype=dtype, quant=quant)

        if verbose:
            libs = detect_libraries()
            print(f"[controlmt] loading {model_id}  ({cfg.describe()})")
            print(f"[controlmt] env  torch={libs.get('torch')}  transformers={libs.get('transformers')}  "
                  f"cuda={libs.get('_cuda')}  device={libs.get('_cuda_device')}")

        torch_dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16,
                       "float16": torch.float16}[cfg.dtype_str]

        load_kw = {"trust_remote_code": True}
        if revision: load_kw["revision"] = revision
        if cfg.quant != "int8-dynamic" and cfg.dtype_str != "float32":
            load_kw["dtype"] = torch_dtype   # let HF cast during load

        tokenizer = AutoTokenizer.from_pretrained(model_id, **{k: v for k, v in load_kw.items() if k != "dtype"})
        model = AutoModelForSeq2SeqLM.from_pretrained(model_id, **load_kw)

        # int8 dynamic quantization (CPU-only — checked in resolve())
        if cfg.quant == "int8-dynamic":
            model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)

        model = model.to(torch.device(cfg.device)).eval()
        return cls(hf_model=model, tokenizer=tokenizer, config=cfg, model_id=model_id)

    # ──────────────────────────────────────────────────────────────
    #  Inference
    # ──────────────────────────────────────────────────────────────
    def translate(
        self,
        text: str,
        *,
        direction: str | None = None,    # "kn2en" | "en2kn" | None=auto-detect
        num_beams: int = 2,
        anti_lm_alpha: float = 0.5,
        max_length: int = 200,
    ) -> str:
        """Translate one sentence. Direction auto-detected from input script if not given."""
        if not text or not text.strip():
            return ""
        if direction is None:
            direction = self.detect_direction(text)
        return self._model.translate(
            text.strip(),
            tokenizer=self._tokenizer,
            direction=direction,
            num_beams=num_beams,
            anti_lm_alpha=anti_lm_alpha,
            max_length=max_length,
        )

    def batch_translate(
        self,
        texts: Sequence[str],
        *,
        batch_size: int | None = None,         # None → 1 (safe), or int N
        auto_batch: bool = False,              # GPU only — auto-fit by VRAM
        direction: str | None = None,
        num_beams: int = 2,
        anti_lm_alpha: float = 0.5,
        max_length: int = 200,
    ) -> list[str]:
        """Translate a batch. User must supply batch_size, or auto_batch=True on GPU.

        Policy (matching DEPLOYMENT.md Section 11):
          - batch_size=None and auto_batch=False  → batch_size=1 (one at a time, safe)
          - batch_size=N                          → uses N concurrent translations
          - auto_batch=True                       → GPU: probe free VRAM, pick N; CPU: warn + N=1

        We use a ThreadPoolExecutor — each thread runs one model.translate() call.
        For our model this is ~the same throughput as a true batched call (beam
        search is per-sentence), with much simpler code.
        """
        if not texts:
            return []

        if auto_batch:
            batch_size = auto_batch_size(
                device=self.config.device,
                dtype=self.config.dtype_str,
                quant=self.config.quant,
            )
        if batch_size is None:
            batch_size = 1
        batch_size = max(1, int(batch_size))

        # If batch_size=1, no need for threadpool overhead
        if batch_size == 1:
            return [self.translate(t, direction=direction, num_beams=num_beams,
                                   anti_lm_alpha=anti_lm_alpha, max_length=max_length)
                    for t in texts]

        def _one(t: str) -> str:
            return self.translate(t, direction=direction, num_beams=num_beams,
                                  anti_lm_alpha=anti_lm_alpha, max_length=max_length)

        with ThreadPoolExecutor(max_workers=batch_size) as ex:
            return list(ex.map(_one, texts))

    # ──────────────────────────────────────────────────────────────
    #  Helpers
    # ──────────────────────────────────────────────────────────────
    @staticmethod
    def detect_direction(text: str) -> str:
        """Heuristic: > 30% Kannada characters → kn2en, else en2kn."""
        if not text: return "en2kn"
        kn_chars = sum(1 for c in text if _KN_RE.match(c))
        total = sum(1 for c in text if not c.isspace() and c.isprintable())
        return "kn2en" if total and kn_chars / total > 0.3 else "en2kn"

    def warmup(self) -> float:
        """JIT/compile kernels on a throwaway translation. Returns elapsed seconds."""
        t0 = time.time()
        self.translate("hello.", direction="en2kn", num_beams=1, max_length=20)
        return time.time() - t0

    def benchmark(self, num_beams: int = 2) -> dict:
        """Run the 6-pair DEPLOYMENT.md verification suite on this loaded model.
        Returns the same shape as scripts/verify_deployment.py's JSON output."""
        TEST_PAIRS = [
            ("kn2en", "ನಾನು ಕನ್ನಡ ಮಾತನಾಡುತ್ತೇನೆ."),
            ("kn2en", "ಬೆಂಗಳೂರಿನಲ್ಲಿ ಮೆಟ್ರೋ ಬಹಳ ಅನುಕೂಲಕರವಾಗಿದೆ."),
            ("kn2en", "ಆಪಲ್ ಹೊಸ ಐಫೋನ್ 17 ಬಿಡುಗಡೆ ಮಾಡಿದೆ."),
            ("en2kn", "I speak Kannada."),
            ("en2kn", "The new metro line opens next month."),
            ("en2kn", "Please transfer money to my UPI ID."),
        ]
        self.warmup()
        rows = []
        for direction, src in TEST_PAIRS:
            t0 = time.time()
            out = self.translate(src, direction=direction, num_beams=num_beams)
            rows.append({"direction": direction, "src": src, "out": out,
                         "latency_s": round(time.time() - t0, 3)})
        lats = sorted([r["latency_s"] for r in rows])
        return {
            "config": self.config.describe(),
            "num_beams": num_beams,
            "rows": rows,
            "median_latency_s": lats[len(lats)//2],
        }

    def __repr__(self) -> str:
        return f"<ControlMT model_id={self.model_id!r} {self.config.describe()}>"