@tuned-tensor/cli 0.19.0 → 0.21.0

package/training/adapter/src/evaluate.py

@@ -10,12 +10,12 @@ from tempfile import TemporaryDirectory
 from typing import Any
 
 from model_contract import (
-    CERTIFIED_BASE_MODELS,
     CERTIFIED_BASE_MODEL,
     assert_certified_base_model,
     assert_certified_base_model_revision,
     assert_certified_model_config,
     chat_template_kwargs,
+    uncertified_load_kwargs,
 )
 
 MAX_ARCHIVE_MEMBERS = 20_000
@@ -145,10 +145,12 @@ def resolve_device(device: str) -> str:
     return device
 
 
-def _assert_certified_model_source(base_model: str) -> None:
-    if not Path(base_model).exists() and base_model not in CERTIFIED_BASE_MODELS:
+def _assert_certified_model_source(base_model: str, model_source: str) -> None:
+    # A hub source must be the requested model; certified or not, it is
+    # checked by assert_certified_base_model.
+    if not Path(model_source).exists() and model_source != base_model:
         raise ValueError(
-            f"The bundled evaluator does not certify {base_model!r}"
+            f"The bundled evaluator does not certify {model_source!r}"
         )
 
 
@@ -160,9 +162,14 @@ def load_text_model(
     base_model = str(payload["base_model"])
     model_source = str(payload.get("model_source", base_model))
     assert_certified_base_model(base_model, "Evaluation base model")
-    _assert_certified_model_source(model_source)
+    _assert_certified_model_source(base_model, model_source)
     revision = payload.get("base_model_revision")
-    assert_certified_base_model_revision(base_model, revision, "Evaluation base model revision")
+    assert_certified_base_model_revision(
+        base_model,
+        revision,
+        "Evaluation base model revision",
+        local_source=Path(model_source).is_dir(),
+    )
     device = resolve_device(str(payload.get("device", "cuda")))
     token = os.getenv("HF_TOKEN")
     source_kwargs: dict[str, Any] = {
@@ -184,7 +191,7 @@ def load_text_model(
     if device == "cuda":
         dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
 
-    model_kwargs: dict[str, Any] = dict(source_kwargs)
+    model_kwargs: dict[str, Any] = {**source_kwargs, **uncertified_load_kwargs(base_model)}
     if dtype is not None:
         model_kwargs["dtype"] = dtype
     if device == "cuda":

package/training/adapter/src/model_contract.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from typing import Any
 
 
@@ -15,6 +16,10 @@ NEMOTRON_BASE_MODEL_REVISION = "ce38b6ab8b252b4b8ee7165b4605e93191cafd73"
 # Immutable Hugging Face revision reviewed and certified for Muse Glimmer.
 MUSE_GLIMMER_BASE_MODEL_REVISION = "a4e59da52a7bc87ae7251dd5545c0dd437c44b68"
 CERTIFIED_BASE_MODELS = (CERTIFIED_BASE_MODEL, NEMOTRON_BASE_MODEL, MUSE_GLIMMER_BASE_MODEL)
+# Any other Hugging Face repo may be trained as an uncertified causal LM when
+# the run pins an immutable commit. Node enforces the same contract.
+HUGGING_FACE_REPO_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}")
+COMMIT_SHA = re.compile(r"[0-9a-fA-F]{40}")
 CERTIFIED_QWEN_TEXT_CONFIG = {
     "model_type": "qwen3_5_text",
     "hidden_size": 2048,
@@ -58,6 +63,9 @@ def assert_certified_model_config(
     label: str = "base-model config",
     expected_model_id: str | None = None,
 ) -> None:
+    if expected_model_id and not is_certified_base_model(expected_model_id):
+        # No reviewed architecture exists for an uncertified model.
+        return
     model_type = _field(value, "model_type")
     architectures = _field(value, "architectures")
 
@@ -126,17 +134,30 @@ def assert_certified_model_config(
         )
 
 
+def is_certified_base_model(model_id: str) -> bool:
+    return model_id in CERTIFIED_BASE_MODELS
+
+
 def assert_certified_base_model(model_id: str, label: str = "base model") -> None:
-    if model_id not in CERTIFIED_BASE_MODELS:
-        supported = ", ".join(CERTIFIED_BASE_MODELS)
-        raise ValueError(f"{label} must be one of {supported}; got {model_id!r}")
+    """Accept a certified model or an uncertified Hugging Face repo id."""
+    if is_certified_base_model(model_id) or HUGGING_FACE_REPO_ID.fullmatch(model_id):
+        return
+    supported = ", ".join(CERTIFIED_BASE_MODELS)
+    raise ValueError(f"{label} must be one of {supported} or a Hugging Face repo id; got {model_id!r}")
 
 
-def assert_certified_base_model_revision(base_model: str, revision: str | None, label: str = "base model revision") -> None:
+def assert_certified_base_model_revision(
+    base_model: str,
+    revision: str | None,
+    label: str = "base model revision",
+    *,
+    local_source: bool = False,
+) -> None:
     """Require the certified immutable revision for pinned training models.
 
     Qwen, Nemotron, and Muse Glimmer loads are bound here; local snapshot
-    contents are additionally verified by the Node runtime.
+    contents are additionally verified by the Node runtime. An uncertified
+    model must pin a commit SHA unless it loads from a local directory.
     """
     expected = {
         CERTIFIED_BASE_MODEL: QWEN_BASE_MODEL_REVISION,
@@ -144,6 +165,13 @@ def assert_certified_base_model_revision(base_model: str, revision: str | None,
         MUSE_GLIMMER_BASE_MODEL: MUSE_GLIMMER_BASE_MODEL_REVISION,
     }.get(base_model)
     if expected is None:
+        if revision is None and local_source:
+            return
+        if revision is None or not COMMIT_SHA.fullmatch(str(revision)):
+            raise ValueError(
+                f"{label} must be a 40-character commit SHA for uncertified base model {base_model}; "
+                f"got {revision!r}"
+            )
         return
     if revision is None:
         raise ValueError(f"{label} is required for {base_model}")
@@ -167,3 +195,8 @@ def chat_template_kwargs(model_id: str) -> dict[str, Any]:
     if model_id == NEMOTRON_BASE_MODEL:
         return {"enable_thinking": False}
     return {}
+
+
+def uncertified_load_kwargs(model_id: str) -> dict[str, Any]:
+    """Refuse pickle weights for models TT has not reviewed."""
+    return {} if is_certified_base_model(model_id) else {"use_safetensors": True}

package/training/adapter/src/prefetch.py

@@ -10,6 +10,7 @@ from model_contract import (
     assert_certified_base_model,
     assert_certified_base_model_revision,
     assert_certified_model_config,
+    is_certified_base_model,
 )
 
 ALLOW_PATTERNS = [
@@ -99,6 +100,10 @@ def verify_snapshot(snapshot: Path, expected_model_id: str | None = None) -> tup
     except (OSError, json.JSONDecodeError) as exc:
         raise ValueError(f"Cached snapshot has an invalid config.json: {snapshot}") from exc
 
+    if expected_model_id and not is_certified_base_model(expected_model_id):
+        if not any(path.stat().st_size > 0 for path in snapshot.glob("*.safetensors")):
+            raise ValueError("Uncertified base models require non-empty safetensors weights; pickle-only checkpoints are unsupported")
+
     weight_files = sorted(snapshot.glob("*.safetensors")) + sorted(snapshot.glob("pytorch_model*.bin"))
     index_files = sorted(snapshot.glob("*.safetensors.index.json")) + sorted(snapshot.glob("pytorch_model*.bin.index.json"))
     for index in index_files:
@@ -175,7 +180,7 @@ def main() -> None:
         revision=payload.get("revision"),
         token=os.getenv("HF_TOKEN"),
         allow_patterns=ALLOW_PATTERNS,
-        ignore_patterns=IGNORE_PATTERNS,
+        ignore_patterns=IGNORE_PATTERNS + ([] if is_certified_base_model(base_model) else ["pytorch_model*.bin"]),
         local_files_only=bool(payload.get("local_files_only", False)),
     )
     snapshot = Path(snapshot_path)

package/training/adapter/src/serve.py

@@ -11,7 +11,11 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 
 from evaluate import configure_hugging_face_cache, resolve_adapter_path
-from model_contract import assert_certified_base_model_revision, assert_certified_model_config
+from model_contract import (
+    assert_certified_base_model_revision,
+    assert_certified_model_config,
+    is_certified_base_model,
+)
 
 class BearerAuthMiddleware:
     """Keep TT's auth boundary on ALL routes, including upstream admin APIs.
@@ -38,6 +42,27 @@ class BearerAuthMiddleware:
         await self.app(scope, receive, send)
 
 
+def native_chat_template(model_source: str) -> str:
+    source = Path(model_source)
+    template_file = source / "chat_template.jinja"
+    template = (template_file.read_text(encoding="utf-8") if template_file.is_file()
+                else json.loads((source / "tokenizer_config.json").read_text())["chat_template"])
+    if not isinstance(template, str):
+        raise ValueError("Serving requires an unambiguous string chat template.")
+    return template
+
+
+def pin_native_content_format(args: list[str], model_source: str) -> None:
+    if "--chat-template" not in args:
+        return
+    # The wrapper handles system content parts, which confuses upstream's AST
+    # detection for string-only native templates. Detect the ORIGINAL template
+    # using the detector from our pinned vLLM runtime, then pin that format.
+    from vllm.renderers.hf import _detect_content_format
+    content_format = _detect_content_format(native_chat_template(model_source), default="string")
+    args += ["--chat-template-content-format", content_format]
+
+
 def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path) -> list[str]:
     args = [
         "serve", model_source,
@@ -58,6 +83,8 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
     ]
     if os.environ.get("TT_API_KEY"):
         args += ["--middleware", "serve.BearerAuthMiddleware"]
+    if not is_certified_base_model(os.environ.get("TT_BASE_MODEL", "")):
+        args += ["--load-format", "safetensors"]
     base = os.environ["TT_BASE_MODEL"]
     # Upstream model-specific parsers, not a TT parser or model-size heuristic.
     parser = {
@@ -81,14 +108,7 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
                  "path": adapter_path, "base_model_name": base_name})]
     prompt = os.environ.get("TT_SYSTEM_PROMPT", "").strip()
     if prompt:
-        source = Path(model_source)
-        template_file = source / "chat_template.jinja"
-        if template_file.is_file():
-            template = template_file.read_text(encoding="utf-8")
-        else:
-            template = json.loads((source / "tokenizer_config.json").read_text())["chat_template"]
-        if not isinstance(template, str):
-            raise ValueError("Serving requires an unambiguous string chat template.")
+        template = native_chat_template(model_source)
         # Literal data, never prompt text interpolated as Jinja source. Preserve
         # tool calls/results verbatim while merging leading system context.
         prefix = "{%- set tt = namespace(system=" + json.dumps(prompt) + ", history=[]) -%}"
@@ -111,9 +131,14 @@ def build_vllm_args(model_source: str, adapter_path: str | None, temp_dir: Path)
 def prepare_model_source() -> str:
     base = os.environ["TT_BASE_MODEL"]
     revision = os.environ.get("TT_BASE_MODEL_REVISION")
-    assert_certified_base_model_revision(base, revision, "Serving base model revision")
-    configure_hugging_face_cache(os.environ.get("HF_HOME"))
     source = os.environ.get("TT_MODEL_SOURCE", base)
+    assert_certified_base_model_revision(
+        base,
+        revision,
+        "Serving base model revision",
+        local_source=source != base and Path(source).is_dir(),
+    )
+    configure_hugging_face_cache(os.environ.get("HF_HOME"))
     if not Path(source).is_dir():
         if source != base:
             raise ValueError("Explicit serving snapshot is not an existing directory.")
@@ -143,6 +168,7 @@ def run_server() -> None:
             model_source = prepare_model_source()
             adapter = resolve_adapter_path(os.environ.get("TT_MODEL_ARTIFACT"), Path(temp))
         args = build_vllm_args(model_source, adapter, Path(temp))
+        pin_native_content_format(args, model_source)
         from vllm.entrypoints.cli.main import main
         previous = sys.argv
         try:

package/training/adapter/src/train.py

@@ -27,6 +27,7 @@ from model_contract import (
     assert_certified_model_config,
     chat_template_kwargs,
     parse_lora_target_modules,
+    uncertified_load_kwargs,
 )
 from sft_data import IGNORE_INDEX, build_assistant_only_example
 
@@ -82,6 +83,7 @@ def assert_certified_request() -> None:
         str(requested_model),
         hp("base_model_revision"),
         "Training base model revision",
+        local_source=Path(resolve_model_source()).is_dir(),
     )
     loader = hp("model_loader", "causal_lm")
     if loader not in ("causal_lm", "image_text_to_text"):
@@ -199,6 +201,11 @@ def create_model_and_tokenizer(model_source: str) -> tuple[Any, Any, torch.dtype
         raise ValueError(f"{hp('base_model', CERTIFIED_BASE_MODEL)} tokenizer has no EOS token")
     if tokenizer.pad_token_id is None:
         tokenizer.pad_token = tokenizer.eos_token
+    if not tokenizer.chat_template:
+        raise ValueError(
+            f"{hp('base_model', CERTIFIED_BASE_MODEL)} has no chat template; "
+            "choose an instruct or chat variant of the model"
+        )
 
     dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
     # Multimodal checkpoints whose text tower is not registered in the
@@ -212,6 +219,7 @@ def create_model_and_tokenizer(model_source: str) -> tuple[Any, Any, torch.dtype
     model = model_class.from_pretrained(
         model_source,
         **source_kwargs,
+        **uncertified_load_kwargs(str(hp("base_model", CERTIFIED_BASE_MODEL))),
         dtype=dtype,
         device_map={"": torch.cuda.current_device()},
     )

package/training/decision/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "tuned-tensor-decision"
+version = "0.1.0"
+requires-python = ">=3.11"
+dependencies = [
+  "huggingface-hub>=0.36.0",
+  "laya==0.3.22",
+  "numpy>=2.0.0",
+  "safetensors>=0.4.5",
+  "torch>=2.11.0",
+  "transformers>=4.48.0",
+]
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu130", marker = "sys_platform == 'linux'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[tool.uv]
+package = false

package/training/decision/src/common.py

@@ -0,0 +1,127 @@
+"""Shared helpers for the typed decision-model engine (Laya-format checkpoints).
+
+Node owns the behavior spec, splits, scoring and reports. This package only
+loads a decision checkpoint, predicts labels for opaque IDs, and fine-tunes the
+checkpoint's encoder and decision head on labelled rows.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import torch
+
+CHECKPOINT_FILES = ("rl_agent_config.json", "model.safetensors", "tokenizer/*", "encoder/*")
+QUESTION_ID = "decision"
+
+
+def load_config(argv: list[str] | None = None) -> dict[str, Any]:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", required=True)
+    args = parser.parse_args(argv)
+    return json.loads(Path(args.config).read_text(encoding="utf-8"))
+
+
+def ensure_private_directory(path: str | Path) -> Path:
+    destination = Path(path)
+    destination.mkdir(parents=True, exist_ok=True, mode=0o700)
+    destination.chmod(0o700)
+    return destination
+
+
+def write_json(path: str | Path, value: Any) -> None:
+    write_text(path, json.dumps(value, indent=2) + "\n")
+
+
+def write_jsonl(path: str | Path, rows: list[Any]) -> None:
+    write_text(path, "".join(json.dumps(row, ensure_ascii=False) + "\n" for row in rows))
+
+
+def write_text(path: str | Path, content: str) -> None:
+    destination = Path(path)
+    ensure_private_directory(destination.parent)
+    descriptor, temporary_name = tempfile.mkstemp(dir=destination.parent, prefix=f".{destination.name}.", suffix=".tmp")
+    temporary = Path(temporary_name)
+    try:
+        with os.fdopen(descriptor, "w", encoding="utf-8") as output:
+            os.fchmod(output.fileno(), 0o600)
+            output.write(content)
+            output.flush()
+            os.fsync(output.fileno())
+        os.replace(temporary, destination)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def read_jsonl(path: str | Path) -> list[dict[str, Any]]:
+    rows = []
+    for number, line in enumerate(Path(path).read_text(encoding="utf-8").splitlines(), start=1):
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        if not isinstance(row, dict):
+            raise ValueError(f"{path}:{number} must be a JSON object")
+        rows.append(row)
+    return rows
+
+
+def pick_device(requested: str | None) -> torch.device:
+    if requested and requested != "auto":
+        device = torch.device(requested)
+        if device.type == "cuda" and not torch.cuda.is_available():
+            raise RuntimeError("Decision device cuda was requested, but PyTorch cannot see a CUDA GPU.")
+        if device.type == "mps" and not torch.backends.mps.is_available():
+            raise RuntimeError("Decision device mps was requested, but PyTorch cannot use Apple MPS.")
+        return device
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def option_labels(question: dict[str, Any]) -> list[str]:
+    """Answer labels in the model's option order; these are the strings Node scores."""
+    kind = question["type"]
+    criteria = question.get("criteria")
+    if kind == "choice":
+        return [str(label) for label in (criteria.keys() if isinstance(criteria, dict) else criteria)]
+    if kind == "noul":
+        return ["false", "true"]
+    if kind == "score":
+        return [str(index) for index in range(len(criteria))]
+    raise ValueError(f"Unsupported decision type: {kind!r}")
+
+
+def label_index(question: dict[str, Any], output: str) -> int:
+    labels = option_labels(question)
+    normalized = output.strip().lower()
+    for index, label in enumerate(labels):
+        if label.strip().lower() == normalized:
+            return index
+    raise ValueError(f"Expected output {output!r} is not one of the decision labels {labels}")
+
+
+def resolve_checkpoint_dir(model: str, revision: str | None) -> Path:
+    """Return a local checkpoint directory, downloading only the files a load needs."""
+    local = Path(model)
+    if local.is_dir():
+        return local
+    if local.is_absolute() or model.startswith(("./", "../")):
+        raise FileNotFoundError(f"Decision checkpoint directory not found: {model}")
+    from huggingface_hub import snapshot_download
+
+    return Path(snapshot_download(model, revision=revision, allow_patterns=list(CHECKPOINT_FILES)))
+
+
+def load_agent(model: str, revision: str | None, device: torch.device):
+    from laya import Agent
+
+    checkpoint = resolve_checkpoint_dir(model, revision)
+    # Load the resolved directory so Hub and local checkpoints take the same path.
+    agent = Agent(str(checkpoint), device=str(device))
+    return agent, checkpoint

package/training/decision/src/evaluate.py

@@ -0,0 +1,70 @@
+"""Predict one typed decision per input row.
+
+The config names only opaque row IDs and inputs. Reference labels stay in Node,
+which joins these predictions to trusted expected outputs before scoring.
+"""
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import numpy as np
+
+from common import QUESTION_ID, load_agent, load_config, option_labels, pick_device, read_jsonl, write_json, write_jsonl
+
+
+def answer_distribution(question: dict[str, Any], answer: dict[str, Any]) -> dict[str, float]:
+    labels = option_labels(question)
+    if question["type"] == "noul":
+        p_true = float(answer["noul"])
+        return {"false": 1.0 - p_true, "true": p_true}
+    # Laya keys choice probabilities by label and score probabilities by level index.
+    return {label: float(answer["probabilities"][label]) for label in labels}
+
+
+def predict_rows(agent, question: dict[str, Any], rows: list[dict[str, Any]], batch_size: int) -> list[dict[str, Any]]:
+    predictions: list[dict[str, Any]] = []
+    for start in range(0, len(rows), batch_size):
+        batch = rows[start:start + batch_size]
+        began = time.perf_counter()
+        results = agent.predict_batch([row["input"] for row in batch], {QUESTION_ID: question}, batch_size=batch_size)
+        latency_ms = int(round((time.perf_counter() - began) * 1000 / max(1, len(batch))))
+        for row, result in zip(batch, results):
+            distribution = answer_distribution(question, result["answers"][QUESTION_ID])
+            labels = list(distribution)
+            values = np.array([distribution[label] for label in labels])
+            predictions.append({
+                "id": row["id"],
+                "prediction": labels[int(values.argmax())],
+                "probabilities": {label: round(float(value), 6) for label, value in distribution.items()},
+                "confidence": round(float(values.max()), 6),
+                "latency_ms": latency_ms,
+            })
+    return predictions
+
+
+def main() -> None:
+    config = load_config()
+    device = pick_device(config.get("device"))
+    rows = read_jsonl(config["inputs_path"])
+    for row in rows:
+        if set(row) != {"id", "input"}:
+            raise ValueError("Decision evaluation rows must contain only id and input.")
+    began = time.perf_counter()
+    agent, checkpoint = load_agent(config["model"], config.get("revision"), device)
+    load_seconds = time.perf_counter() - began
+    predictions = predict_rows(agent, config["question"], rows, int(config.get("batch_size", 16)))
+    write_jsonl(config["output_path"], predictions)
+    write_json(config["metrics_path"], {
+        "ok": True,
+        "model": config["model"],
+        "revision": config.get("revision"),
+        "checkpoint_dir": str(checkpoint),
+        "device": str(agent.device),
+        "rows": len(predictions),
+        "load_seconds": round(load_seconds, 3),
+    })
+
+
+if __name__ == "__main__":
+    main()