dsh-router-laya 2.1.0

package/service/laya/router.py

@@ -0,0 +1,447 @@
+"""Route a request to the Laya checkpoint best suited to it.
+
+Three checkpoints, measured on a shared benchmark (17,416 questions, one T4, identical questions
+per model -- see the repository's benchmark notebook):
+
+  english          convaiinnovations/laya                421M  ModernBERT-large, 512 tokens
+  multilingual     convaiinnovations/laya-multilingual   322M  mmBERT-base, 1024 tokens, 100+ langs
+  typed-decisions  convaiinnovations/laya-typed-decisions 421M  ModernBERT-large, 1024 tokens,
+                                                                fine-tuned on the typed-decisions
+                                                                workflows
+
+Why routing is worth it -- accuracy by language family:
+
+                      english   multilingual
+  MASSIVE intent  en    0.783       0.657        <- English checkpoint wins
+  MASSIVE intent  non-en 0.306      0.451
+  XNLI            en    0.860       0.843
+  XNLI            non-en 0.521      0.731        <- +21 points for multilingual
+  English suites        0.684       0.619
+
+The English checkpoint does not gently degrade off English, it collapses: on 20-option MASSIVE
+intent it scores 0.100 on Hindi and 0.103 on Korean, against 0.050 for random guessing -- and it
+reports high confidence while doing so (ECE 0.855 on Hindi). Script detection is therefore the
+primary routing signal.
+
+`typed-decisions` is never selected automatically unless you opt in with
+`auto_task_detection=True` or pass `task="typed_decisions"`: it is fine-tuned on four specific
+synthetic workflows and should not be a silent default.
+"""
+import gc
+import os
+import threading
+from typing import Any, Dict, List, Optional, Union
+
+from .lang import analyse
+
+# The hub repo bundles all three checkpoints; only the requested subfolder is downloaded.
+BUNDLE_REPO = "convaiinnovations/laya"
+DEFAULT_MODELS = {
+    "english": (BUNDLE_REPO, None),
+    "multilingual": (BUNDLE_REPO, "multilingual"),
+    "typed-decisions": (BUNDLE_REPO, "typed-decisions"),
+}
+
+# The same checkpoints also live in their own repos, for anyone who prefers them.
+STANDALONE_MODELS = {
+    "english": "convaiinnovations/laya",
+    "multilingual": "convaiinnovations/laya-multilingual",
+    "typed-decisions": "convaiinnovations/laya-typed-decisions",
+}
+
+
+def _repo_str(spec):
+    """Human-readable id for a model spec: 'repo' or 'repo/subfolder'."""
+    repo, sub = _split(spec)
+    return "%s/%s" % (repo, sub) if sub else repo
+
+
+def _split(spec):
+    """Normalise a model spec to (repo_or_path, subfolder)."""
+    if isinstance(spec, (tuple, list)):
+        repo, sub = (list(spec) + [None])[:2]
+        return repo, sub
+    return spec, None
+
+# Aliases people are likely to type.
+_ALIASES = {
+    "en": "english", "laya": "english", "default": "english",
+    "multi": "multilingual", "ml": "multilingual", "laya-multilingual": "multilingual",
+    "typed": "typed-decisions", "typed_decisions": "typed-decisions",
+    "laya-typed-decisions": "typed-decisions", "decisions": "typed-decisions",
+}
+
+# Question-id signatures of the four typed-decisions workflows, used only when
+# auto_task_detection is enabled.
+_TYPED_DECISION_WORKFLOWS = {
+    "agent_trace_observability": {"action", "needs_review", "outcome", "risk", "urgency"},
+    "customer_service": {"action", "category", "churn_risk", "needs_human", "urgency"},
+    "invoice_processing": {"discrepancy_severity", "disposition", "duplicate", "matches_order", "urgency"},
+    "security_incidents": {"credential_compromise", "disposition", "severity", "true_positive", "urgency"},
+}
+
+
+class RouteDecision(dict):
+    """The routing outcome: which model, why, and what was detected.
+
+    Behaves as a dict so it serialises straight into an API response.
+    """
+
+    @property
+    def model(self) -> str:
+        return self["model"]
+
+    @property
+    def reason(self) -> str:
+        return self["reason"]
+
+    def __repr__(self):
+        return "RouteDecision(model=%r, reason=%r)" % (self["model"], self["reason"])
+
+
+def normalise_name(name: str) -> str:
+    key = str(name).strip().lower()
+    key = _ALIASES.get(key, key)
+    if key not in DEFAULT_MODELS:
+        raise ValueError("unknown model %r; choose one of %s (or an alias: %s)"
+                         % (name, sorted(DEFAULT_MODELS), sorted(_ALIASES)))
+    return key
+
+
+def match_typed_decisions_workflow(questions: Dict[str, Any]) -> Optional[str]:
+    """Name of the typed-decisions workflow whose question ids these are, else None.
+
+    Requires an exact id-set match, so an unrelated schema that happens to contain 'urgency'
+    is never captured.
+    """
+    ids = set(questions or {})
+    for wf, sig in _TYPED_DECISION_WORKFLOWS.items():
+        if ids == sig:
+            return wf
+    return None
+
+
+# Subtags that mean "the English checkpoint can read this". Routing needs one bit -- is this
+# English Latin text, or something the English checkpoint cannot read -- not a language id, so
+# every other code resolves to the multilingual checkpoint.
+_ENGLISH_SUBTAGS = ("en", "eng", "english")
+
+
+def _english_from_code(value: Any) -> Optional[bool]:
+    """True/False for a language code, or None when the code identifies nothing.
+
+    Accepts the forms a caller is likely to have to hand: `"en"`, `"EN"`, `"en-US"`, the
+    POSIX `"en_US"` (which `$LANG` holds), and `"en_US.UTF-8"`. `None` here means "no usable
+    hint", which is what lets a language-identification model abstain.
+    """
+    if value is None:
+        return None
+    code = str(value).strip().lower()
+    if not code:
+        return None
+    code = code.split(".", 1)[0]                       # en_US.UTF-8 -> en_US
+    primary = code.replace("_", "-").split("-", 1)[0]  # en_US -> en
+    if not primary:
+        return None
+    return primary in _ENGLISH_SUBTAGS
+
+
+class Router:
+    """Lazily loads Laya checkpoints and sends each request to the right one.
+
+        from laya import Router
+
+        r = Router()
+        r.predict({"message": "Mein Konto wurde zweimal belastet"}, questions)   # -> multilingual
+        r.predict({"message": "I was charged twice"}, questions)                 # -> english
+        r.predict(state, questions, model="typed-decisions")                     # explicit
+
+    Models are downloaded and built on first use. `max_loaded` caps how many stay resident
+    (least-recently-used is evicted), because all three together are ~1.16B parameters.
+
+    The default is 2, because automatic routing only ever chooses between `english` and
+    `multilingual`: a cap of one rebuilds the checkpoint it just evicted on every script switch,
+    which is seconds per request on exactly the traffic the Router exists for. Traffic that only
+    ever sees one language never builds the second checkpoint, so the default costs it nothing.
+    Lower it to 1 for a memory-constrained host, and raise it to 3 (or preload) when
+    `auto_task_detection`, an explicit `model=` or an explicit `task=` can reach
+    `typed-decisions` as well.
+
+    For a server or a demo, preload instead: a cold load costs seconds, while detection costs
+    microseconds, so even the default still pays a load the first time a language appears.
+
+        r = Router(preload=True)                    # all three resident, routing is free
+        r = Router(preload=True, device="cuda")
+        r.preload(["english", "multilingual"])      # or just the two you serve
+    """
+
+    def __init__(
+        self,
+        models: Optional[Dict[str, str]] = None,
+        device: Optional[str] = None,
+        token: Optional[str] = None,
+        max_loaded: int = 2,
+        default: str = "english",
+        auto_task_detection: bool = False,
+        standalone_repos: bool = False,
+        preload: bool = False,
+        lang_guess: Optional[Any] = None,
+    ):
+        self.models = dict(STANDALONE_MODELS if standalone_repos else DEFAULT_MODELS)
+        if models:
+            self.models.update({normalise_name(k): v for k, v in models.items()})
+        self.device = device
+        self.token = token or os.environ.get("HF_TOKEN")
+        self.max_loaded = max(1, int(max_loaded))
+        self.default = normalise_name(default)
+        self.auto_task_detection = bool(auto_task_detection)
+        # An opt-in language hint installed for every request: a code, or a callable taking the
+        # state and returning one (or None to abstain). Checked before the built-in detection,
+        # never before an explicit `model`, `task` or `lang`. The default path is unchanged, so
+        # the heuristic stays dependency-free; this is the seam for a real LID model.
+        self.lang_guess = lang_guess
+        self._agents: Dict[str, Any] = {}
+        self._order: List[str] = []          # least-recently-used first
+        # Re-entrant lock guarding model lifecycle (load/unload/attach/preload) and the
+        # LRU bookkeeping. RLock so the public methods can call the private `_touch`/`_evict`
+        # helpers without deadlocking. Inference (`Agent.system_one`) is deliberately left
+        # outside the lock so concurrent predictions share a checkpoint without serialising.
+        self._lock = threading.RLock()
+        if preload:
+            self.preload()
+
+    # ------------------------------------------------------------------ loading
+    def load(self, name: str):
+        """Return the Agent for `name`, downloading and building it on first use.
+
+        Concurrent callers share a single Agent instead of building duplicates.
+        """
+        key = normalise_name(name)
+        with self._lock:
+            if key in self._agents:
+                self._touch(key)
+                return self._agents[key]
+            from .agent import Agent
+            repo, sub = _split(self.models[key])
+            agent = Agent(repo, device=self.device, token=self.token, subfolder=sub)
+            self._agents[key] = agent
+            self._order.append(key)
+            self._evict()
+            return agent
+
+    def _touch(self, key: str):
+        with self._lock:
+            if key in self._order:
+                self._order.remove(key)
+            self._order.append(key)
+
+    def _evict(self):
+        with self._lock:
+            evicted = False
+            while len(self._order) > self.max_loaded:
+                victim = self._order.pop(0)
+                agent = self._agents.pop(victim, None)
+                if agent is not None:
+                    evicted = True
+                    del agent
+            if len(self._order) < len(self._agents):     # keep the two views consistent
+                for k in list(self._agents):
+                    if k not in self._order:
+                        agent = self._agents.pop(k, None)
+                        if agent is not None:
+                            evicted = True
+                            del agent
+            if evicted:
+                gc.collect()
+                try:
+                    import torch
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                except Exception:
+                    pass
+
+    def attach(self, name: str, agent: Any):
+        """Register an already-built Agent under `name` instead of loading a second copy.
+
+        Useful when the process has a checkpoint loaded for other reasons: a demo that already
+        built `convaiinnovations/laya` can hand it to the router rather than pay for -- and hold
+        in memory -- a duplicate 421M parameters.
+        """
+        key = normalise_name(name)
+        with self._lock:
+            self._agents[key] = agent
+            self._touch(key)
+            self.max_loaded = max(self.max_loaded, len(self._agents))
+        return agent
+
+    def preload(self, names: Optional[List[str]] = None):
+        """Download and build checkpoints up front so no request ever pays a model load.
+
+        A cold load costs seconds; language detection costs microseconds. With every
+        checkpoint resident, routing is effectively free -- which is what you want in a
+        server or a demo. `max_loaded` is raised to fit both the requested checkpoints and
+        all already-resident agents, so incremental preloading does not evict either.
+        """
+        names = [normalise_name(n) for n in (names or list(self.models))]
+        with self._lock:
+            self.max_loaded = max(self.max_loaded, len(set(names) | set(self._agents)))
+            for n in names:
+                if n not in self._agents:      # an attached agent is already built
+                    self.load(n)
+        return self
+
+    def unload(self, name: Optional[str] = None):
+        """Free one model, or all of them."""
+        with self._lock:
+            if name is None:
+                self._agents.clear()
+                self._order.clear()
+            else:
+                key = normalise_name(name)
+                agent = self._agents.pop(key, None)
+                if key in self._order:
+                    self._order.remove(key)
+                del agent
+            gc.collect()
+            try:
+                import torch
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
+            except Exception:
+                pass
+
+    @property
+    def loaded(self) -> List[str]:
+        with self._lock:
+            return list(self._order)
+
+    def _resolve_hint(self, hint: Any, state: Union[str, dict, list, None]) -> Optional[bool]:
+        """True/False for a hint about whether the English checkpoint can read `state`.
+
+        `hint` is either a language code or a callable taking the state. Anything the hint
+        cannot answer returns None, which makes `route` fall through to detection rather than
+        picking a checkpoint on no evidence.
+        """
+        if hint is None:
+            return None
+        if callable(hint):
+            hint = hint(state)
+        return _english_from_code(hint)
+
+    # ------------------------------------------------------------------ routing
+    def route(
+        self,
+        state: Union[str, dict, list, None],
+        questions: Optional[Dict[str, Any]] = None,
+        model: Optional[str] = None,
+        task: Optional[str] = None,
+        lang: Optional[str] = None,
+        lang_guess: Optional[Any] = None,
+    ) -> RouteDecision:
+        """Decide which checkpoint to use, without loading or running anything.
+
+        Precedence: explicit `model` > explicit `task` > detected workflow (opt-in) >
+        explicit `lang` > `lang_guess` > detected script/language > default.
+
+        `lang_guess` is an opt-in hint -- a language code or a callable taking the state --
+        checked after an explicit `lang` and before the built-in detection. It only answers
+        "can the English checkpoint read this?", so any non-English code routes to the
+        multilingual checkpoint. A hint that resolves to nothing falls through to detection,
+        which lets a language-identification model abstain. Pass one here, or set
+        `Router(lang_guess=...)` to apply it to every request.
+        """
+        if model is not None:
+            key = normalise_name(model)
+            return RouteDecision(model=key, repo=_repo_str(self.models[key]), reason="explicit model=%r" % model,
+                                 detection=None, workflow=None)
+
+        if task is not None:
+            key = normalise_name("typed-decisions" if str(task).lower().replace("-", "_") == "typed_decisions" else task)
+            return RouteDecision(model=key, repo=_repo_str(self.models[key]), reason="explicit task=%r" % task,
+                                 detection=None, workflow=None)
+
+        workflow = match_typed_decisions_workflow(questions or {})
+        if workflow and self.auto_task_detection:
+            return RouteDecision(model="typed-decisions", repo=_repo_str(self.models["typed-decisions"]),
+                                 reason="question ids match the %r typed-decisions workflow" % workflow,
+                                 detection=None, workflow=workflow)
+
+        if lang is not None:
+            key = "english" if _english_from_code(lang) else "multilingual"
+            return RouteDecision(model=key, repo=_repo_str(self.models[key]), reason="explicit lang=%r" % lang,
+                                 detection=None, workflow=workflow)
+
+        # Caller-supplied hint, per-call first then the one installed on the Router. Only a hint
+        # that actually answers the question routes here; anything else falls through.
+        for source, hint in (("lang_guess", lang_guess), ("Router(lang_guess=...)", self.lang_guess)):
+            resolved = self._resolve_hint(hint, state)
+            if resolved is not None:
+                key = "english" if resolved else "multilingual"
+                return RouteDecision(
+                    model=key, repo=_repo_str(self.models[key]),
+                    reason="%s: the caller identified this as %s text" % (
+                        source, "English" if resolved else "non-English"),
+                    detection=None, workflow=workflow)
+
+        det = analyse(state)
+        if det["script"] == "unknown":
+            key = self.default
+            reason = "no letters detected in state; using default (%s)" % key
+        elif det["script"] != "latin":
+            key = "multilingual"
+            reason = "non-Latin script (%s, %.0f%% of letters); the English checkpoint cannot read it" % (
+                det["script"], 100 * float(det["non_latin_fraction"]))
+        elif not det["is_english"]:
+            key = "multilingual"
+            if det["language"]:
+                reason = "Latin script but language looks like %r, not English" % det["language"]
+            else:
+                # Unidentified Latin-script language: routed on the non-English letters alone,
+                # because no stopword list here covers it.
+                reason = ("Latin script, language not identified but %.0f%% non-English letters; "
+                          "not safe for the English checkpoint" % (100 * float(det["diacritic_rate"])))
+        elif det["language_undecided"]:
+            # Nothing identifies the language: too short, or only content words ("Quero cancelar",
+            # "Esqueci minha senha"). That is no evidence of English either, so it takes the same
+            # `default` as a state with no letters. A deployment that serves mostly non-English
+            # traffic sets `Router(default="multilingual")`; the stock default keeps it English.
+            key = self.default
+            reason = ("Latin script, language not identified and no non-English letters; "
+                      "using default (%s)" % key)
+        else:
+            key = "english"
+            reason = "English Latin text"
+        return RouteDecision(model=key, repo=_repo_str(self.models[key]), reason=reason,
+                             detection=det, workflow=workflow)
+
+    # ------------------------------------------------------------------ running
+    def predict(
+        self,
+        state: Union[str, dict, list],
+        questions: Dict[str, Any],
+        model: Optional[str] = None,
+        task: Optional[str] = None,
+        lang: Optional[str] = None,
+        lang_guess: Optional[Any] = None,
+    ) -> Dict[str, Any]:
+        """Route, then answer every question in one forward pass on the chosen checkpoint.
+
+        The result is the usual `system_one` payload plus a `routing` key recording the decision.
+        """
+        decision = self.route(state, questions, model=model, task=task, lang=lang, lang_guess=lang_guess)
+        agent = self.load(decision["model"])
+        result = agent.system_one(state, questions)
+        result["routing"] = dict(decision)
+        return result
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.unload()
+        return False
+
+    system_one = predict
+
+    def __repr__(self):
+        return "Router(loaded=%s, max_loaded=%d, default=%r)" % (self.loaded, self.max_loaded, self.default)