"""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""