thx01 1.0.0__py3-none-any.whl

thx01/__init__.py

@@ -0,0 +1,60 @@
+"""THX-01: fast, non-autoregressive decision engine with calibrated probabilities."""
+
+from .agent import Agent, RLAgent, load
+from .decide import decide as _decide
+from .common import (
+    QTYPES,
+    QTYPE_NAMES,
+    confidence_from_probs,
+    ece_score,
+    proper_reward,
+    render_options,
+    td_lambda_targets,
+)
+from .email import clean_email_body, email_state
+from .lang import analyse as detect_language
+from .lang import detect_script, is_english
+from .presets import (
+    email_questions,
+    guard_questions,
+    moderation_questions,
+    router_questions,
+    triage_questions,
+)
+from .shortlist import embed_fn_from_agent, predict_shortlist, shortlist_choice
+
+__version__ = "1.0.0"
+__all__ = [
+    "Agent",
+    "RLAgent",
+    "load",
+    "shortlist_choice",
+    "predict_shortlist",
+    "embed_fn_from_agent",
+    "detect_language",
+    "detect_script",
+    "is_english",
+    "clean_email_body",
+    "email_questions",
+    "email_state",
+    "guard_questions",
+    "moderation_questions",
+    "router_questions",
+    "triage_questions",
+    "proper_reward",
+    "td_lambda_targets",
+    "ece_score",
+    "confidence_from_probs",
+    "render_options",
+    "QTYPES",
+    "QTYPE_NAMES",
+    "__version__",
+]
+
+
+def _agent_decide(self, state, questions):
+    """Answer any THX-01 question set (choice, noul, score, number, excerpt, cite) exactly like the HTTP API."""
+    return _decide(self, state, questions)
+
+
+Agent.decide = _agent_decide

thx01/agent.py

