| """THX-02 runtime: one decision call with option-order averaging, the joint-read rule and self-gated thinking. |
| |
| from thx02 import THX02 |
| m = THX02.load("HAL-X-AI/THX-02", thinker_url="http://127.0.0.1:8121") # or a local folder; thinker optional |
| m.decide("I was charged twice, please refund", {"refund": {"type": "noul", "instructions": "Refund request?"}}) |
| |
| The decision logic is a generator (`plan`) that yields batches of sub-requests and thought jobs; `THX02.decide` drives |
| it synchronously and `thx02.serve` drives many of them concurrently through one batched forward pass. |
| """ |
| from __future__ import annotations |
|
|
| import concurrent.futures |
| import json |
| import sys |
| import time |
| import urllib.request |
| from pathlib import Path |
|
|
| from .engine import (ALWAYS, AUTO, THINK_SYSTEM, distribution, facet_joint, gate_request, geo_mean, orders, permuted, |
| think_body, with_thought) |
| from .protocol import answer, options |
|
|
|
|
| def resolve(path) -> Path: |
| """A local folder as is; otherwise a Hugging Face repo id, downloaded without the thinker and mobile files.""" |
| local = Path(str(path)).expanduser() |
| if local.exists(): |
| return local |
| from huggingface_hub import snapshot_download |
| return Path(snapshot_download(str(path), ignore_patterns=["thinker/*", "mobile/*", "assets/*"])) |
|
|
|
|
| def load_decider(path, device="cuda", fast=False): |
| if device == "cpu" and "transformers.models.qwen3_5.modeling_qwen3_5" not in sys.modules: |
| sys.modules.setdefault("fla", None) |
| from .model import THXDecisionModel |
| return THXDecisionModel.from_pretrained(resolve(path), device=device, fast=fast) |
|
|
|
|
| def plan(req: dict, *, think_available: bool, default_tta: int = 1, max_questions: int = 64): |
| """Generator: yields ("decide", [requests]) or ("think", [single-question requests]); receives the outputs |
| (decision responses or thought strings) and returns the response dict.""" |
| state, questions = req.get("state", ""), req.get("questions") |
| if not isinstance(questions, dict) or not questions or len(questions) > max_questions: |
| raise ValueError(f"questions must be an object with 1 to {max_questions} entries") |
| images = [i for i in (req.get("images") or []) if isinstance(i, str)][:4] |
| base = {"state": state, **({"images": images} if images else {})} |
| think = req.get("think", False) |
| cfg = ALWAYS if think is True else AUTO if think == "auto" else None |
| if not think_available: |
| cfg = None |
| k = max(1, min(int(req.get("tta", default_tta)), 8)) |
| keys = list(questions) |
|
|
| if not images and facet_joint(questions): |
| out, = yield ("decide", [{**base, "questions": questions}]) |
| return {"model": "THX-02", "answers": out["answers"]} |
|
|
| subs, spans = [], [] |
| for key in keys: |
| q = questions[key] |
| crit = q.get("criteria") if q.get("type") == "choice" else None |
| vs = orders(len(crit), k) if k > 1 and isinstance(crit, dict) and len(crit) > 1 else [None] |
| s = len(subs) |
| subs += [{**base, "questions": {key: q if o is None else permuted(q, o)}} for o in vs] |
| spans.append((s, len(subs))) |
| n_main = len(subs) |
| if cfg and cfg["gate"] > 0: |
| subs += [gate_request(state, questions[key]) for key in keys] |
| outs = yield ("decide", subs) |
|
|
| answers, p0 = {}, {} |
| for key, (s, e) in zip(keys, spans): |
| ks = options(questions[key])[0] |
| ps = [distribution(outs[j]["answers"][key], ks) for j in range(s, e)] |
| p0[key] = ps[0] if len(ps) == 1 else geo_mean(ps) |
| answers[key] = outs[s]["answers"][key] if len(ps) == 1 else answer(questions[key], p0[key]) |
|
|
| if cfg: |
| gates = [o["answers"]["g"]["noul"] for o in outs[n_main:]] if cfg["gate"] > 0 else [1.0] * len(keys) |
| unsure = [key for key, g in zip(keys, gates) if max(p0[key]) < cfg["tau"] and g >= cfg["gate"]] |
| if unsure: |
| thoughts = yield ("think", [{**base, "questions": {key: questions[key]}} for key in unsure], cfg["cap"]) |
| pairs = [(key, th) for key, th in zip(unsure, thoughts) if th] |
| if pairs: |
| re = yield ("decide", [{**base, "questions": {key: with_thought(questions[key], th)}} for key, th in pairs]) |
| for (key, th), o in zip(pairs, re): |
| if isinstance(o, Exception): |
| continue |
| p1 = distribution(o["answers"][key], options(questions[key])[0]) |
| answers[key] = answer(questions[key], [cfg["alpha"] * x + (1 - cfg["alpha"]) * y for x, y in zip(p0[key], p1)]) |
| if req.get("return_thoughts"): |
| answers[key]["thought"] = th |
| return {"model": "THX-02", "answers": answers} |
|
|
|
|
| def thinker_payload(enc, request: dict, cap: int) -> dict: |
| content = [{"type": "image_url", "image_url": {"url": u}} for u in request.get("images", [])] |
| content.append({"type": "text", "text": think_body(enc, request)}) |
| return {"model": "thx02-thinker", "messages": [{"role": "system", "content": THINK_SYSTEM}, {"role": "user", "content": content}], |
| "max_tokens": cap, "temperature": 0, "stop": ["</think>"], "chat_template_kwargs": {"enable_thinking": True}} |
|
|
|
|
| def thought_text(response: dict) -> str: |
| m = response["choices"][0]["message"] |
| return ((m.get("reasoning_content") or m.get("content") or "").replace("<think>", "")).strip() |
|
|
|
|
| class THX02: |
| """Local THX-02: decisions in-process; thoughts from a vLLM thinker server when thinker_url is set.""" |
|
|
| def __init__(self, decider, thinker_url: str | None = None, default_tta: int = 1): |
| self.d, self.thinker_url, self.default_tta = decider, thinker_url, default_tta |
|
|
| @classmethod |
| def load(cls, path="HAL-X-AI/THX-02", device="cuda", fast=False, thinker_url=None, default_tta=1): |
| """path: a local THX-02 folder, or a Hugging Face repo id (downloaded once, about 9 GB).""" |
| return cls(load_decider(path, device=device, fast=fast), thinker_url, default_tta) |
|
|
| def _thoughts(self, requests, cap): |
| def one(r): |
| try: |
| body = json.dumps(thinker_payload(self.d.encoder, r, cap)).encode() |
| q = urllib.request.Request(self.thinker_url.rstrip("/") + "/v1/chat/completions", body, {"Content-Type": "application/json"}) |
| return thought_text(json.load(urllib.request.urlopen(q, timeout=30))) |
| except Exception: |
| return "" |
| with concurrent.futures.ThreadPoolExecutor(16) as ex: |
| return list(ex.map(one, requests)) |
|
|
| def decide(self, state, questions, images=None, think=False, tta=None, return_thoughts=False) -> dict: |
| """think: False | "auto" (only when unsure and reasoning helps) | True (nearly always). tta: option orders (1 or 4).""" |
| t0 = time.time() |
| req = {"state": state, "questions": questions, "images": images or [], "think": think, |
| "tta": self.default_tta if tta is None else tta, "return_thoughts": return_thoughts} |
| g = plan(req, think_available=bool(self.thinker_url), default_tta=self.default_tta) |
| msg = next(g) |
| while True: |
| try: |
| msg = g.send(self.d.decide(msg[1]) if msg[0] == "decide" else self._thoughts(msg[1], msg[2])) |
| except StopIteration as stop: |
| out = stop.value |
| out["latency_ms"] = round((time.time() - t0) * 1000, 1) |
| return out |
|
|