@@ -0,0 +1,519 @@
+"""High-level inference runtime for thx01 System 1 decision models."""
+import json
+import os
+import warnings
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+
+from .common import (
+    QTYPES,
+    TEMP_MAX,
+    TEMP_MIN,
+    amp_dtype,
+    build_model,
+    build_sequence,
+    clamp_temperature,
+    collate_items,
+    confidence_from_probs,
+    render_options,
+    temp_bucket,
+)
+
+
+def _fix_tokenizer_config(path: str):
+    """Ensure tokenizer_config.json can be loaded across all transformers versions."""
+    cfg_file = os.path.join(path, "tokenizer", "tokenizer_config.json")
+    if not os.path.exists(cfg_file):
+        return
+    try:
+        with open(cfg_file) as f:
+            tcfg = json.load(f)
+        changed = False
+        if tcfg.get("tokenizer_class") in (None, "TokenizersBackend"):
+            tcfg["tokenizer_class"] = "PreTrainedTokenizerFast"
+            tcfg.pop("backend", None)
+            tcfg.pop("is_local", None)
+            changed = True
+        # Checkpoints built on the mmBERT/Gemma tokenizer store extra_special_tokens as a list;
+        # transformers expects a mapping and raises "'list' object has no attribute 'keys'",
+        # which makes AutoTokenizer -- and so the whole model -- fail to load.
+        extra = tcfg.get("extra_special_tokens")
+        if isinstance(extra, list):
+            tcfg["extra_special_tokens"] = {"extra_%d" % i: t for i, t in enumerate(extra)}
+            changed = True
+        if changed:
+            with open(cfg_file, "w") as f:
+                json.dump(tcfg, f, indent=2)
+    except Exception:
+        pass
+
+
+def _verify_compatibility(model: torch.nn.Module, cfg: Dict, weights: Dict[str, torch.Tensor], model_id: str):
+    """Verify that the loaded checkpoint weights and config strictly match the expected architecture."""
+    # 1. Verify required configuration attributes
+    required_cfg = ["encoder", "head_layers"]
+    missing_cfg = [k for k in required_cfg if k not in cfg]
+    if missing_cfg:
+        raise ValueError(
+            f"Incompatible model config for {model_id!r}: missing configuration keys {missing_cfg}. "
+            f"Ensure this is a valid RL Agent decision model."
+        )
+
+    # 2. Check for required component prefixes
+    required_prefixes = ("encoder.", "type_emb.", "scorer.", "act_head.")
+    for prefix in required_prefixes:
+        if not any(k.startswith(prefix) for k in weights.keys()):
+            raise ValueError(
+                f"Incompatible model weights for {model_id!r}: checkpoint is missing '{prefix}' parameters. "
+                f"Expected an RL Agent decision model with encoder and decision heads."
+            )
+
+    # 3. Check for parameter shape mismatches
+    model_sd = model.state_dict()
+    shape_mismatches = []
+    missing_keys = []
+
+    for name, param in model.named_parameters():
+        if name not in weights:
+            missing_keys.append(name)
+        elif tuple(weights[name].shape) != tuple(param.shape):
+            shape_mismatches.append(f"  - {name}: expected {tuple(param.shape)}, found {tuple(weights[name].shape)}")
+
+    if shape_mismatches:
+        err_details = "\n".join(shape_mismatches[:5])
+        if len(shape_mismatches) > 5:
+            err_details += f"\n  ... and {len(shape_mismatches) - 5} more mismatched layers."
+        raise ValueError(
+            f"Model architecture mismatch for {model_id!r}:\n{err_details}\n"
+            f"The checkpoint weights do not match the configured model architecture."
+        )
+
+    if missing_keys:
+        raise ValueError(
+            f"Model weights incomplete for {model_id!r}: missing {len(missing_keys)} parameter tensors "
+            f"(e.g. {missing_keys[:3]})."
+        )
+
+
+class Agent:
+    """System 1 decision model runtime: fast, non-autoregressive, calibrated decisions."""
+
+    def __init__(
+        self,
+        model_id_or_path: str = "doofz/THX-01",
+        device: Optional[str] = None,
+        token: Optional[str] = None,
+        subfolder: Optional[str] = None,
+        half: bool = True,
+    ):
+        """Load a THX-01 checkpoint.
+
+        `subfolder` selects one checkpoint from a repo that bundles several, e.g.
+        `Agent("doofz/THX-01", subfolder="multilingual")`. Only that subfolder is
+        downloaded, so bundling does not cost every user the whole family.
+        """
+        from safetensors.torch import load_file
+        from transformers import AutoTokenizer
+
+        model_dir = model_id_or_path
+        if not os.path.exists(model_dir):
+            if model_id_or_path.startswith(("/", "./", "../")) or os.path.isabs(model_id_or_path):
+                raise FileNotFoundError(
+                    f"Local model path not found: {model_id_or_path!r}. "
+                    f"Check that the directory exists and that training saved the model successfully."
+                )
+            from huggingface_hub import snapshot_download
+
+            # Restrict root checkpoints too: the default repo also contains sibling
+            # checkpoints, which an unfiltered snapshot would unnecessarily download.
+            prefix = f"{subfolder}/" if subfolder else ""
+            kw = {
+                "token": token or os.environ.get("HF_TOKEN"),
+                "allow_patterns": [prefix + name for name in (
+                    "config.json", "rl_agent_config.json", "model.safetensors", "tokenizer/*", "encoder/*",
+                )],
+            }
+            model_dir = snapshot_download(model_id_or_path, **kw)
+
+        if subfolder:
+            model_dir = os.path.join(model_dir, subfolder)
+            if not os.path.isdir(model_dir):
+                raise FileNotFoundError(
+                    f"Subfolder {subfolder!r} not found in {model_id_or_path!r}."
+                )
+
+        _fix_tokenizer_config(model_dir)
+
+        cfg_path = os.path.join(model_dir, "rl_agent_config.json")
+        if not os.path.exists(cfg_path):
+            raise FileNotFoundError(
+                f"Incompatible model: {model_id_or_path!r} does not contain 'rl_agent_config.json'. "
+                f"That file ships with the weights of a THX-01 checkpoint, so load one of those "
+                f"(e.g. 'doofz/THX-01') or a directory your own training run wrote."
+            )
+
+        with open(cfg_path) as f:
+            self.cfg = json.load(f)
+
+        weights_path = os.path.join(model_dir, "model.safetensors")
+        if not os.path.exists(weights_path):
+            raise FileNotFoundError(
+                f"Incompatible model: 'model.safetensors' not found in {model_id_or_path!r}."
+            )
+
+        # 1. Device resolution with automatic fallback
+        if device is not None:
+            target_device = torch.device(device)
+            if target_device.type == "cuda" and not torch.cuda.is_available():
+                print("Warning: CUDA requested but not available. Falling back to CPU.")
+                self.device = torch.device("cpu")
+            elif target_device.type == "mps" and not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
+                print("Warning: MPS requested but not available. Falling back to CPU.")
+                self.device = torch.device("cpu")
+            else:
+                self.device = target_device
+        else:
+            if torch.cuda.is_available():
+                self.device = torch.device("cuda")
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                self.device = torch.device("mps")
+            else:
+                self.device = torch.device("cpu")
+
+        tok_dir = os.path.join(model_dir, "tokenizer")
+        self.tok = AutoTokenizer.from_pretrained(tok_dir if os.path.exists(tok_dir) else self.cfg.get("encoder"))
+
+        enc_dir = os.path.join(model_dir, "encoder")
+        self.model = build_model(self.cfg, encoder_dir=enc_dir if os.path.exists(enc_dir) else None)
+
+        # Load weights and verify architectural compatibility
+        weights = load_file(weights_path)
+        _verify_compatibility(self.model, self.cfg, weights, model_id_or_path)
+
+        self.model.load_state_dict(weights, strict=True)
+
+        # ModernBERT's reference_compile defaults to "auto" and will torch.compile the encoder.
+        # That is a loss for the batch sizes THX-01 runs (a handful of questions per call) and can
+        # hang on some platforms, so keep the eager path.
+        try:
+            self.model.encoder.config.reference_compile = False
+        except Exception:
+            pass
+
+        # Keep what the checkpoint shipped for inspection, but only ever apply clamped values:
+        # some buckets are fitted to sharpen rather than soften (see clamp_temperature).
+        self.temperature_raw = self.cfg.get("temperature", [1.0, 1.0, 1.0])
+        self.temperature_by_options_raw = self.cfg.get("temperature_by_options", {})
+        self.temperature = [clamp_temperature(t) for t in self.temperature_raw]
+        self.temperature_by_options = {k: clamp_temperature(v)
+                                       for k, v in self.temperature_by_options_raw.items()}
+        rejected = ["%s=%.4g" % (k, float(v)) for k, v in self.temperature_by_options_raw.items()
+                    if clamp_temperature(v) != float(v)]
+        rejected += ["temperature[%d]=%.4g" % (i, float(t)) for i, t in enumerate(self.temperature_raw)
+                     if clamp_temperature(t) != float(t)]
+        if rejected:
+            warnings.warn(
+                "thx01: this checkpoint ships temperatures outside [%g, %g] which would distort "
+                "confidence; clamping %s. Treat confidence from the affected buckets as uncalibrated."
+                % (TEMP_MIN, TEMP_MAX, ", ".join(rejected)),
+                RuntimeWarning, stacklevel=2)
+        self.dtype = amp_dtype(self.cfg.get("amp_dtype", "fp16"))
+
+        if self.device.type == "cuda" and torch.cuda.get_device_capability(self.device)[0] < 8:
+            self.dtype = torch.float16
+        elif self.device.type in ("cpu", "mps"):
+            self.dtype = torch.float32
+
+        # 2. Place on device with graceful fallback to CPU on memory error
+        fell_back_from = fell_back_why = None
+        try:
+            self.model.to(self.device).eval()
+            if self.device.type == "cuda" and half:
+                # Store the weights in the autocast dtype instead of fp32: halves VRAM
+                # (421M -> ~0.85 GB, 322M -> ~0.65 GB) with no change to what autocast computes.
+                self.model.to(self.dtype)
+        except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
+            if self.device.type != "cpu":
+                # Record what actually went wrong: the reason matters more than the symptom,
+                # and it is the only place the underlying exception is ever surfaced.
+                fell_back_from, fell_back_why = self.device, e
+                self.device = torch.device("cpu")
+                self.dtype = torch.float32
+                self.model.to(self.device).eval()
+            else:
+                raise e
+
+        if fell_back_from is not None:
+            print(
+                "\n[thx01] Warning: could not place the model on %s, so it is running on CPU.\n"
+                "  Reason: %s\n"
+                "  Inference will be roughly 10-15x slower (~200-500 ms rather than ~35 ms).\n"
+                "  If this is a newer NVIDIA GPU (Blackwell / RTX 50-series), your PyTorch build\n"
+                "  may not support its CUDA architecture:\n"
+                "    pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu128\n"
+                "  See https://pytorch.org/get-started/locally/\n"
+                % (fell_back_from, fell_back_why), flush=True)
+
+    @staticmethod
+    def _to_internal(qdef: Dict) -> Dict:
+        t = qdef["type"]
+        crit = qdef.get("criteria")
+        if t == "choice" and isinstance(crit, list):
+            crit = {c: None for c in crit}
+        ins = qdef["instructions"]
+        if not isinstance(ins, str):
+            ins = json.dumps(ins)
+        return {"t": t, "ins": ins, "crit": crit}
+
+    def _head_budget(self, q: Dict) -> int:
+        """Head (instruction + options) token budget for one question.
+
+        The checkpoint default is kept for ordinary questions. When the options need more room
+        (e.g. 77 intents), the budget grows to fit them -- up to half the context -- instead of
+        truncating every label to 3-4 tokens, which is what made large label sets collapse.
+        """
+        base = self.cfg.get("head_max_len", 192)
+        if not self.cfg.get("auto_head_budget", True):
+            return base
+        opts = render_options(q)
+        need = 24 + sum(min(49, len(self.tok(" " + o, add_special_tokens=False)["input_ids"]) + 1) for o in opts)
+        return int(min(max(base, need), self.cfg.get("max_len", 512) // 2))
+
+    def _decode(self, questions: Dict[str, Dict[str, Any]], items, logits: np.ndarray, act: np.ndarray,
+                offset: int = 0) -> Dict[str, Any]:
+        """Turn raw marker logits for one state's questions into calibrated typed answers."""
+        ids = list(questions.keys())
+        answers = {}
+        for r, qid in enumerate(ids):
+            q = self._to_internal(questions[qid])
+            k = len(items[r]["markers"])
+            qt = QTYPES[q["t"]]
+            t_scale = self.temperature_by_options.get(temp_bucket(qt, k), self.temperature[qt])
+            z = logits[offset + r, :k] / t_scale
+            p = np.exp(z - z.max())
+            p = p / p.sum()
+
+            conf_score = round(confidence_from_probs(p, k), 4)
+            ext = {"act_probability": round(float(act[offset + r, 0]), 4)}
+
+            if q["t"] == "choice":
+                keys = list(q["crit"].keys())
+                answers[qid] = {
+                    "type": "choice",
+                    "choice": keys[int(p.argmax())],
+                    "probabilities": {kk: round(float(v), 4) for kk, v in zip(keys, p)},
+                    "confidence": conf_score,
+                    "action": ext,
+                }
+            elif q["t"] == "score":
+                exp_score = float((np.arange(k) * p).sum())
+                answers[qid] = {
+                    "type": "score",
+                    "score": round(exp_score, 4),
+                    "legend": {str(i): c for i, c in enumerate(q["crit"])},
+                    "probabilities": {str(i): round(float(v), 4) for i, v in enumerate(p)},
+                    "confidence": conf_score,
+                    "action": ext,
+                }
+            else:
+                answers[qid] = {
+                    "type": "noul",
+                    "noul": round(float(p[1]), 4),
+                    "confidence": round(max(float(p[1]), 1.0 - float(p[1])), 4),
+                    "action": ext,
+                }
+
+        return answers
+
+    def _tournament(self, states, qdefs, chunk: int = 20, keep: int = 3):
+        """Coarse-to-fine answer for choice questions with many options (two batched passes).
+
+        Round 1 scores each state against every chunk of ~`chunk` options and keeps the top `keep`
+        of each; round 2 chooses among those finalists. Returns one answer dict per state, with
+        probabilities over the finalists (every other option gets 0).
+        """
+        rng = __import__("random").Random(0)
+        s1, q1, spans = [], [], []
+        for st, q in zip(states, qdefs):
+            keys = list(q["criteria"])
+            rng.shuffle(keys)
+            chunks = [keys[i:i + chunk] for i in range(0, len(keys), chunk)]
+            spans.append(len(chunks))
+            for c in chunks:
+                s1.append(st)
+                q1.append({"q": {**q, "criteria": {k: q["criteria"][k] for k in c}}})
+        o1 = self.predict_batch(s1, q1, _tournament=False) if s1 else []
+        s2, q2, off = [], [], 0
+        for st, q, n in zip(states, qdefs, spans):
+            fin = []
+            for o in o1[off:off + n]:
+                pr = o["answers"]["q"]["probabilities"]
+                fin += sorted(pr, key=pr.get, reverse=True)[:keep]
+            off += n
+            s2.append(st)
+            q2.append({"q": {**q, "criteria": {k: q["criteria"][k] for k in fin}}})
+        o2 = self.predict_batch(s2, q2, _tournament=False)
+        out = []
+        for q, o in zip(qdefs, o2):
+            a = o["answers"]["q"]
+            a["probabilities"] = {k: a["probabilities"].get(k, 0.0) for k in q["criteria"]}
+            a["mode"] = "tournament"
+            out.append(a)
+        return out
+
+    @torch.no_grad()
+    def predict_batch(self, states, questions, batch_size: int = 64, _tournament: bool = True):
+        """Answer many states at once (throughput mode for evals and bulk jobs).
+
+        `questions` is either one schema shared by every state or a list with one schema per
+        state. Returns one `system_one`-style result per state, in order.
+        """
+        from .common import collate_items as _collate
+        qlist = questions if isinstance(questions, list) else [questions] * len(states)
+        limit = self.cfg.get("tournament_above", 24)
+        if _tournament and limit:
+            big = [(i, qid) for i, qs in enumerate(qlist) for qid, q in qs.items()
+                   if q.get("type") == "choice" and isinstance(q.get("criteria"), (dict, list)) and len(q["criteria"]) > limit]
+            if big:
+                norm = lambda q: {**q, "criteria": q["criteria"] if isinstance(q["criteria"], dict) else {c: None for c in q["criteria"]}}
+                answers = self._tournament([states[i] for i, _ in big], [norm(qlist[i][qid]) for i, qid in big])
+                rest = [{k: v for k, v in qs.items() if (i, k) not in set(big)} for i, qs in enumerate(qlist)]
+                keep = [i for i, r in enumerate(rest) if r]
+                sub = self.predict_batch([states[i] for i in keep], [rest[i] for i in keep], batch_size, _tournament=False) if keep else []
+                res = [{"model": "THX-01", "answers": {}, "usage": {"input_tokens": 0, "output_tokens": 0}} for _ in states]
+                for i, r in zip(keep, sub):
+                    res[i] = r
+                for (i, qid), a in zip(big, answers):
+                    res[i]["answers"][qid] = a
+                for i, qs in enumerate(qlist):                   # restore the caller's question order
+                    res[i]["answers"] = {k: res[i]["answers"][k] for k in qs}
+                return res
+        max_len = self.cfg.get("max_len", 512)
+        head_max_len = self.cfg.get("head_max_len", 192)
+        per_state = []
+        flat = []
+        for st, qs in zip(states, qlist):
+            its = []
+            for qid in qs:
+                q = self._to_internal(qs[qid])
+                head_max_len = self._head_budget(q)
+                seq, markers = build_sequence(self.tok, st, q, max_len, head_max_len)
+                if len(markers) != len(render_options(q)):
+                    raise ValueError("question %r options exceed head_max_len=%d" % (qid, head_max_len))
+                its.append({"ids": seq, "markers": markers, "qtype": QTYPES[q["t"]]})
+            per_state.append(its)
+            flat.extend(its)
+        # sort by length so each chunk pads as little as possible
+        order = sorted(range(len(flat)), key=lambda i: len(flat[i]["ids"]))
+        kmax = max(len(it["markers"]) for it in flat)
+        all_logits = np.full((len(flat), kmax), -1e4, dtype=np.float32)
+        all_act = np.zeros((len(flat), 2), dtype=np.float32)
+        use_amp = self.device.type == "cuda"
+        for c in range(0, len(order), batch_size):
+            idx = order[c:c + batch_size]
+            bt = _collate([[flat[i] for i in idx]], self.tok.pad_token_id)
+            with torch.autocast(device_type=self.device.type, dtype=self.dtype, enabled=use_amp):
+                lg, ac = self.model(bt["input_ids"].to(self.device), bt["attention_mask"].to(self.device),
+                                    bt["marker_pos"].to(self.device), bt["marker_mask"].to(self.device),
+                                    bt["qtype"].to(self.device))
+            lg = lg.float().cpu().numpy()
+            ac = torch.softmax(ac.float(), -1).cpu().numpy()
+            for j, i in enumerate(idx):
+                all_logits[i, :lg.shape[1]] = lg[j]
+                all_act[i, :ac.shape[1]] = ac[j]
+        out, off = [], 0
+        for qs, its in zip(qlist, per_state):
+            out.append({"model": "THX-01",
+                        "answers": self._decode(qs, its, all_logits, all_act, offset=off),
+                        "usage": {"input_tokens": sum(len(it["ids"]) for it in its), "output_tokens": 0}})
+            off += len(its)
+        return out
+
+    @torch.no_grad()
+    def system_one(self, state: Union[str, dict, list], questions: Dict[str, Dict[str, Any]]) -> Dict[str, Any]:
+        """Evaluate typed questions across state in a single, parallel forward pass.
+
+        Args:
+            state: Text string, JSON dict, or conversation turn list.
+            questions: Dictionary mapping question_id -> question definition.
+                - choice: {"type": "choice", "instructions": "...", "criteria": {"optA": "...", ...}}
+                - score:  {"type": "score",  "instructions": "...", "criteria": ["lvl0", "lvl1", ...]}
+                - noul:   {"type": "noul",   "instructions": "..."}
+
+        Returns:
+            Dictionary with answers, probabilities, calibrated confidence, and token usage.
+        """
+        limit = self.cfg.get("tournament_above", 24)
+        if limit and any(q.get("type") == "choice" and isinstance(q.get("criteria"), (dict, list))
+                         and len(q["criteria"]) > limit for q in questions.values()):
+            return self.predict_batch([state], questions)[0]
+        ids = list(questions.keys())
+        items = []
+        max_len = self.cfg.get("max_len", 512)
+        head_max_len = self.cfg.get("head_max_len", 192)
+
+        for qid in ids:
+            q = self._to_internal(questions[qid])
+            head_max_len = self._head_budget(q)
+            seq, markers = build_sequence(self.tok, state, q, max_len, head_max_len)
+            if len(markers) != len(render_options(q)):
+                raise ValueError("question %r options exceed head_max_len=%d" % (qid, head_max_len))
+            items.append({"ids": seq, "markers": markers, "qtype": QTYPES[q["t"]]})
+
+        b = collate_items([items], self.tok.pad_token_id)
+        use_amp = self.device.type == "cuda"
+
+        try:
+            with torch.autocast(device_type=self.device.type, dtype=self.dtype, enabled=use_amp):
+                logits, act = self.model(
+                    b["input_ids"].to(self.device),
+                    b["attention_mask"].to(self.device),
+                    b["marker_pos"].to(self.device),
+                    b["marker_mask"].to(self.device),
+                    b["qtype"].to(self.device),
+                )
+        except (RuntimeError, torch.cuda.OutOfMemoryError) as e:
+            if self.device.type != "cpu" and ("memory" in str(e).lower() or "cuda" in str(e).lower()):
+                print("Warning: GPU memory exceeded during inference. Falling back to CPU...")
+                self.device = torch.device("cpu")
+                self.dtype = torch.float32
+                self.model.to(self.device)
+                logits, act = self.model(
+                    b["input_ids"].to(self.device),
+                    b["attention_mask"].to(self.device),
+                    b["marker_pos"].to(self.device),
+                    b["marker_mask"].to(self.device),
+                    b["qtype"].to(self.device),
+                )
+            else:
+                raise e
+
+        logits = logits.float().cpu().numpy()
+        act = torch.softmax(act.float(), -1).cpu().numpy()
+        n_tokens = int(b["attention_mask"].sum())
+        answers = self._decode(questions, items, logits, act)
+
+        return {
+            "model": "THX-01",
+            "answers": answers,
+            "usage": {"input_tokens": n_tokens, "output_tokens": 0},
+        }
+
+    predict = system_one
+
+
+RLAgent = Agent
+
+
+def load(model_id_or_path: str = "doofz/THX-01", device: Optional[str] = None,
+         token: Optional[str] = None, subfolder: Optional[str] = None, half: bool = True) -> Agent:
+    """Load a THX-01 agent.
+
+    `subfolder` picks one checkpoint out of a repo that bundles several:
+
+        thx01.load("doofz/THX-01")                           # English (repo root)
+        thx01.load("doofz/THX-01", subfolder="multilingual")
+    """
+    return Agent(model_id_or_path, device=device, token=token, subfolder=subfolder, half=half)