Upload Tercet-R-1.1 (stage-3 EMA @ step 5300)

README.md

@@ -0,0 +1,218 @@
+---
+license: apache-2.0
+language:
+  - en
+tags:
+  - text-generation
+  - causal-lm
+  - pytorch
+  - sft
+  - instruction-tuned
+  - chat
+  - reasoning
+  - tool-use
+  - hybrid
+  - gated-deltanet
+  - gqa
+  - tercet
+pipeline_tag: text-generation
+library_name: tiny_gdn
+datasets:
+  - HuggingFaceFW/fineweb-edu
+  - HuggingFaceTB/smoltalk2
+  - HuggingFaceTB/smoltalk
+  - nvidia/Llama-Nemotron-Post-Training-Dataset
+  - nvidia/Nemotron-Post-Training-Dataset-v1
+  - nvidia/Nemotron-Agentic-v1
+  - nvidia/Nemotron-SFT-Agentic-v2
+  - nvidia/Nemotron-SFT-Instruction-Following-Chat-v2
+  - nvidia/Nemotron-Instruction-Following-Chat-v1
+  - nvidia/Nemotron-Cascade-2-SFT-Data
+  - Agent-Ark/Toucan-1.5M
+  - interstellarninja/hermes_reasoning_tool_use
+  - NousResearch/Hermes-3-Dataset
+  - allenai/WildChat-4.8M
+  - open-thoughts/OpenThoughts3-1.2M
+  - teknium/OpenHermes-2.5
+  - open-r1/Mixture-of-Thoughts
+  - CohereLabs/aya_dataset
+  - NousResearch/hermes-function-calling-v1
+  - simplescaling/s1K-1.1
+  - allenai/tulu-3-sft-personas-instruction-following
+  - Salesforce/xlam-function-calling-60k
+  - THUDM/LongAlign-10k
+base_model: kerzgrr/Tercet-R-1.0
+---
+
+<div align="center">
+
+# Tercet-R-1.1
+
+### Reasoning + tool-call chat model (~502M) — latest Tercet-R snapshot
+
+[![Model](https://img.shields.io/badge/Model-~502M_params-blue)](.)
+[![Stage](https://img.shields.io/badge/Stage-SFT_stage--3-green.svg)](.)
+[![License](https://img.shields.io/badge/License-Apache_2.0-green.svg)](LICENSE)
+[![Demo](https://img.shields.io/badge/Space-tercet--r--1.1--demo-indigo.svg)](https://huggingface.co/spaces/kerzgrr/tercet-r-1.1-demo)
+
+*A ~502M hybrid GDN-2 + GQA model, continued on multi-turn tools and instruction following*
+
+</div>
+
+---
+
+## What this is
+
+**Tercet-R-1.1** is the latest public reasoning / tool-use checkpoint in the Tercet-R line.
+
+- Previous public snapshot: [`kerzgrr/Tercet-R-1.0`](https://huggingface.co/kerzgrr/Tercet-R-1.0)
+- Base (pretrain): [`kerzgrr/Tercet-base`](https://huggingface.co/kerzgrr/Tercet-base)
+- Live demo: [`kerzgrr/tercet-r-1.1-demo`](https://huggingface.co/spaces/kerzgrr/tercet-r-1.1-demo)
+- Hub weights are the **EMA** snapshot in **bfloat16**
+
+This upload is the latest **complete** stage-3 SFT checkpoint (optimizer step **5,300**). The stage-3 run is still in progress toward a much longer schedule; 1.1 is the snapshot that existed when this repo was published.
+
+Initialization was the instruct-run EMA at step 5,000 (`tercet-r-instruct`), not the public 1.0 file (that release packaged instruct step 4,500). The chat contract is unchanged.
+
+---
+
+## Chat contract
+
+### Thinking
+
+Each assistant turn is prefixed with a **zero-loss control token**:
+
+| Mode | Prefix | Typical body |
+|------|--------|----------------|
+| think | `<\|think\|>\n` | `<think>…</think>` then the answer |
+| no-think | `<\|no_think\|>\n` | answer only |
+
+`inference.py` streams the `<think>` region live (dim yellow) and hides the control tokens.
+
+### Tool calls (SmolTalk JSON)
+
+```
+<tool_call>
+{"name": "web-search", "arguments": {"query": "…"}}
+</tool_call>
+```
+
+Stage-3 also saw `stateful_python_code_exec` leftovers from earlier Cascade-2 SFT. `inference.py` auto-runs:
+
+| Built-in | Tool name | Observation |
+|----------|-----------|-------------|
+| `--tools web-search` | `web-search` | Tavily-shaped JSON (Tavily if `TAVILY_API_KEY` is set, else DuckDuckGo + Wikipedia) |
+| `--tools python` | `stateful_python_code_exec` | Jupyter-style stdout / last value from a restricted interpreter |
+| `--tools calculator` | `calculator` | Numeric result of a math expression |
+
+### Tool results
+
+Each observation is a `tool` (or `user`) turn prefixed with:
+
+```
+<|tool_response|>
+{observation}
+```
+
+---
+
+## Install & run
+
+```bash
+pip install torch safetensors tokenizers huggingface_hub
+hf download kerzgrr/Tercet-R-1.1 inference.py --local-dir .
+python inference.py
+```
+
+```bash
+python inference.py --prompt "What is the capital of France?"
+python inference.py --tools web-search,python,calculator
+python inference.py --no-think --prompt "Reply in one sentence."
+```
+
+`inference.py` auto-downloads weights / tokenizer / `tiny_gdn/` and auto-installs pinned `flash-linear-attention`. Git is required on `PATH`.
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--prompt` | — | One-shot user message |
+| `--system` | — | System prompt, used verbatim |
+| `--think` / `--no-think` | think | Assistant control prefix |
+| `--tools` | — | Built-ins: `web-search`, `python`, `calculator` (comma-separated) |
+| `--temperature` | `0.7` | Sampling temperature |
+| `--max-new-tokens` | `4096` | Max generation length |
+| `--device` | `cuda` if available | `cuda` / `cpu` |
+
+Interactive commands: `/think` `/no_think` `/system …` `/reset` `/exit`.
+
+---
+
+## Model architecture
+
+Same TinyGDN hybrid as [`Tercet-base`](https://huggingface.co/kerzgrr/Tercet-base) (501,635,264 parameters):
+
+| | |
+|--|--|
+| **Layers** | 32 (GDN-2 ×3 + GQA every 4th) |
+| **Hidden** | 1,024 |
+| **MLP** | SwiGLU 2,624 |
+| **Attention** | 8 Q / 2 KV, head dim 128, partial RoPE |
+| **Linear** | Gated DeltaNet-2, 8 heads × 128 |
+| **Vocab** | 49,152 BPE |
+| **Context** | 16,384 |
+
+---
+
+## Training
+
+| Stage | Details |
+|-------|---------|
+| **Base** | [HuggingFaceFW/fineweb-edu](https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu) → [`Tercet-base`](https://huggingface.co/kerzgrr/Tercet-base) |
+| **Mid + instruct SFT** | [HuggingFaceTB/smoltalk2](https://huggingface.co/datasets/HuggingFaceTB/smoltalk2) Mid / SFT → instruct EMA @ step 5,000 |
+| **Stage-3 SFT** | Multi-turn mix: Nemotron PTD v1, Agentic v1/v2 (including `web-search`), IF-Chat v1/v2, Cascade-2 instruction following, Toucan 1.5M, Hermes-3, Hermes reasoning tool-use, WildChat-4.8M. Seq 16,384, AdamW 3×10⁻⁵, 27.5 hours so far, 602M training tokens |
+| **Checkpoint** | optimizer step 5,300 (latest complete stage-3 snapshot) |
+| **Weights** | EMA (this repo's `model.safetensors`) |
+| **Val loss (EMA)** | 2.027 (ppl 7.59) on the stage-3 holdout — a harder mix than the instruct-only 1.0 val set |
+
+Stage-3 does **not** re-include Cascade-2 `math/math_tool.jsonl`. The restricted Python tool in `inference.py` is there so leftover `stateful_python_code_exec` calls still compute instead of failing open.
+
+---
+
+## Limitations
+
+- **In-progress run**: this is a mid-stage-3 snapshot, not a finished 1.1 train
+- **Scale**: ~502M is a research / edge model, not a frontier system
+- Requires `flash-linear-attention`; not GGUF / llama.cpp compatible today
+- The Python tool is a restricted interpreter (math-oriented imports only)
+
+---
+
+## Model family
+
+| Model | Stage | Hub |
+|-------|-------|-----|
+| Tercet-base | Pretrain | [`kerzgrr/Tercet-base`](https://huggingface.co/kerzgrr/Tercet-base) |
+| Tercet | SFT chat | [`kerzgrr/Tercet`](https://huggingface.co/kerzgrr/Tercet) |
+| Tercet-R-1.0 | SFT reasoning + tools | [`kerzgrr/Tercet-R-1.0`](https://huggingface.co/kerzgrr/Tercet-R-1.0) |
+| **Tercet-R-1.1** | **Stage-3 SFT (latest)** | **this repo** |
+| Demo | ZeroGPU Space | [`kerzgrr/tercet-r-1.1-demo`](https://huggingface.co/spaces/kerzgrr/tercet-r-1.1-demo) |
+
+---
+
+## Citation
+
+```bibtex
+@misc{tercetr11,
+  title={Tercet-R-1.1: A 502M Hybrid GDN-2 + GQA Reasoning Model},
+  author={kerzgrr},
+  year={2026},
+  url={https://huggingface.co/kerzgrr/Tercet-R-1.1}
+}
+```
+
+---
+
+<div align="center">
+
+*R is for reasoning.*
+
+</div>

chat_template.jinja

@@ -0,0 +1,86 @@
+{%- set ns = namespace(xml_tools=none, tools_emitted=false) -%}
+{%- if xml_tools is defined and xml_tools -%}
+{%- set ns.xml_tools = xml_tools -%}
+{%- elif tools is defined and tools -%}
+{%- set ns.xml_tools = tools -%}
+{%- endif -%}
+{%- set tools_preamble = 'You may call one or more functions to assist with the user query.\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>\n' -%}
+{%- set tools_epilogue = '</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{"name": <function-name>, "arguments": <args-json-object>}\n</tool_call>' -%}
+{%- for message in messages -%}
+{%- if loop.first -%}{{- bos_token -}}{%- endif -%}
+{%- if loop.first and ns.xml_tools and message['role'] != 'system' -%}
+{{- '<|im_start|>system\n' + tools_preamble -}}
+{%- for tool in ns.xml_tools -%}
+{%- if tool is string -%}{{- (tool | replace('<tools>', '') | replace('</tools>', '') | trim) + '\n' -}}
+{%- else -%}{{- tool | tojson + '\n' -}}
+{%- endif -%}
+{%- endfor -%}
+{{- tools_epilogue + '<|im_end|>\n' -}}
+{%- set ns.tools_emitted = true -%}
+{%- endif -%}
+{%- set raw_role = message['role'] -%}
+{%- set role = 'user' if raw_role == 'tool' or raw_role == 'function' else raw_role -%}
+{%- set content_text = namespace(value='') -%}
+{%- if message['content'] is string -%}
+{%- set content_text.value = message['content'] -%}
+{%- elif message['content'] is iterable -%}
+{%- for item in message['content'] -%}
+{%- if item['type'] == 'text' -%}{%- set content_text.value = content_text.value + item['text'] -%}{%- endif -%}
+{%- endfor -%}
+{%- endif -%}
+{%- set is_tool = raw_role == 'tool' or raw_role == 'function' -%}
+{%- set prev_is_tool = loop.previtem is defined and (loop.previtem['role'] == 'tool' or loop.previtem['role'] == 'function') -%}
+{%- set next_is_tool = loop.nextitem is defined and (loop.nextitem['role'] == 'tool' or loop.nextitem['role'] == 'function') -%}
+{%- if raw_role == 'system' and not (content_text.value | trim) and not (ns.xml_tools and not ns.tools_emitted) -%}
+{%- else -%}
+{%- if is_tool and prev_is_tool -%}
+{{- '\n' -}}
+{%- else -%}
+{{- '<|im_start|>' + role + '\n' -}}
+{%- endif -%}
+{%- if raw_role == 'assistant' and enable_thinking is defined -%}
+{{- ('<|think|>\n' if enable_thinking else '<|no_think|>\n') -}}
+{%- endif -%}
+{%- if is_tool -%}
+{{- '<|tool_response|>\n' + (content_text.value | replace('<tool_response>', '') | replace('</tool_response>', '') | trim) -}}
+{%- elif raw_role == 'system' -%}
+{{- content_text.value | replace('/system_override', '') | replace('/no_think', '') | replace('/think', '') | trim -}}
+{%- else -%}
+{{- content_text.value -}}
+{%- endif -%}
+{%- if raw_role == 'system' and ns.xml_tools and not ns.tools_emitted and '<tools>' not in content_text.value -%}
+{%- if content_text.value | trim -%}{{- '\n\n' -}}{%- endif -%}
+{{- tools_preamble -}}
+{%- for tool in ns.xml_tools -%}
+{%- if tool is string -%}{{- (tool | replace('<tools>', '') | replace('</tools>', '') | trim) + '\n' -}}
+{%- else -%}{{- tool | tojson + '\n' -}}
+{%- endif -%}
+{%- endfor -%}
+{{- tools_epilogue -}}
+{%- set ns.tools_emitted = true -%}
+{%- endif -%}
+{%- if raw_role == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] -%}
+{%- if '<tool_call>' not in content_text.value -%}
+{%- for tool_call in message['tool_calls'] -%}
+{%- set fn = tool_call['function'] if tool_call['function'] is defined else tool_call -%}
+{%- if loop.first and not (content_text.value | trim) -%}
+{{- '<tool_call>\n{"name": "' + fn['name'] + '", "arguments": ' -}}
+{%- else -%}
+{{- '\n<tool_call>\n{"name": "' + fn['name'] + '", "arguments": ' -}}
+{%- endif -%}
+{%- if fn['arguments'] is string -%}{{- fn['arguments'] -}}
+{%- else -%}{{- fn['arguments'] | tojson -}}
+{%- endif -%}
+{{- '}\n</tool_call>' -}}
+{%- endfor -%}
+{%- endif -%}
+{%- endif -%}
+{%- if not (is_tool and next_is_tool) -%}
+{{- '<|im_end|>\n' -}}
+{%- endif -%}
+{%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+{{- '<|im_start|>assistant\n' -}}
+{%- if enable_thinking is defined -%}{{- ('<|think|>\n' if enable_thinking else '<|no_think|>\n') -}}{%- endif -%}
+{%- endif -%}

config.json

@@ -0,0 +1,86 @@
+{
+  "allow_negative_eigenvalues": false,
+  "architecture": "TinyGDNForCausalLM",
+  "architectures": [
+    "TinyGDNForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attention_head_dim": 128,
+  "base_model": "kerzgrr/Tercet-R-1.0",
+  "bos_token_id": 0,
+  "chat_template": "smollm3",
+  "checkpoint_step": 5300,
+  "eos_token_id": 1,
+  "full_attention_interval": 4,
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "intermediate_size": 2624,
+  "layer_types": [
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention",
+    "gdn2",
+    "gdn2",
+    "gdn2",
+    "full_attention"
+  ],
+  "linear_conv_kernel_dim": 4,
+  "linear_expand_v": 1.0,
+  "linear_head_dim": 128,
+  "linear_num_heads": 8,
+  "linear_num_value_heads": 8,
+  "max_position_embeddings": 32768,
+  "model_family": "Tercet-R",
+  "model_name": "Tercet-R-1.1",
+  "model_type": "tiny_gdn",
+  "mtp_adapter_rank": 128,
+  "mtp_loss_weight": 0.0,
+  "mtp_num_heads": 0,
+  "num_attention_heads": 8,
+  "num_hidden_layers": 32,
+  "num_key_value_heads": 2,
+  "pad_token_id": 2,
+  "partial_rotary_factor": 0.5,
+  "reasoning": true,
+  "release": "1.1",
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sft_hours": 27.491864080583593,
+  "sft_val_loss_ema": 2.0271428378822587,
+  "sft_val_ppl_ema": 7.592362728419498,
+  "shared_layer_indices": [],
+  "stage": "sft",
+  "tie_word_embeddings": true,
+  "tool_calls": "smoltalk_json",
+  "torch_dtype": "bfloat16",
+  "training_sequence_length": 16384,
+  "transformers_version": "4.45.0",
+  "unk_token_id": 3,
+  "vocab_size": 49152,
+  "weights": "ema"
+}

inference.py

@@ -0,0 +1,544 @@
+#!/usr/bin/env python3
+"""Standalone chat inference for kerzgrr/Tercet-R-1.1.
+
+Streams a live reasoning trace (`<think>` / `<|think|>`) and runs SmolTalk
+tool-call rounds. Built-in `web-search`, `python`, and `calculator` tools
+execute automatically. Other observations are sent as `<|tool_response|>`.
+
+Examples:
+  python inference.py --prompt "What is the capital of France?"
+  python inference.py --tools web-search,python,calculator
+  python inference.py --no-think --prompt "Reply in one sentence."
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import platform
+import shutil
+import subprocess
+import sys
+import time
+import warnings
+from dataclasses import fields
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import load_file
+from tokenizers import Tokenizer
+
+
+def _silence_runtime_warnings() -> None:
+    patterns = (
+        r"tl\.make_block_ptr is deprecated",
+        r"Memory efficient kernel not used because",
+        r"Memory Efficient attention has been runtime disabled",
+        r"Flash attention kernel not used because",
+        r"Torch was not compiled with flash attention",
+        r"cuDNN attention kernel not used because",
+        r"cuDNN attention has been runtime disabled",
+    )
+    for pattern in patterns:
+        warnings.filterwarnings("ignore", message=pattern)
+
+
+REPO_ID = "kerzgrr/Tercet-R-1.1"
+FLA_COMMIT = "cbb0a72efb55c18ca0ef4f298298317573ad2cb3"
+FLA_REPO = "https://github.com/fla-org/flash-linear-attention.git"
+PATCH_FILES = (
+    "fla/__init__.py",
+    "fla/ops/__init__.py",
+    "fla/layers/__init__.py",
+    "fla/ops/simple_gla/__init__.py",
+)
+TINY_GDN_FILES = (
+    "tiny_gdn/__init__.py",
+    "tiny_gdn/config.py",
+    "tiny_gdn/model.py",
+    "tiny_gdn/tools.py",
+    "tiny_gdn/smoltalk_chat.py",
+    "tiny_gdn/chatml.py",
+    "tiny_gdn/cli_chat.py",
+    "tiny_gdn/detokenize.py",
+    "tiny_gdn/code_exec.py",
+    "tiny_gdn/web_search.py",
+)
+
+
+def _download(filename: str, local_dir: Path | None) -> Path:
+    from huggingface_hub import hf_hub_download
+
+    return Path(
+        hf_hub_download(
+            repo_id=REPO_ID,
+            filename=filename,
+            local_dir=str(local_dir) if local_dir else None,
+        )
+    )
+
+
+def _run(cmd: list[str], *, cwd: Path | None = None, env: dict | None = None) -> None:
+    print("+", " ".join(cmd), flush=True)
+    merged = os.environ.copy()
+    if env:
+        merged.update(env)
+    merged.setdefault("PYTHONUTF8", "1")
+    merged.setdefault("PYTHONIOENCODING", "utf-8")
+    subprocess.check_call(cmd, cwd=str(cwd) if cwd else None, env=merged)
+
+
+def _pip_install(*args: str) -> None:
+    _run([sys.executable, "-m", "pip", "install", *args])
+
+
+def _fla_importable() -> tuple[bool, str]:
+    try:
+        from fla.layers.gdn2 import GatedDeltaNet2  # noqa: F401
+    except Exception as error:  # noqa: BLE001
+        return False, str(error)
+    return True, ""
+
+
+def _cache_root() -> Path:
+    override = os.environ.get("MONOSTICH_CACHE")
+    if override:
+        path = Path(override).expanduser().resolve()
+    else:
+        path = Path.home() / ".cache" / "tercet-r"
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def _ensure_git() -> None:
+    if shutil.which("git") is None:
+        raise RuntimeError(
+            "git is required to auto-install flash-linear-attention. "
+            "Install Git and ensure it is on PATH."
+        )
+
+
+def _apply_windows_fla_patches(fla_root: Path, local_dir: Path | None) -> None:
+    print("Applying Windows FLA import patches from the Hub …", flush=True)
+    for relative in PATCH_FILES:
+        source = _download(f"windows_fla_patches/{relative}", local_dir)
+        target = fla_root / relative
+        target.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(source, target)
+        print(f"  patched {relative}", flush=True)
+
+
+def _install_fla(local_dir: Path | None) -> None:
+    print("flash-linear-attention missing/broken — installing automatically …", flush=True)
+    _pip_install("einops", "numpy")
+    if platform.system() != "Windows":
+        _pip_install("--no-deps", f"git+{FLA_REPO}@{FLA_COMMIT}")
+        return
+
+    _ensure_git()
+    fla_root = _cache_root() / "flash-linear-attention"
+    if (fla_root / ".git").is_dir():
+        _run(["git", "fetch", "--depth", "1", "origin", FLA_COMMIT], cwd=fla_root)
+        _run(["git", "checkout", "--force", FLA_COMMIT], cwd=fla_root)
+    else:
+        if fla_root.exists():
+            shutil.rmtree(fla_root)
+        _run(["git", "clone", "--filter=blob:none", FLA_REPO, str(fla_root)])
+        _run(["git", "fetch", "--depth", "1", "origin", FLA_COMMIT], cwd=fla_root)
+        _run(["git", "checkout", "--force", FLA_COMMIT], cwd=fla_root)
+
+    _apply_windows_fla_patches(fla_root, local_dir)
+    _pip_install("--no-build-isolation", "--no-deps", "-e", str(fla_root))
+
+
+def _ensure_fla(local_dir: Path | None) -> None:
+    ok, error = _fla_importable()
+    if ok:
+        return
+    print(f"FLA not ready ({error})", flush=True)
+    try:
+        _install_fla(local_dir)
+    except Exception as install_error:  # noqa: BLE001
+        raise RuntimeError(
+            "Automatic flash-linear-attention install failed.\n"
+            f"Original import error: {error}\n"
+            f"Install error: {install_error}"
+        ) from install_error
+
+    for name in list(sys.modules):
+        if name == "fla" or name.startswith("fla."):
+            del sys.modules[name]
+
+    ok, error = _fla_importable()
+    if not ok:
+        raise RuntimeError(
+            "flash-linear-attention installed but still failed to import "
+            f"GatedDeltaNet2: {error}"
+        )
+    print("flash-linear-attention ready.", flush=True)
+
+
+def _ensure_tiny_gdn(local_dir: Path | None) -> Path:
+    here = Path(__file__).resolve().parent
+    if (here / "tiny_gdn" / "cli_chat.py").is_file():
+        return here
+    if local_dir and (local_dir / "tiny_gdn" / "cli_chat.py").is_file():
+        return local_dir
+    for name in TINY_GDN_FILES:
+        _download(name, local_dir)
+    return _download("tiny_gdn/__init__.py", local_dir).parent.parent
+
+
+def _sample(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    generator: torch.Generator,
+) -> int:
+    logits = logits.float()
+    if temperature <= 1e-5:
+        return int(torch.argmax(logits).item())
+    logits = logits / temperature
+    if 0 < top_k < logits.shape[-1]:
+        threshold = torch.topk(logits, top_k).values[-1]
+        logits = logits.masked_fill(logits < threshold, -torch.inf)
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        probs = torch.softmax(sorted_logits, dim=-1)
+        remove = torch.cumsum(probs, dim=-1) > top_p
+        remove[1:] = remove[:-1].clone()
+        remove[0] = False
+        sorted_logits = sorted_logits.masked_fill(remove, -torch.inf)
+        logits = torch.full_like(logits, -torch.inf)
+        logits.scatter_(0, sorted_indices, sorted_logits)
+    probs = torch.softmax(logits, dim=-1).cpu()
+    return int(torch.multinomial(probs, 1, generator=generator).item())
+
+
+def _apply_repetition_penalty(
+    logits: torch.Tensor,
+    token_ids: list[int],
+    penalty: float,
+    window: int,
+) -> torch.Tensor:
+    if penalty == 1.0 or not token_ids:
+        return logits
+    recent = token_ids[-window:] if window > 0 else token_ids
+    unique = torch.tensor(list(set(recent)), dtype=torch.long, device=logits.device)
+    score = logits[unique]
+    logits[unique] = torch.where(score > 0, score / penalty, score * penalty)
+    return logits
+
+
+@torch.inference_mode()
+def generate_turn(
+    model,
+    tokenizer: Tokenizer,
+    encode_chatml_messages,
+    IncrementalUtf8Decoder,
+    parse_tool_calls,
+    streamer,
+    messages: list[dict[str, Any]],
+    *,
+    tools: Any | None,
+    enable_thinking: bool,
+    max_new_tokens: int,
+    context_length: int,
+    temperature: float,
+    top_p: float,
+    top_k: int,
+    repetition_penalty: float,
+    repetition_window: int,
+    seed: int,
+    device: torch.device,
+) -> tuple[str, tuple]:
+    prompt_ids = encode_chatml_messages(
+        tokenizer,
+        messages,
+        tools=tools,
+        enable_thinking=enable_thinking,
+    )
+    eos_id = int(model.config.eos_token_id)
+    im_end = tokenizer.token_to_id("<|im_end|>")
+    stop = {eos_id}
+    if im_end is not None:
+        stop.add(im_end)
+
+    token_ids = list(prompt_ids[-context_length:])
+    generated: list[int] = []
+    decoder = IncrementalUtf8Decoder(tokenizer)
+    stop_reason = "max_new_tokens"
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed)
+    started = time.perf_counter()
+
+    input_ids = torch.tensor([token_ids], dtype=torch.long, device=device)
+    output = model(
+        input_ids,
+        return_logits=True,
+        logits_to_keep=1,
+        use_cache=True,
+    )
+
+    for step in range(max_new_tokens):
+        if output.logits is None:
+            raise RuntimeError("Model returned no logits")
+        next_logits = _apply_repetition_penalty(
+            output.logits[0, -1],
+            token_ids,
+            repetition_penalty,
+            repetition_window,
+        )
+        next_id = _sample(
+            next_logits,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            generator=generator,
+        )
+        if next_id in stop:
+            stop_reason = "stop"
+            break
+        token_ids.append(next_id)
+        generated.append(next_id)
+        decoder.push(next_id)
+        if streamer is not None:
+            streamer.update(decoder.text)
+        if step + 1 < max_new_tokens:
+            if output.past_key_values is None:
+                raise RuntimeError("Model returned no decode cache")
+            input_ids = torch.tensor([[next_id]], dtype=torch.long, device=device)
+            output = model(
+                input_ids,
+                return_logits=True,
+                logits_to_keep=1,
+                past_key_values=output.past_key_values,
+                use_cache=True,
+            )
+
+    decoded = decoder.finalize()
+    if streamer is not None:
+        streamer.finish()
+    elapsed = time.perf_counter() - started
+    tps = len(generated) / max(elapsed, 1e-9)
+    print(
+        f"[done] tokens={len(generated)} stop={stop_reason} {tps:.1f} tok/s",
+        file=sys.stderr,
+        flush=True,
+    )
+    return decoded, tuple(parse_tool_calls(decoded))
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Tercet-R-1.1 chat inference")
+    parser.add_argument("--prompt", default=None, help="Single user prompt")
+    parser.add_argument("--system", default="", help="Optional system prompt")
+    parser.add_argument(
+        "--think",
+        dest="enable_thinking",
+        action="store_true",
+        default=True,
+    )
+    parser.add_argument(
+        "--no-think",
+        dest="enable_thinking",
+        action="store_false",
+    )
+    parser.add_argument(
+        "--tools",
+        default="",
+        help="Built-in tools: web-search, python, calculator (comma-separated)",
+    )
+    parser.add_argument("--tools-json", default=None)
+    parser.add_argument("--max-tool-rounds", type=int, default=8)
+    parser.add_argument("--max-new-tokens", type=int, default=4096)
+    parser.add_argument("--temperature", type=float, default=0.7)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--top-k", type=int, default=50)
+    parser.add_argument("--repetition-penalty", type=float, default=1.08)
+    parser.add_argument("--repetition-window", type=int, default=256)
+    parser.add_argument("--context-length", type=int, default=16384)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument(
+        "--device",
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        choices=["cuda", "cpu"],
+    )
+    parser.add_argument("--no-stream", action="store_true")
+    parser.add_argument("--no-color", action="store_true")
+    parser.add_argument("--local-dir", default=None)
+    parser.add_argument("--repo-id", default=REPO_ID)
+    return parser.parse_args()
+
+
+def main() -> int:
+    _silence_runtime_warnings()
+    args = parse_args()
+    global REPO_ID
+    REPO_ID = args.repo_id
+    local_dir = Path(args.local_dir).resolve() if args.local_dir else None
+
+    print(f"Loading Tercet-R-1.1 from huggingface.co/{REPO_ID} …", flush=True)
+    try:
+        package_root = _ensure_tiny_gdn(local_dir)
+    except Exception as error:  # noqa: BLE001
+        print(f"Failed to resolve tiny_gdn package: {error}", file=sys.stderr)
+        return 1
+
+    if str(package_root) not in sys.path:
+        sys.path.insert(0, str(package_root))
+
+    try:
+        _ensure_fla(local_dir)
+    except Exception as error:  # noqa: BLE001
+        print(str(error), file=sys.stderr)
+        return 1
+
+    try:
+        from tiny_gdn import TinyGDNConfig, TinyGDNForCausalLM
+        from tiny_gdn.chatml import encode_chatml_messages
+        from tiny_gdn.cli_chat import (
+            ChatLoopState,
+            LiveReasoningStreamer,
+            append_tool_round,
+            apply_slash_command,
+            prompt_tool_results,
+            resolve_cli_tools,
+        )
+        from tiny_gdn.detokenize import IncrementalUtf8Decoder
+        from tiny_gdn.tools import parse_tool_calls
+    except ImportError as error:
+        print(f"Could not import tiny_gdn chat stack: {error}", file=sys.stderr)
+        return 1
+
+    try:
+        raw_json = (
+            Path(args.tools_json).read_text(encoding="utf-8")
+            if args.tools_json
+            else None
+        )
+        tools = resolve_cli_tools(args.tools, raw_json)
+    except (OSError, ValueError, json.JSONDecodeError) as error:
+        print(error, file=sys.stderr)
+        return 1
+
+    weights_path = _download("model.safetensors", local_dir)
+    tok_path = _download("tokenizer.json", local_dir)
+    cfg_path = _download("config.json", local_dir)
+
+    raw = json.loads(cfg_path.read_text(encoding="utf-8"))
+    allowed = {item.name for item in fields(TinyGDNConfig)}
+    payload = {key: value for key, value in raw.items() if key in allowed}
+    if "shared_layer_indices" in payload:
+        payload["shared_layer_indices"] = tuple(payload["shared_layer_indices"])
+    config = TinyGDNConfig(**payload)
+
+    device = torch.device(args.device)
+    if device.type == "cuda" and not torch.cuda.is_available():
+        print("CUDA requested but unavailable; falling back to CPU.", flush=True)
+        device = torch.device("cpu")
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+
+    print(
+        f"Building TinyGDN ({config.num_hidden_layers}L / {config.hidden_size}d) "
+        f"on {device} …",
+        flush=True,
+    )
+    model = TinyGDNForCausalLM(config)
+    state = load_file(str(weights_path), device="cpu")
+    model.load_state_dict(state, strict=True)
+    del state
+    model = model.to(device=device, dtype=dtype)
+    model.eval()
+    model.requires_grad_(False)
+
+    tokenizer = Tokenizer.from_file(str(tok_path))
+    context_length = min(args.context_length, config.max_position_embeddings)
+    color = not args.no_color and sys.stdout.isatty()
+
+    def read_line(prompt: str) -> str:
+        return input(prompt)
+
+    def run_turn(state: ChatLoopState) -> None:
+        for round_index in range(max(1, args.max_tool_rounds + 1)):
+            streamer = (
+                None
+                if args.no_stream
+                else LiveReasoningStreamer(sys.stdout, color=color)
+            )
+            if streamer is not None:
+                print("assistant> ", end="", flush=True)
+            text, calls = generate_turn(
+                model,
+                tokenizer,
+                encode_chatml_messages,
+                IncrementalUtf8Decoder,
+                parse_tool_calls,
+                streamer,
+                state.messages,
+                tools=state.tools,
+                enable_thinking=state.enable_thinking,
+                max_new_tokens=args.max_new_tokens,
+                context_length=context_length,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                top_k=args.top_k,
+                repetition_penalty=args.repetition_penalty,
+                repetition_window=args.repetition_window,
+                seed=args.seed,
+                device=device,
+            )
+            if args.no_stream:
+                print(f"assistant> {text}")
+            if not calls or round_index >= args.max_tool_rounds:
+                state.messages.append({"role": "assistant", "content": text})
+                return
+            results = prompt_tool_results(calls, read_line=read_line)
+            append_tool_round(state.messages, text, results)
+
+    state = ChatLoopState(
+        system=args.system,
+        enable_thinking=args.enable_thinking,
+        tools=tools,
+    )
+    state.reset()
+
+    if args.prompt is not None:
+        state.add_user(args.prompt)
+        run_turn(state)
+        return 0
+
+    print(
+        "Interactive Tercet-R-1.1 chat. "
+        "Commands: /think  /no_think  /system …  /reset  /exit\n"
+        "web-search, python, and calculator run automatically when enabled. "
+        "Other <tool_call> observations can be pasted at the prompt.",
+        flush=True,
+    )
+    while True:
+        try:
+            user_input = input("user> ")
+        except (EOFError, KeyboardInterrupt):
+            print()
+            return 0
+        text = user_input.strip()
+        if not text:
+            continue
+        status = apply_slash_command(state, text)
+        if status == "exit":
+            return 0
+        if status is not None:
+            print(f"({status})", flush=True)
+            continue
+        state.add_user(text)
+        run_turn(state)
+        print(flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a3bd09e50ef1d40c70ad4fd1ab7bfca3a553f600bd08a0fffec0ea5e47bda886
+size 1003334200

requirements.txt

@@ -0,0 +1,9 @@
+torch>=2.4.0
+safetensors>=0.4.0
+tokenizers>=0.20.0
+huggingface_hub>=0.26.0
+einops>=0.8.0
+numpy>=1.26.0
+sympy>=1.13.0
+
+# flash-linear-attention is auto-installed by inference.py on first run.

sft_config.json

@@ -0,0 +1,38 @@
+{
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.95,
+  "adam_epsilon": 1e-08,
+  "cpu_prefetch_factor": 4,
+  "cpu_workers": 2,
+  "dataset_manifest": "data/tokenized/stage3-sft/manifest.json",
+  "ema_inv_gamma": 1.0,
+  "ema_max_decay": 0.9999,
+  "ema_power": 0.75,
+  "enable_gradient_checkpointing": true,
+  "full_shuffle": true,
+  "gradient_accumulation_steps": 16,
+  "gradient_clip_norm": 1.0,
+  "initial_run_dir": "runs/tercet-r-instruct",
+  "initial_weights": "ema",
+  "learning_rate": 3e-05,
+  "log_every_optimizer_steps": 1,
+  "logit_z_loss_coefficient": 0.0001,
+  "maximum_checkpoints": 5,
+  "maximum_optimizer_steps": null,
+  "micro_batch_size": 1,
+  "minimum_learning_rate_ratio": 0.1,
+  "muon_learning_rate": 0.01,
+  "muon_momentum": 0.95,
+  "optimizer": "adamw",
+  "output_dir": "runs/tercet-r-stage3",
+  "require_complete_pretraining": false,
+  "require_complete_sft_source": false,
+  "save_every_optimizer_steps": 100,
+  "seed": 20260902,
+  "sequence_length": 16384,
+  "shuffle_block_sequences": 4096,
+  "validation_batches": 16,
+  "warmup_ratio": 0.01,
+  "warmup_steps": 200,
+  "weight_decay": 0.0
+}

special_token_ids.json

@@ -0,0 +1,15 @@
+{
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "im_start": "<|im_start|>",
+  "im_end": "<|im_end|>",
+  "think_control": "<|think|>",
+  "no_think_control": "<|no_think|>",
+  "tool_response": "<|tool_response|>",
+  "bos_token_id": 0,
+  "eos_token_id": 1,
+  "pad_token_id": 2,
+  "unk_token_id": 3
+}

special_tokens_map.json

@@ -0,0 +1,68 @@
+{
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|tool_call|>",
+    "<|tool_response|>",
+    "<think>",
+    "</think>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|fim_pad|>",
+    "<|no_think|>",
+    "<|think|>",
+    "<|reserved_002|>",
+    "<|reserved_003|>",
+    "<|reserved_004|>",
+    "<|reserved_005|>",
+    "<|reserved_006|>",
+    "<|reserved_007|>",
+    "<|reserved_008|>",
+    "<|reserved_009|>",
+    "<|reserved_010|>",
+    "<|reserved_011|>",
+    "<|reserved_012|>",
+    "<|reserved_013|>",
+    "<|reserved_014|>",
+    "<|reserved_015|>",
+    "<|reserved_016|>",
+    "<|reserved_017|>",
+    "<|reserved_018|>",
+    "<|reserved_019|>",
+    "<|reserved_020|>",
+    "<|reserved_021|>",
+    "<|reserved_022|>",
+    "<|reserved_023|>",
+    "<|reserved_024|>",
+    "<|reserved_025|>",
+    "<|reserved_026|>",
+    "<|reserved_027|>",
+    "<|reserved_028|>",
+    "<|reserved_029|>",
+    "<|reserved_030|>",
+    "<|reserved_031|>",
+    "<|reserved_032|>",
+    "<|reserved_033|>",
+    "<|reserved_034|>",
+    "<|reserved_035|>",
+    "<|reserved_036|>",
+    "<|reserved_037|>",
+    "<|reserved_038|>",
+    "<|reserved_039|>",
+    "<|reserved_040|>",
+    "<|reserved_041|>",
+    "<|reserved_042|>",
+    "<|reserved_043|>",
+    "<|reserved_044|>",
+    "<|reserved_045|>",
+    "<|reserved_046|>",
+    "<|reserved_047|>",
+    "<|reserved_048|>",
+    "<|reserved_049|>"
+  ]
+}

tiny_gdn/__init__.py

@@ -0,0 +1,8 @@
+from tiny_gdn.config import TinyGDNConfig
+from tiny_gdn.model import TinyGDNForCausalLM, TinyGDNOutput
+
+__all__ = [
+    "TinyGDNConfig",
+    "TinyGDNForCausalLM",
+    "TinyGDNOutput",
+]

tiny_gdn/chatml.py

@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+from typing import Any
+
+from tiny_gdn.smoltalk_chat import (
+    NO_THINK_PREFIX,
+    THINK_CONTROL_PREFIX,
+    kwargs_from_inference_request,
+    materialize_smoltalk_messages,
+    merge_kwargs_with_system_message,
+    resolve_reasoning_mode,
+)
+from tiny_gdn.tools import prepare_inference_messages
+
+THINK_START = "<think>"
+THINK_END = "</think>"
+
+
+def decode_chat_completion(tokenizer: Any, token_ids: list[int]) -> str:
+    return tokenizer.decode(token_ids, skip_special_tokens=False)
+
+
+def append_assistant_generation_prompt(
+    tokenizer: Any,
+    token_ids: list[int],
+    *,
+    enable_thinking: bool | None = None,
+) -> None:
+    im_start = tokenizer.token_to_id("<|im_start|>")
+    if im_start is None:
+        raise RuntimeError("Tokenizer is missing <|im_start|>")
+    token_ids.append(im_start)
+    token_ids.extend(
+        tokenizer.encode("assistant\n", add_special_tokens=False).ids
+    )
+    if enable_thinking is True:
+        token_ids.extend(
+            tokenizer.encode(THINK_CONTROL_PREFIX, add_special_tokens=False).ids
+        )
+    elif enable_thinking is False:
+        token_ids.extend(
+            tokenizer.encode(NO_THINK_PREFIX, add_special_tokens=False).ids
+        )
+
+
+def encode_chatml_messages(
+    tokenizer: Any,
+    messages: list[dict[str, Any]],
+    *,
+    tools: Any | None = None,
+    enable_thinking: bool | None = None,
+    xml_tools: Any | None = None,
+    python_tools: Any | None = None,
+) -> list[int]:
+    if enable_thinking is None:
+        prepared = prepare_inference_messages(messages, tools=tools)
+        assistant_thinking_control = None
+    else:
+        kwargs = kwargs_from_inference_request(
+            enable_thinking=enable_thinking,
+            xml_tools=xml_tools,
+            python_tools=python_tools,
+            tools=tools,
+        )
+        prepared = materialize_smoltalk_messages(messages, kwargs)
+        _, resolved = merge_kwargs_with_system_message(messages, kwargs)
+        assistant_thinking_control = (
+            resolve_reasoning_mode(
+                resolved.enable_thinking,
+                resolved.custom_instructions,
+            )
+            == "/think"
+        )
+    required_tokens = {
+        token: tokenizer.token_to_id(token)
+        for token in (
+            "<|begin_of_text|>",
+            "<|im_start|>",
+            "<|im_end|>",
+        )
+    }
+    if any(token_id is None for token_id in required_tokens.values()):
+        raise RuntimeError("Tokenizer is missing ChatML special tokens")
+    bos_id = required_tokens["<|begin_of_text|>"]
+    im_start_id = required_tokens["<|im_start|>"]
+    im_end_id = required_tokens["<|im_end|>"]
+    if bos_id is None or im_start_id is None or im_end_id is None:
+        raise RuntimeError("Tokenizer ChatML IDs could not be resolved")
+
+    token_ids = [bos_id]
+    newline_ids = tokenizer.encode("\n", add_special_tokens=False).ids
+    allowed_roles = {"system", "user", "assistant"}
+    for message_index, message in enumerate(prepared):
+        role = message["role"]
+        content = message["content"]
+        masked_prefix = message.get("masked_prefix", "")
+        if role not in allowed_roles:
+            raise ValueError(
+                f"Unsupported chat role at index {message_index}: {role!r}"
+            )
+        if not isinstance(masked_prefix, str):
+            raise ValueError(
+                f"Masked prefix at index {message_index} must be a string"
+            )
+        if not content.strip():
+            raise ValueError(
+                f"Chat content at index {message_index} must be non-empty"
+            )
+        token_ids.append(im_start_id)
+        token_ids.extend(
+            tokenizer.encode(f"{role}\n", add_special_tokens=False).ids
+        )
+        if masked_prefix:
+            token_ids.extend(
+                tokenizer.encode(masked_prefix, add_special_tokens=False).ids
+            )
+        token_ids.extend(
+            tokenizer.encode(content, add_special_tokens=False).ids
+        )
+        token_ids.append(im_end_id)
+        token_ids.extend(newline_ids)
+    if prepared[-1]["role"] != "user":
+        raise ValueError("The final chat message must have role 'user'")
+    append_assistant_generation_prompt(
+        tokenizer,
+        token_ids,
+        enable_thinking=assistant_thinking_control,
+    )
+    return token_ids

tiny_gdn/cli_chat.py

@@ -0,0 +1,484 @@
+"""Terminal chat: live reasoning stream plus SmolTalk tool-call rounds.
+
+Tercet-R assistant turns may contain:
+
+- a zero-loss `<|think|>` / `<|no_think|>` control prefix
+- a `<think>…</think>` reasoning block
+- one or more SmolTalk JSON `<tool_call>` blocks (NVIDIA XML is also parsed)
+
+This module colours those regions as tokens arrive and, after a completed
+turn, collects tool observations (prefixed with `<|tool_response|>`) so the
+model can continue the same conversation.
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal, TextIO
+
+from tiny_gdn.smoltalk_chat import (
+    TOOL_RESPONSE_TOKEN,
+    format_smoltalk_tool_call,
+    wrap_smoltalk_tool_result,
+)
+from tiny_gdn.code_exec import (
+    CALCULATOR_TOOL,
+    NEMOTRON_PYTHON_EXEC_TOOL,
+    execute_math_tool,
+    is_auto_math_tool_name,
+    reset_default_python_session,
+)
+from tiny_gdn.tools import ParsedToolCall, parse_tool_calls
+from tiny_gdn.web_search import (
+    NEMOTRON_WEB_SEARCH_TOOL,
+    is_web_search_tool_name,
+    query_from_arguments,
+    search_web,
+)
+
+
+SegmentKind = Literal["answer", "think", "tool_call"]
+MarkerKind = Literal[
+    "think_open",
+    "think_close",
+    "think_control",
+    "no_think",
+    "tool_open",
+    "tool_close",
+]
+
+THINK_OPEN = "<think>"
+THINK_CLOSE = "</think>"
+TOOL_CALL_OPEN = "<tool_call>"
+TOOL_CALL_CLOSE = "</tool_call>"
+THINK_CONTROL = "<|think|>"
+NO_THINK_CONTROL = "<|no_think|>"
+
+MARKERS: tuple[tuple[str, MarkerKind], ...] = (
+    (THINK_CLOSE, "think_close"),
+    (TOOL_CALL_CLOSE, "tool_close"),
+    (THINK_OPEN, "think_open"),
+    (TOOL_CALL_OPEN, "tool_open"),
+    (NO_THINK_CONTROL, "no_think"),
+    (THINK_CONTROL, "think_control"),
+)
+
+ANSI = {
+    "think": "\033[2;33m",
+    "tool_call": "\033[36m",
+    "answer": "\033[0m",
+    "reset": "\033[0m",
+}
+
+DEFAULT_MAX_TOOL_ROUNDS = 8
+
+
+@dataclass(frozen=True)
+class ContentSegment:
+    kind: SegmentKind
+    text: str
+    open: bool = False
+
+
+@dataclass(frozen=True)
+class GeneratedTurn:
+    text: str
+    token_count: int
+    stop_reason: str
+    tool_calls: tuple[ParsedToolCall, ...]
+
+
+def first_marker(text: str) -> tuple[int, str, MarkerKind] | None:
+    best: tuple[int, str, MarkerKind] | None = None
+    for marker, kind in MARKERS:
+        at = text.find(marker)
+        if at < 0:
+            continue
+        if (
+            best is None
+            or at < best[0]
+            or (at == best[0] and len(marker) > len(best[1]))
+        ):
+            best = (at, marker, kind)
+    return best
+
+
+def holdback_prefix_length(text: str) -> int:
+    if not text:
+        return 0
+    keep = 0
+    for marker, _kind in MARKERS:
+        limit = min(len(marker) - 1, len(text))
+        for size in range(1, limit + 1):
+            if marker.startswith(text[-size:]):
+                keep = max(keep, size)
+    return keep
+
+
+def mode_after_marker(kind: MarkerKind, current: SegmentKind) -> SegmentKind:
+    if kind in {"think_open", "think_control"}:
+        return "think"
+    if kind in {"think_close", "no_think", "tool_close"}:
+        return "answer"
+    if kind == "tool_open":
+        return "tool_call"
+    return current
+
+
+def _segment_text(segment: ContentSegment) -> str:
+    return segment.text.strip()
+
+
+def format_assistant_markdown(source: str) -> str:
+    """Render a streamed assistant turn as a single markdown block."""
+
+    think_parts: list[str] = []
+    answer_parts: list[str] = []
+    tool_parts: list[str] = []
+    for segment in split_assistant_segments(source):
+        text = _segment_text(segment)
+        if not text:
+            continue
+        if segment.kind == "think":
+            think_parts.append(text)
+        elif segment.kind == "tool_call":
+            tool_parts.append(text)
+        else:
+            answer_parts.append(text)
+    blocks: list[str] = []
+    if think_parts:
+        blocks.append("**Reasoning**\n\n" + "\n\n".join(think_parts))
+    if answer_parts:
+        blocks.append("\n\n".join(answer_parts))
+    for tool in tool_parts:
+        blocks.append(f"```tool_call\n{tool}\n```")
+    return "\n\n".join(blocks).strip()
+
+
+def format_assistant_chat_messages(source: str) -> list[dict[str, Any]]:
+    """Split a turn into Gradio thoughts plus a normal assistant chat message.
+
+    Messages with ``metadata.title`` render as collapsible thoughts. The reply
+    has no metadata so Gradio shows it as the chat bubble.
+    """
+
+    think_parts: list[str] = []
+    answer_parts: list[str] = []
+    tool_parts: list[str] = []
+    think_pending = False
+    for segment in split_assistant_segments(source):
+        text = _segment_text(segment)
+        if segment.kind == "think":
+            think_pending = segment.open
+            if text:
+                think_parts.append(text)
+            continue
+        think_pending = False
+        if not text:
+            continue
+        if segment.kind == "tool_call":
+            tool_parts.append(text)
+        else:
+            answer_parts.append(text)
+
+    messages: list[dict[str, Any]] = []
+    if think_parts:
+        messages.append(
+            {
+                "role": "assistant",
+                "content": "\n\n".join(think_parts),
+                "metadata": {
+                    "title": "Reasoning",
+                    "status": "pending" if think_pending else "done",
+                },
+            }
+        )
+    for tool in tool_parts:
+        messages.append(
+            {
+                "role": "assistant",
+                "content": f"```json\n{tool}\n```",
+                "metadata": {"title": "Tool call", "status": "done"},
+            }
+        )
+    answer = "\n\n".join(answer_parts)
+    if answer or not messages:
+        messages.append({"role": "assistant", "content": answer})
+    return messages
+
+
+def split_assistant_segments(source: str) -> list[ContentSegment]:
+    """Split a completed (or in-progress) assistant turn for tests / replay."""
+
+    segments: list[ContentSegment] = []
+    mode: SegmentKind = "answer"
+    cursor = 0
+    while cursor < len(source):
+        found = first_marker(source[cursor:])
+        if found is None:
+            tail = source[cursor:]
+            if tail:
+                segments.append(ContentSegment(kind=mode, text=tail, open=True))
+            break
+        at, marker, kind = found
+        at += cursor
+        if at > cursor:
+            segments.append(
+                ContentSegment(kind=mode, text=source[cursor:at], open=False)
+            )
+        mode = mode_after_marker(kind, mode)
+        cursor = at + len(marker)
+    return [segment for segment in segments if segment.text]
+
+
+class LiveReasoningStreamer:
+    """Colour reasoning and tool-call regions as decoded text grows."""
+
+    def __init__(
+        self,
+        writer: TextIO | None = None,
+        *,
+        color: bool | None = None,
+    ) -> None:
+        self.writer = writer if writer is not None else sys.stdout
+        if color is None:
+            color = bool(getattr(self.writer, "isatty", lambda: False)())
+        self.color = color
+        self._seen = ""
+        self._hold = ""
+        self._mode: SegmentKind = "answer"
+        self._style: SegmentKind | None = None
+        self._emitted_think_label = False
+        self._emitted_tool_label = False
+
+    def update(self, decoded: str) -> None:
+        if decoded.startswith(self._seen):
+            delta = decoded[len(self._seen) :]
+        else:
+            delta = decoded
+        self._seen = decoded
+        if delta:
+            self._consume(delta, final=False)
+
+    def finish(self) -> str:
+        if self._hold:
+            self._emit(self._hold)
+            self._hold = ""
+        self._set_style(None)
+        self.writer.write("\n")
+        self.writer.flush()
+        return self._seen
+
+    def _consume(self, delta: str, *, final: bool) -> None:
+        buffer = self._hold + delta
+        self._hold = ""
+        while buffer:
+            found = first_marker(buffer)
+            if found is None:
+                keep = 0 if final else holdback_prefix_length(buffer)
+                if keep:
+                    self._emit(buffer[:-keep])
+                    self._hold = buffer[-keep:]
+                else:
+                    self._emit(buffer)
+                return
+            at, marker, kind = found
+            if at:
+                self._emit(buffer[:at])
+            self._switch(kind)
+            buffer = buffer[at + len(marker) :]
+        if final:
+            return
+
+    def _switch(self, kind: MarkerKind) -> None:
+        nxt = mode_after_marker(kind, self._mode)
+        if nxt != self._mode and nxt == "answer":
+            self._emit_plain("\n")
+        self._mode = nxt
+        if nxt == "think" and not self._emitted_think_label:
+            self._emit_plain("\n")
+            self._set_style("think")
+            self._emit_plain("reasoning  ")
+            self._emitted_think_label = True
+        elif nxt == "tool_call" and not self._emitted_tool_label:
+            self._emit_plain("\n")
+            self._set_style("tool_call")
+            self._emit_plain("tool_call  ")
+            self._emitted_tool_label = True
+        elif nxt == "answer":
+            self._set_style("answer")
+
+    def _emit(self, text: str) -> None:
+        if not text:
+            return
+        self._set_style(self._mode)
+        self.writer.write(text)
+        self.writer.flush()
+
+    def _emit_plain(self, text: str) -> None:
+        if not text:
+            return
+        self._set_style(None)
+        self.writer.write(text)
+        self.writer.flush()
+
+    def _set_style(self, kind: SegmentKind | None) -> None:
+        if not self.color:
+            self._style = kind
+            return
+        if kind == self._style:
+            return
+        self.writer.write(ANSI["reset"])
+        if kind in {"think", "tool_call"}:
+            self.writer.write(ANSI[kind])
+        self._style = kind
+
+
+def prompt_tool_results(
+    calls: Sequence[ParsedToolCall],
+    *,
+    read_line: Callable[[str], str],
+    writer: TextIO | None = None,
+    execute_web_search: bool = True,
+) -> list[str]:
+    out = writer if writer is not None else sys.stdout
+    results: list[str] = []
+    auto_search = sum(
+        1
+        for call in calls
+        if execute_web_search and is_web_search_tool_name(call.name)
+    )
+    auto_math = sum(1 for call in calls if is_auto_math_tool_name(call.name))
+    if auto_search:
+        out.write(
+            f"\n{auto_search} web-search call(s) will run automatically "
+            f"(Tavily JSON, sent as {TOOL_RESPONSE_TOKEN}).\n"
+        )
+        out.flush()
+    if auto_math:
+        out.write(
+            f"\n{auto_math} python/calculator call(s) will run automatically "
+            f"(sent as {TOOL_RESPONSE_TOKEN}).\n"
+        )
+        out.flush()
+    manual = len(calls) - auto_search - auto_math
+    if manual:
+        out.write(
+            f"\n{manual} tool call(s). Paste each observation; "
+            f"it is sent as {TOOL_RESPONSE_TOKEN}.\n"
+        )
+        out.flush()
+    for index, call in enumerate(calls, start=1):
+        out.write(
+            f"\n[{index}/{len(calls)}] "
+            f"{format_smoltalk_tool_call(call.name, call.arguments)}\n"
+        )
+        out.flush()
+        if execute_web_search and is_web_search_tool_name(call.name):
+            query = query_from_arguments(call.arguments)
+            out.write(f"searching {query!r}…\n")
+            out.flush()
+            results.append(search_web(query))
+            continue
+        if is_auto_math_tool_name(call.name):
+            out.write("running python…\n")
+            out.flush()
+            try:
+                results.append(execute_math_tool(call.name, call.arguments))
+            except Exception as error:
+                results.append(f"Error: {error}")
+            continue
+        results.append(read_line(f"result[{call.name}]> "))
+    return results
+
+
+def append_tool_round(
+    messages: list[dict[str, Any]],
+    assistant_text: str,
+    raw_results: Sequence[str],
+) -> None:
+    messages.append({"role": "assistant", "content": assistant_text})
+    for raw in raw_results:
+        wrapped = wrap_smoltalk_tool_result(raw)
+        if not wrapped:
+            raise ValueError("Tool result cannot be empty")
+        messages.append({"role": "tool", "content": raw})
+
+
+def resolve_cli_tools(spec: str | None, tools_json: str | None) -> list[dict[str, Any]] | None:
+    tools: list[dict[str, Any]] = []
+    if spec:
+        for name in spec.split(","):
+            key = name.strip().lower()
+            if not key or key in {"none", "off"}:
+                continue
+            if is_web_search_tool_name(key):
+                tools.append(NEMOTRON_WEB_SEARCH_TOOL)
+                continue
+            if key in {"python", "python-exec", "code-interpreter"}:
+                tools.append(NEMOTRON_PYTHON_EXEC_TOOL)
+                continue
+            if key in {"calculator", "calc"}:
+                tools.append(CALCULATOR_TOOL)
+                continue
+            raise ValueError(
+                f"Unknown built-in tool {name!r}. "
+                "Use web-search, python, calculator, or --tools-json."
+            )
+    if tools_json:
+        payload = json.loads(tools_json)
+        if isinstance(payload, dict):
+            tools.append(payload)
+        elif isinstance(payload, list):
+            tools.extend(payload)
+        else:
+            raise ValueError("tools JSON must be an object or array")
+    return tools or None
+
+
+@dataclass
+class ChatLoopState:
+    messages: list[dict[str, Any]] = field(default_factory=list)
+    system: str = ""
+    enable_thinking: bool = True
+    tools: list[dict[str, Any]] | None = None
+
+    def reset(self) -> None:
+        self.messages = []
+        reset_default_python_session()
+        if self.system.strip():
+            self.messages.append({"role": "system", "content": self.system.strip()})
+
+    def add_user(self, text: str) -> None:
+        self.messages.append({"role": "user", "content": text})
+
+
+def apply_slash_command(state: ChatLoopState, text: str) -> str | None:
+    """Return a status string if `text` is a slash command, else None."""
+
+    command = text.strip()
+    lowered = command.lower()
+    if lowered in {"/exit", "/quit"}:
+        return "exit"
+    if lowered == "/reset":
+        state.reset()
+        return "history cleared"
+    if lowered == "/think":
+        if state.messages:
+            return "thinking can only be changed on a fresh conversation (/reset first)"
+        state.enable_thinking = True
+        return "thinking on — next assistant turn is prefixed with <|think|>"
+    if lowered in {"/no_think", "/nothink"}:
+        if state.messages:
+            return "thinking can only be changed on a fresh conversation (/reset first)"
+        state.enable_thinking = False
+        return "thinking off — next assistant turn is prefixed with <|no_think|>"
+    if lowered.startswith("/system"):
+        rest = command[len("/system") :].strip()
+        state.system = rest
+        state.reset()
+        return "system prompt updated" if rest else "system prompt cleared"
+    return None

tiny_gdn/code_exec.py

@@ -0,0 +1,497 @@
+"""Restricted stateful Python for Nemotron Cascade-2 math tool calls.
+
+Tercet-R saw `stateful_python_code_exec` on `math/math_tool.jsonl` during
+the Cascade-2 SFT stage. Observations are the Jupyter-style stdout / last
+expression value, not a JSON envelope.
+
+The worker process only allows a math-oriented import whitelist and rejects
+dunder access, so a tool call cannot read the filesystem or start a shell.
+"""
+
+from __future__ import annotations
+
+import ast
+import io
+import json
+import math
+import os
+import subprocess
+import sys
+import threading
+from collections.abc import Mapping
+from pathlib import Path
+from queue import Empty, Queue
+from typing import Any
+
+
+NEMOTRON_PYTHON_EXEC_TOOL: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "stateful_python_code_exec",
+        "description": (
+            "Call this function to execute Python code in a stateful Jupyter "
+            "notebook environment. Python will respond with the output of the "
+            "execution or time out after 120.0 seconds."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "code": {
+                    "type": "string",
+                    "description": "Code to execute",
+                }
+            },
+            "required": ["code"],
+        },
+    },
+}
+
+CALCULATOR_TOOL: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "calculator",
+        "description": "Evaluate a math expression and return the numeric result.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "expression": {
+                    "type": "string",
+                    "description": "Math expression, for example 12.5 * (3 + 4)",
+                }
+            },
+            "required": ["expression"],
+        },
+    },
+}
+
+PYTHON_EXEC_TOOL_NAMES = frozenset(
+    {
+        "stateful_python_code_exec",
+        "python_code_exec",
+        "code_interpreter",
+        "python_exec",
+        "python",
+    }
+)
+CALCULATOR_TOOL_NAMES = frozenset({"calculator", "calc"})
+DEFAULT_EXEC_TIMEOUT_SECONDS = 30.0
+_SRC_ROOT = Path(__file__).resolve().parents[1]
+
+_session_lock = threading.Lock()
+_default_session: StatefulPythonSession | None = None
+
+
+def _normalize_tool_name(name: str) -> str:
+    return name.strip().lower().replace("-", "_")
+
+
+def is_python_exec_tool_name(name: str) -> bool:
+    return _normalize_tool_name(name) in PYTHON_EXEC_TOOL_NAMES
+
+
+def is_calculator_tool_name(name: str) -> bool:
+    return _normalize_tool_name(name) in CALCULATOR_TOOL_NAMES
+
+
+def is_auto_math_tool_name(name: str) -> bool:
+    return is_python_exec_tool_name(name) or is_calculator_tool_name(name)
+
+
+def code_from_arguments(arguments: Mapping[str, Any] | None) -> str:
+    if not arguments:
+        return ""
+    for key in ("code", "expression", "expr", "source"):
+        raw = arguments.get(key)
+        if isinstance(raw, str) and raw.strip():
+            return raw
+        if raw is not None and key != "code":
+            return str(raw)
+    return ""
+
+
+class CodeExecError(ValueError):
+    """Rejected or failed tool code."""
+
+
+class StatefulPythonSession:
+    """One long-lived restricted interpreter, matching the training tool."""
+
+    def __init__(self, *, timeout_seconds: float = DEFAULT_EXEC_TIMEOUT_SECONDS) -> None:
+        self.timeout_seconds = timeout_seconds
+        self._lock = threading.Lock()
+        self._process: subprocess.Popen[str] | None = None
+
+    def close(self) -> None:
+        with self._lock:
+            self._kill_locked()
+
+    def reset(self) -> None:
+        with self._lock:
+            self._kill_locked()
+
+    def run(self, code: str) -> str:
+        text = code.strip()
+        if not text:
+            raise CodeExecError("code is empty")
+        with self._lock:
+            return self._run_locked(text)
+
+    def _run_locked(self, code: str) -> str:
+        process = self._ensure_process_locked()
+        try:
+            process.stdin.write(json.dumps({"code": code}, ensure_ascii=False) + "\n")
+            process.stdin.flush()
+            line = self._readline_locked(process, self.timeout_seconds)
+        except (BrokenPipeError, OSError) as error:
+            self._kill_locked()
+            raise CodeExecError(f"python worker died: {error}") from error
+        if not line:
+            stderr = ""
+            if process.stderr is not None:
+                try:
+                    stderr = process.stderr.read()
+                except OSError:
+                    stderr = ""
+            self._kill_locked()
+            detail = stderr.strip() or "python worker closed stdout"
+            raise CodeExecError(detail)
+        try:
+            payload = json.loads(line)
+        except json.JSONDecodeError as error:
+            self._kill_locked()
+            raise CodeExecError(f"python worker returned invalid JSON: {line!r}") from error
+        if not isinstance(payload, dict):
+            raise CodeExecError("python worker returned a non-object")
+        output = str(payload.get("output") or "")
+        if payload.get("ok") is True:
+            return output if output else "None"
+        raise CodeExecError(output or "execution failed")
+
+    def _readline_locked(self, process: subprocess.Popen[str], timeout: float) -> str:
+        if process.stdout is None:
+            raise CodeExecError("python worker has no stdout")
+        lines: Queue[str] = Queue()
+
+        def _read() -> None:
+            lines.put(process.stdout.readline() if process.stdout is not None else "")
+
+        reader = threading.Thread(target=_read, daemon=True)
+        reader.start()
+        reader.join(timeout)
+        if reader.is_alive():
+            self._kill_locked()
+            raise CodeExecError(f"timed out after {timeout:g}s")
+        try:
+            return lines.get_nowait()
+        except Empty:
+            return ""
+
+    def _ensure_process_locked(self) -> subprocess.Popen[str]:
+        process = self._process
+        if process is not None and process.poll() is None:
+            return process
+        env = os.environ.copy()
+        pythonpath = env.get("PYTHONPATH", "")
+        env["PYTHONPATH"] = (
+            str(_SRC_ROOT) if not pythonpath else f"{_SRC_ROOT}{os.pathsep}{pythonpath}"
+        )
+        self._process = subprocess.Popen(
+            [sys.executable, "-m", "tiny_gdn.code_exec"],
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1,
+            env=env,
+        )
+        if self._process.stdin is None or self._process.stdout is None:
+            self._kill_locked()
+            raise CodeExecError("failed to start python worker")
+        return self._process
+
+    def _kill_locked(self) -> None:
+        process = self._process
+        self._process = None
+        if process is None:
+            return
+        try:
+            process.kill()
+            process.wait(timeout=2)
+        except (OSError, subprocess.TimeoutExpired):
+            pass
+
+
+def default_python_session() -> StatefulPythonSession:
+    global _default_session
+    with _session_lock:
+        if _default_session is None:
+            _default_session = StatefulPythonSession()
+        return _default_session
+
+
+def reset_default_python_session() -> None:
+    global _default_session
+    with _session_lock:
+        session = _default_session
+        _default_session = None
+    if session is not None:
+        session.close()
+
+
+def execute_calculator(expression: str) -> str:
+    tree = ast.parse(expression, mode="eval")
+    _reject_unsafe_ast(tree)
+    value = eval(  # noqa: S307 — AST already rejected unsafe nodes
+        compile(tree, "<calculator>", "eval"),
+        {"__builtins__": {}},
+        _calculator_namespace(),
+    )
+    return _format_result(value)
+
+
+def execute_math_tool(name: str, arguments: Mapping[str, Any] | None) -> str:
+    source = code_from_arguments(arguments)
+    if is_calculator_tool_name(name) and "\n" not in source and ";" not in source:
+        try:
+            return execute_calculator(source)
+        except (SyntaxError, CodeExecError, TypeError, ValueError, ZeroDivisionError) as error:
+            raise CodeExecError(str(error)) from error
+    return default_python_session().run(source)
+
+
+def _calculator_namespace() -> dict[str, Any]:
+    names = {
+        key: getattr(math, key)
+        for key in dir(math)
+        if not key.startswith("_")
+    }
+    names.update(
+        {
+            "abs": abs,
+            "min": min,
+            "max": max,
+            "round": round,
+            "pow": pow,
+            "pi": math.pi,
+            "e": math.e,
+            "tau": math.tau,
+            "inf": math.inf,
+        }
+    )
+    return names
+
+
+def _reject_unsafe_ast(tree: ast.AST) -> None:
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.ClassDef, ast.AsyncFunctionDef)):
+            raise CodeExecError("class definitions are blocked")
+        if isinstance(node, ast.Attribute) and node.attr.startswith("_"):
+            raise CodeExecError("dunder attribute access is blocked")
+        if isinstance(node, ast.Name) and node.id.startswith("_"):
+            raise CodeExecError("dunder names are blocked")
+        if isinstance(node, ast.Call) and isinstance(node.func, ast.Name):
+            if node.func.id in _BANNED_CALLS:
+                raise CodeExecError(f"{node.func.id}() is blocked")
+
+
+_BANNED_CALLS = frozenset(
+    {
+        "eval",
+        "exec",
+        "compile",
+        "open",
+        "input",
+        "breakpoint",
+        "getattr",
+        "setattr",
+        "delattr",
+        "globals",
+        "locals",
+        "vars",
+        "dir",
+        "help",
+        "__import__",
+        "memoryview",
+        "exit",
+        "quit",
+    }
+)
+
+
+def _format_result(value: Any) -> str:
+    if value is None:
+        return "None"
+    if isinstance(value, bool):
+        return str(value)
+    if isinstance(value, float):
+        if value.is_integer() and abs(value) < 1e15:
+            return str(int(value))
+        return format(value, ".12g")
+    return str(value)
+
+
+_ALLOWED_IMPORT_ROOTS = frozenset(
+    {
+        "cmath",
+        "collections",
+        "copy",
+        "decimal",
+        "fractions",
+        "functools",
+        "itertools",
+        "json",
+        "math",
+        "mpmath",
+        "numbers",
+        "numpy",
+        "operator",
+        "re",
+        "statistics",
+        "string",
+        "sympy",
+        "textwrap",
+        "unicodedata",
+    }
+)
+
+_WORKER_NAMESPACE: dict[str, Any] | None = None
+
+
+def _allowed_import(
+    name: str,
+    globals: dict[str, Any] | None = None,
+    locals: dict[str, Any] | None = None,
+    fromlist: tuple[str, ...] = (),
+    level: int = 0,
+) -> Any:
+    root = name.split(".")[0]
+    if root not in _ALLOWED_IMPORT_ROOTS:
+        raise ImportError(f"import of {name!r} is blocked")
+    return __import__(name, globals, locals, fromlist, level)
+
+
+def _worker_namespace() -> dict[str, Any]:
+    global _WORKER_NAMESPACE
+    if _WORKER_NAMESPACE is None:
+        builtins = {
+            "abs": abs,
+            "all": all,
+            "any": any,
+            "bin": bin,
+            "bool": bool,
+            "bytes": bytes,
+            "chr": chr,
+            "complex": complex,
+            "dict": dict,
+            "divmod": divmod,
+            "enumerate": enumerate,
+            "filter": filter,
+            "float": float,
+            "format": format,
+            "frozenset": frozenset,
+            "hex": hex,
+            "int": int,
+            "isinstance": isinstance,
+            "issubclass": issubclass,
+            "iter": iter,
+            "len": len,
+            "list": list,
+            "map": map,
+            "max": max,
+            "min": min,
+            "next": next,
+            "oct": oct,
+            "ord": ord,
+            "pow": pow,
+            "print": print,
+            "range": range,
+            "repr": repr,
+            "reversed": reversed,
+            "round": round,
+            "set": set,
+            "slice": slice,
+            "sorted": sorted,
+            "str": str,
+            "sum": sum,
+            "tuple": tuple,
+            "zip": zip,
+            "True": True,
+            "False": False,
+            "None": None,
+            "__import__": _allowed_import,
+        }
+        _WORKER_NAMESPACE = {
+            "__builtins__": builtins,
+            "__name__": "__tool__",
+            "math": math,
+        }
+    return _WORKER_NAMESPACE
+
+
+def run_cell(code: str) -> str:
+    tree = ast.parse(code)
+    _reject_unsafe_import_roots(tree)
+    _reject_unsafe_ast(tree)
+    namespace = _worker_namespace()
+    buffer = io.StringIO()
+    previous = sys.stdout
+    sys.stdout = buffer
+    try:
+        if tree.body and isinstance(tree.body[-1], ast.Expr):
+            body = tree.body[:-1]
+            last = tree.body[-1]
+            if body:
+                exec(  # noqa: S102
+                    compile(ast.Module(body, type_ignores=[]), "<tool>", "exec"),
+                    namespace,
+                    namespace,
+                )
+            value = eval(  # noqa: S307
+                compile(ast.Expression(last.value), "<tool>", "eval"),
+                namespace,
+                namespace,
+            )
+            printed = buffer.getvalue()
+            if value is None:
+                return printed if printed else "None"
+            rendered = _format_result(value)
+            return f"{printed}{rendered}" if printed else rendered
+        exec(compile(tree, "<tool>", "exec"), namespace, namespace)  # noqa: S102
+        printed = buffer.getvalue()
+        return printed if printed else "None"
+    finally:
+        sys.stdout = previous
+
+
+def _reject_unsafe_import_roots(tree: ast.AST) -> None:
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                root = alias.name.split(".")[0]
+                if root not in _ALLOWED_IMPORT_ROOTS:
+                    raise CodeExecError(f"import of {alias.name!r} is blocked")
+        elif isinstance(node, ast.ImportFrom):
+            root = (node.module or "").split(".")[0]
+            if root not in _ALLOWED_IMPORT_ROOTS:
+                raise CodeExecError(f"import of {node.module!r} is blocked")
+
+
+def _worker_loop() -> None:
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            request = json.loads(line)
+            code = str(request.get("code") or "")
+            output = run_cell(code)
+            sys.stdout.write(json.dumps({"ok": True, "output": output}, ensure_ascii=False) + "\n")
+        except Exception as error:
+            sys.stdout.write(
+                json.dumps({"ok": False, "output": str(error)}, ensure_ascii=False) + "\n"
+            )
+        sys.stdout.flush()
+
+
+if __name__ == "__main__":
+    _worker_loop()

tiny_gdn/config.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+import json
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass(frozen=True)
+class TinyGDNConfig:
+    architecture: str = "TinyGDNForCausalLM"
+    model_type: str = "tiny_gdn"
+
+    vocab_size: int = 49_152
+    # Deep-thin sizing is deliberate: controlled sub-billion studies find
+    # depth materially more valuable than width around the 125M-150M scale.
+    hidden_size: int = 512
+    intermediate_size: int = 1_472
+    num_hidden_layers: int = 32
+
+    num_attention_heads: int = 4
+    num_key_value_heads: int = 1
+    attention_head_dim: int = 128
+    full_attention_interval: int = 4
+    attention_dropout: float = 0.0
+    partial_rotary_factor: float = 0.5
+    rope_theta: float = 1_000_000.0
+
+    linear_num_heads: int = 4
+    linear_num_value_heads: int = 4
+    linear_head_dim: int = 128
+    linear_expand_v: float = 1.0
+    linear_conv_kernel_dim: int = 4
+    allow_negative_eigenvalues: bool = False
+
+    max_position_embeddings: int = 32_768
+    training_sequence_length: int = 2_048
+    rms_norm_eps: float = 1e-6
+    initializer_range: float = 0.02
+    tie_word_embeddings: bool = True
+    shared_layer_indices: tuple[int, ...] = ()
+
+    # MTP is an opt-in ablation at this scale; static MTP is not assumed to
+    # improve a 150M model without a controlled pilot.
+    mtp_num_heads: int = 0
+    mtp_adapter_rank: int = 128
+    mtp_loss_weight: float = 0.0
+
+    bos_token_id: int = 0
+    eos_token_id: int = 1
+    pad_token_id: int = 2
+    unk_token_id: int = 3
+
+    def __post_init__(self) -> None:
+        if self.vocab_size <= 0 or self.vocab_size > 65_536:
+            raise ValueError("vocab_size must fit the uint16 token dataset")
+        if self.hidden_size != self.num_attention_heads * self.attention_head_dim:
+            raise ValueError("hidden_size must equal num_attention_heads * attention_head_dim")
+        if self.hidden_size != self.linear_num_heads * self.linear_head_dim:
+            raise ValueError("hidden_size must equal linear_num_heads * linear_head_dim")
+        if self.linear_num_value_heads < self.linear_num_heads:
+            raise ValueError("linear_num_value_heads must be at least linear_num_heads")
+        if self.linear_num_value_heads % self.linear_num_heads != 0:
+            raise ValueError("linear_num_value_heads must be divisible by linear_num_heads")
+        if self.num_attention_heads % self.num_key_value_heads != 0:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        if self.num_hidden_layers % self.full_attention_interval != 0:
+            raise ValueError("num_hidden_layers must be divisible by full_attention_interval")
+        if not 0.0 < self.partial_rotary_factor <= 1.0:
+            raise ValueError("partial_rotary_factor must be in (0, 1]")
+        rotary_dim = int(self.attention_head_dim * self.partial_rotary_factor)
+        if rotary_dim <= 0 or rotary_dim % 2:
+            raise ValueError("The partial rotary dimension must be positive and even")
+        if self.training_sequence_length > self.max_position_embeddings:
+            raise ValueError("training_sequence_length exceeds max_position_embeddings")
+        if len(set(self.shared_layer_indices)) != len(self.shared_layer_indices):
+            raise ValueError("shared_layer_indices must be unique")
+        if any(
+            index < 0 or index >= self.num_hidden_layers
+            for index in self.shared_layer_indices
+        ):
+            raise ValueError("shared_layer_indices contains an invalid layer")
+        if self.mtp_num_heads < 0:
+            raise ValueError("mtp_num_heads cannot be negative")
+        if self.mtp_num_heads and self.mtp_adapter_rank <= 0:
+            raise ValueError("mtp_adapter_rank must be positive when MTP is enabled")
+        if not 0.0 <= self.mtp_loss_weight <= 1.0:
+            raise ValueError("mtp_loss_weight must be between zero and one")
+        for token_id in (
+            self.bos_token_id,
+            self.eos_token_id,
+            self.pad_token_id,
+            self.unk_token_id,
+        ):
+            if not 0 <= token_id < self.vocab_size:
+                raise ValueError(f"Special token ID {token_id} is outside the vocabulary")
+
+    @property
+    def layer_types(self) -> tuple[str, ...]:
+        return tuple(
+            "full_attention" if (index + 1) % self.full_attention_interval == 0 else "gdn2"
+            for index in range(self.num_hidden_layers)
+        )
+
+    @property
+    def rotary_dim(self) -> int:
+        return int(self.attention_head_dim * self.partial_rotary_factor)
+
+    @property
+    def effective_num_layers(self) -> int:
+        return self.num_hidden_layers + len(self.shared_layer_indices)
+
+    def to_dict(self) -> dict[str, Any]:
+        payload = asdict(self)
+        payload["layer_types"] = list(self.layer_types)
+        return payload
+
+    def save_json(self, path: Path) -> None:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(
+            json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n",
+            encoding="utf-8",
+        )
+
+    @classmethod
+    def from_json(cls, path: Path) -> TinyGDNConfig:
+        payload = json.loads(path.read_text(encoding="utf-8"))
+        payload.pop("layer_types", None)
+        if "shared_layer_indices" in payload:
+            payload["shared_layer_indices"] = tuple(payload["shared_layer_indices"])
+        return cls(**payload)

tiny_gdn/detokenize.py

@@ -0,0 +1,131 @@
+"""Byte-accurate incremental decode for GPT-2-style byte-level BPE.
+
+`tokenizer.decode(ids)` UTF-8-replaces incomplete sequences. Streaming that
+as a string prefix then breaks on the next byte: decode([20492]) is a
+replacement character and decode([20492, 294]) is `` ≈``, so the
+accumulated text is no longer a prefix and the UI re-emits the whole
+think block.
+
+This decoder concatenates raw BPE bytes and only emits complete UTF-8
+characters. The completing byte of ``≈`` yields `` ≈``; nothing is rewound.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from typing import Any
+
+
+def gpt2_bytes_to_unicode() -> dict[int, str]:
+    raw = (
+        list(range(ord("!"), ord("~") + 1))
+        + list(range(ord("¡"), ord("¬") + 1))
+        + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    mapped = raw[:]
+    extra = 0
+    for byte in range(256):
+        if byte not in raw:
+            raw.append(byte)
+            mapped.append(256 + extra)
+            extra += 1
+    return dict(zip(raw, [chr(code) for code in mapped], strict=True))
+
+
+def unicode_to_gpt2_bytes() -> dict[str, int]:
+    return {char: byte for byte, char in gpt2_bytes_to_unicode().items()}
+
+
+def utf8_char_length(lead: int) -> int | None:
+    if lead < 0x80:
+        return 1
+    if 0xC2 <= lead <= 0xDF:
+        return 2
+    if 0xE0 <= lead <= 0xEF:
+        return 3
+    if 0xF0 <= lead <= 0xF4:
+        return 4
+    return None
+
+
+def split_complete_utf8(buffer: bytes) -> tuple[str, bytes]:
+    """Return (decoded prefix, leftover incomplete or invalid-lead bytes)."""
+    index = 0
+    end = len(buffer)
+    while index < end:
+        size = utf8_char_length(buffer[index])
+        if size is None:
+            break
+        if index + size > end:
+            break
+        chunk = buffer[index : index + size]
+        try:
+            chunk.decode("utf-8")
+        except UnicodeDecodeError:
+            break
+        index += size
+    if index == 0:
+        return "", buffer
+    return buffer[:index].decode("utf-8"), buffer[index:]
+
+
+def piece_to_bytes(piece: str, unicode_to_byte: Mapping[str, int]) -> bytes:
+    raw = bytearray()
+    for char in piece:
+        mapped = unicode_to_byte.get(char)
+        if mapped is None:
+            raw.extend(char.encode("utf-8"))
+        else:
+            raw.append(mapped)
+    return bytes(raw)
+
+
+class IncrementalUtf8Decoder:
+    """Push token ids; read only complete Unicode as it becomes available."""
+
+    def __init__(
+        self,
+        tokenizer: Any,
+        *,
+        skip_special_ids: Iterable[int] | None = None,
+    ) -> None:
+        vocab = tokenizer.get_vocab()
+        self._id_to_piece = {tid: piece for piece, tid in vocab.items()}
+        self._unicode_to_byte = unicode_to_gpt2_bytes()
+        self._skip_special_ids = set(skip_special_ids or ())
+        self._pending = b""
+        self._text = ""
+
+    @property
+    def text(self) -> str:
+        return self._text
+
+    def reset(self) -> None:
+        self._pending = b""
+        self._text = ""
+
+    def token_bytes(self, token_id: int) -> bytes:
+        piece = self._id_to_piece.get(token_id, "")
+        return piece_to_bytes(piece, self._unicode_to_byte)
+
+    def push(self, token_id: int) -> str:
+        if token_id in self._skip_special_ids:
+            return ""
+        self._pending += self.token_bytes(token_id)
+        complete, self._pending = split_complete_utf8(self._pending)
+        self._text += complete
+        return complete
+
+    def push_many(self, token_ids: Sequence[int]) -> str:
+        delta = []
+        for token_id in token_ids:
+            piece = self.push(token_id)
+            if piece:
+                delta.append(piece)
+        return "".join(delta)
+
+    def finalize(self) -> str:
+        if self._pending:
+            self._text += self._pending.decode("utf-8", errors="replace")
+            self._pending = b""
+        return self._text

tiny_gdn/model.py

@@ -0,0 +1,683 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from safetensors.torch import load_model, save_model
+from torch import nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from torch.utils.checkpoint import checkpoint
+
+from tiny_gdn.config import TinyGDNConfig
+
+try:
+    # Import the module directly — `from fla.layers import GatedDeltaNet2`
+    # executes layers/__init__.py and eagerly loads every attention kernel.
+    from fla.layers.gdn2 import GatedDeltaNet2
+except ImportError as import_error:
+    GatedDeltaNet2 = None
+    FLA_IMPORT_ERROR: ImportError | None = import_error
+else:
+    FLA_IMPORT_ERROR = None
+
+
+def attention_sdpa_backends(device: torch.device) -> list[SDPBackend]:
+    if device.type != "cuda":
+        return [SDPBackend.MATH]
+    # Blackwell (sm_120) prefers cuDNN SDPA. Flash is often missing on
+    # Windows/WSL builds; mem-efficient at 16k is the slow path. MATH is
+    # never listed on CUDA — a silent math fallback is a 10× step-time cliff.
+    return [
+        SDPBackend.CUDNN_ATTENTION,
+        SDPBackend.FLASH_ATTENTION,
+        SDPBackend.EFFICIENT_ATTENTION,
+    ]
+
+
+@dataclass
+class TinyGDNOutput:
+    loss: torch.Tensor | None
+    logits: torch.Tensor | None
+    main_loss: torch.Tensor | None
+    mtp_loss: torch.Tensor | None
+    z_loss: torch.Tensor | None
+    hidden_states: torch.Tensor | None = None
+    past_key_values: Any | None = None
+
+
+class RMSNorm(nn.Module):
+    """Zero-centered RMSNorm as used by Qwen3-Next."""
+
+    def __init__(self, hidden_size: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        normalized = hidden_states.float()
+        normalized = normalized * torch.rsqrt(normalized.square().mean(dim=-1, keepdim=True) + self.eps)
+        normalized = normalized * (1.0 + self.weight.float())
+        return normalized.to(dtype=input_dtype)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, rotary_dim: int, rope_theta: float) -> None:
+        super().__init__()
+        inverse_frequency = 1.0 / (
+            rope_theta
+            ** (
+                torch.arange(0, rotary_dim, 2, dtype=torch.float32)
+                / rotary_dim
+            )
+        )
+        self.rotary_dim = rotary_dim
+        self.register_buffer("inverse_frequency", inverse_frequency, persistent=False)
+
+    def forward(
+        self,
+        sequence_length: int,
+        device: torch.device,
+        dtype: torch.dtype,
+        position_offset: int = 0,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        positions = torch.arange(
+            position_offset,
+            position_offset + sequence_length,
+            device=device,
+            dtype=torch.float32,
+        )
+        frequencies = torch.outer(positions, self.inverse_frequency.float())
+        embeddings = torch.cat((frequencies, frequencies), dim=-1)
+        return embeddings.cos().to(dtype=dtype), embeddings.sin().to(dtype=dtype)
+
+
+def rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
+    first, second = hidden_states.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def apply_rotary_embedding(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    cosine: torch.Tensor,
+    sine: torch.Tensor,
+    rotary_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    cosine = cosine[None, None, :, :]
+    sine = sine[None, None, :, :]
+    query_rotary, query_pass = query[..., :rotary_dim], query[..., rotary_dim:]
+    key_rotary, key_pass = key[..., :rotary_dim], key[..., rotary_dim:]
+    query_rotary = query_rotary * cosine + rotate_half(query_rotary) * sine
+    key_rotary = key_rotary * cosine + rotate_half(key_rotary) * sine
+    return (
+        torch.cat((query_rotary, query_pass), dim=-1),
+        torch.cat((key_rotary, key_pass), dim=-1),
+    )
+
+
+class GatedGroupedQueryAttention(nn.Module):
+    """QK-normalized, partially rotary GQA with a learned sigmoid output gate."""
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.attention_head_dim
+        self.rotary_dim = config.rotary_dim
+        self.dropout = config.attention_dropout
+
+        query_size = self.num_heads * self.head_dim
+        key_value_size = self.num_key_value_heads * self.head_dim
+        self.q_gate_proj = nn.Linear(config.hidden_size, query_size * 2, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, key_value_size, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, key_value_size, bias=False)
+        self.o_proj = nn.Linear(query_size, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.rotary = RotaryEmbedding(self.rotary_dim, config.rope_theta)
+
+    def _attention_mask(
+        self,
+        attention_mask: torch.Tensor | None,
+        sequence_length: int,
+        device: torch.device,
+    ) -> torch.Tensor | None:
+        if attention_mask is None:
+            return None
+        if attention_mask.ndim != 2:
+            raise ValueError("attention_mask must have shape [batch, sequence]")
+        if attention_mask.shape[1] != sequence_length:
+            raise ValueError("attention_mask sequence length does not match input")
+
+        causal = torch.ones(
+            sequence_length,
+            sequence_length,
+            dtype=torch.bool,
+            device=device,
+        ).tril()
+        valid_keys = attention_mask[:, None, None, :].to(dtype=torch.bool, device=device)
+        return causal[None, None, :, :] & valid_keys
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        batch_size, sequence_length, _ = hidden_states.shape
+        past_len = 0 if past_key_value is None else past_key_value[0].shape[2]
+        query_and_gate = self.q_gate_proj(hidden_states)
+        query, output_gate = query_and_gate.chunk(2, dim=-1)
+
+        query = query.view(batch_size, sequence_length, self.num_heads, self.head_dim)
+        key = self.k_proj(hidden_states).view(
+            batch_size,
+            sequence_length,
+            self.num_key_value_heads,
+            self.head_dim,
+        )
+        value = self.v_proj(hidden_states).view(
+            batch_size,
+            sequence_length,
+            self.num_key_value_heads,
+            self.head_dim,
+        )
+
+        query = self.q_norm(query).transpose(1, 2)
+        key = self.k_norm(key).transpose(1, 2)
+        value = value.transpose(1, 2)
+
+        cosine, sine = self.rotary(
+            sequence_length,
+            device=hidden_states.device,
+            dtype=query.dtype,
+            position_offset=past_len,
+        )
+        query, key = apply_rotary_embedding(
+            query,
+            key,
+            cosine,
+            sine,
+            rotary_dim=self.rotary_dim,
+        )
+        if past_key_value is not None:
+            key = torch.cat([past_key_value[0], key], dim=2)
+            value = torch.cat([past_key_value[1], value], dim=2)
+        present = (key, value)
+
+        kv_len = key.shape[2]
+        if attention_mask is not None and past_key_value is None:
+            sdpa_mask = self._attention_mask(
+                attention_mask,
+                sequence_length,
+                hidden_states.device,
+            )
+            is_causal = False
+        elif past_key_value is not None and sequence_length == 1:
+            # Decode step: query attends to the cached key/value prefix. Preserve
+            # the prefill padding mask when decoding a left-padded prompt batch.
+            sdpa_mask = (
+                None
+                if attention_mask is None
+                else attention_mask[:, None, None, :].to(
+                    dtype=torch.bool,
+                    device=hidden_states.device,
+                )
+            )
+            is_causal = False
+        elif past_key_value is not None:
+            # Prefill chunk with cache — build causal mask over kv_len.
+            q_idx = torch.arange(
+                past_len, past_len + sequence_length, device=hidden_states.device
+            )[:, None]
+            k_idx = torch.arange(kv_len, device=hidden_states.device)[None, :]
+            sdpa_mask = (k_idx <= q_idx)[None, None, :, :]
+            is_causal = False
+        else:
+            sdpa_mask = None
+            is_causal = True
+        sdpa_options = {
+            "attn_mask": sdpa_mask,
+            "dropout_p": self.dropout if self.training else 0.0,
+            "is_causal": is_causal,
+            "enable_gqa": True,
+        }
+        with sdpa_kernel(attention_sdpa_backends(query.device)):
+            attention_output = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                **sdpa_options,
+            )
+        attention_output = attention_output.transpose(1, 2).reshape(
+            batch_size,
+            sequence_length,
+            -1,
+        )
+        attention_output = attention_output * torch.sigmoid(output_gate)
+        return self.o_proj(attention_output), present
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.gate_up_proj = nn.Linear(
+            config.hidden_size,
+            config.intermediate_size * 2,
+            bias=False,
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        gate, up = self.gate_up_proj(hidden_states).chunk(2, dim=-1)
+        return self.down_proj(F.silu(gate) * up)
+
+
+class TinyGDNBlock(nn.Module):
+    def __init__(self, config: TinyGDNConfig, layer_index: int) -> None:
+        super().__init__()
+        layer_type = config.layer_types[layer_index]
+        self.layer_type = layer_type
+        self.token_mixer_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mlp_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+        if layer_type == "gdn2":
+            if GatedDeltaNet2 is None:
+                raise ImportError(
+                    "Gated DeltaNet-2 requires the pinned flash-linear-attention dependency"
+                ) from FLA_IMPORT_ERROR
+            self.token_mixer = GatedDeltaNet2(
+                hidden_size=config.hidden_size,
+                expand_v=config.linear_expand_v,
+                head_dim=config.linear_head_dim,
+                num_heads=config.linear_num_heads,
+                num_v_heads=config.linear_num_value_heads,
+                mode="chunk",
+                use_short_conv=True,
+                allow_neg_eigval=config.allow_negative_eigenvalues,
+                conv_size=config.linear_conv_kernel_dim,
+                conv_bias=False,
+                layer_idx=layer_index,
+                norm_eps=config.rms_norm_eps,
+            )
+        elif layer_type == "full_attention":
+            self.token_mixer = GatedGroupedQueryAttention(config)
+        else:
+            raise ValueError(f"Unsupported layer type: {layer_type}")
+
+        self.mlp = SwiGLU(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        *,
+        past_key_values: Any | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, Any]:
+        residual = hidden_states
+        normalized = self.token_mixer_norm(hidden_states)
+        present: Any = None
+        if self.layer_type == "gdn2":
+            mixed, _, past_key_values = self.token_mixer(
+                normalized,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            present = past_key_values
+        else:
+            mixed, present = self.token_mixer(
+                normalized,
+                attention_mask=attention_mask,
+                past_key_value=past_key_value,
+            )
+            if not use_cache:
+                present = None
+        hidden_states = residual + mixed
+        hidden_states = hidden_states + self.mlp(self.mlp_norm(hidden_states))
+        return hidden_states, present
+
+
+class MultiTokenPredictionAdapter(nn.Module):
+    """A lightweight residual adapter for one additional prediction horizon."""
+
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.down_proj = nn.Linear(
+            config.hidden_size,
+            config.mtp_adapter_rank,
+            bias=False,
+        )
+        self.up_proj = nn.Linear(
+            config.mtp_adapter_rank,
+            config.hidden_size,
+            bias=False,
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        adapted = self.up_proj(F.silu(self.down_proj(self.norm(hidden_states))))
+        return hidden_states + adapted
+
+
+class TinyGDNForCausalLM(nn.Module):
+    def __init__(self, config: TinyGDNConfig) -> None:
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            TinyGDNBlock(config, layer_index)
+            for layer_index in range(config.num_hidden_layers)
+        )
+        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mtp_adapters = nn.ModuleList(
+            MultiTokenPredictionAdapter(config)
+            for _ in range(config.mtp_num_heads)
+        )
+        self.gradient_checkpointing = False
+
+        self.apply(self._initialize_module)
+        self._initialize_residual_projections()
+
+    def _initialize_module(self, module: nn.Module) -> None:
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(
+                module.weight,
+                mean=0.0,
+                std=self.config.initializer_range,
+            )
+
+    def _initialize_residual_projections(self) -> None:
+        residual_std = self.config.initializer_range / math.sqrt(
+            2 * self.config.num_hidden_layers
+        )
+        for layer in self.layers:
+            nn.init.normal_(
+                layer.token_mixer.o_proj.weight,
+                mean=0.0,
+                std=residual_std,
+            )
+            nn.init.normal_(
+                layer.mlp.down_proj.weight,
+                mean=0.0,
+                std=residual_std,
+            )
+        for adapter in self.mtp_adapters:
+            nn.init.normal_(adapter.up_proj.weight, mean=0.0, std=residual_std)
+
+    def enable_gradient_checkpointing(self, enabled: bool = True) -> None:
+        self.gradient_checkpointing = enabled
+
+    def project_to_vocabulary(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.linear(hidden_states, self.embed_tokens.weight)
+
+    def _run_layer(
+        self,
+        layer: TinyGDNBlock,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        *,
+        past_key_values: Any | None = None,
+        past_key_value: tuple[torch.Tensor, torch.Tensor] | None = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, Any]:
+        if self.gradient_checkpointing and self.training:
+            hidden_states, present = checkpoint(
+                layer,
+                hidden_states,
+                attention_mask,
+                use_reentrant=False,
+            )
+            return hidden_states, present
+        return layer(
+            hidden_states,
+            attention_mask,
+            past_key_values=past_key_values,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+        )
+
+    def _causal_loss(
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor,
+        target_offset: int,
+        adapter: nn.Module | None = None,
+        compute_z_loss: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if target_offset < 0:
+            raise ValueError("target_offset cannot be negative")
+        if target_offset and hidden_states.shape[1] <= target_offset:
+            raise ValueError(
+                f"Sequence length must exceed target offset {target_offset}"
+            )
+        if target_offset:
+            prediction_states = hidden_states[:, :-target_offset, :]
+            targets = labels[:, target_offset:].contiguous()
+        else:
+            prediction_states = hidden_states
+            targets = labels.contiguous()
+        if adapter is not None:
+            prediction_states = adapter(prediction_states)
+        logits = self.project_to_vocabulary(prediction_states)
+        cross_entropy = F.cross_entropy(
+            logits.reshape(-1, self.config.vocab_size),
+            targets.reshape(-1),
+            ignore_index=-100,
+        )
+        z_loss = None
+        if compute_z_loss:
+            valid_targets = targets.ne(-100)
+            # Keep logits in their training dtype for logsumexp. Casting the
+            # full [B, S, V] tensor to fp32 materializes ~3 GiB in the autograd
+            # graph at 16k; upcast only the reduced [B, S] partition.
+            log_partition = torch.logsumexp(logits, dim=-1).float()
+            z_loss = log_partition.square()[valid_targets].mean()
+        return cross_entropy, z_loss
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Any | None = None,
+        *,
+        use_cache: bool = False,
+        return_logits: bool = True,
+        return_hidden_states: bool = False,
+        labels_are_shifted: bool = False,
+        include_mtp_loss: bool = True,
+        mtp_loss_weight: float | None = None,
+        z_loss_coefficient: float = 0.0,
+        logits_to_keep: int | None = None,
+    ) -> TinyGDNOutput:
+        if input_ids.ndim != 2:
+            raise ValueError("input_ids must have shape [batch, sequence]")
+        if input_ids.shape[1] > self.config.max_position_embeddings:
+            raise ValueError("Input exceeds max_position_embeddings")
+        if labels is not None and labels.shape != input_ids.shape:
+            raise ValueError("labels must have the same shape as input_ids")
+        if z_loss_coefficient < 0.0:
+            raise ValueError("z_loss_coefficient cannot be negative")
+        if logits_to_keep is not None and logits_to_keep <= 0:
+            raise ValueError("logits_to_keep must be positive")
+        if use_cache and labels is not None:
+            raise ValueError("use_cache is not supported with labels")
+        effective_mtp_weight = (
+            self.config.mtp_loss_weight
+            if mtp_loss_weight is None
+            else mtp_loss_weight
+        )
+        if not 0.0 <= effective_mtp_weight <= 1.0:
+            raise ValueError("mtp_loss_weight must be between zero and one")
+
+        if use_cache and past_key_values is None:
+            try:
+                from fla.models.utils import Cache as FlaCache
+            except ImportError as import_error:
+                raise ImportError(
+                    "Cached decode requires flash-linear-attention Cache"
+                ) from import_error
+            past_key_values = {
+                "fla": FlaCache(),
+                "gqa": [None] * len(self.layers),
+            }
+        elif past_key_values is not None and not isinstance(past_key_values, dict):
+            raise TypeError("past_key_values must be a TinyGDN cache dict or None")
+
+        fla_cache = None if past_key_values is None else past_key_values["fla"]
+        gqa_cache = None if past_key_values is None else past_key_values["gqa"]
+
+        hidden_states = self.embed_tokens(input_ids)
+        shared_layer_indices = set(self.config.shared_layer_indices)
+        for layer_index, layer in enumerate(self.layers):
+            layer_gqa = None if gqa_cache is None else gqa_cache[layer_index]
+            hidden_states, present = self._run_layer(
+                layer,
+                hidden_states,
+                attention_mask,
+                past_key_values=fla_cache,
+                past_key_value=layer_gqa,
+                use_cache=use_cache,
+            )
+            if use_cache and layer.layer_type == "full_attention" and gqa_cache is not None:
+                gqa_cache[layer_index] = present
+            if layer_index in shared_layer_indices:
+                layer_gqa = None if gqa_cache is None else gqa_cache[layer_index]
+                hidden_states, present = self._run_layer(
+                    layer,
+                    hidden_states,
+                    attention_mask,
+                    past_key_values=fla_cache,
+                    past_key_value=layer_gqa,
+                    use_cache=use_cache,
+                )
+                if use_cache and layer.layer_type == "full_attention" and gqa_cache is not None:
+                    gqa_cache[layer_index] = present
+        hidden_states = self.final_norm(hidden_states)
+
+        main_loss = None
+        mtp_loss = None
+        z_loss = None
+        total_loss = None
+        if labels is not None:
+            main_target_offset = 0 if labels_are_shifted else 1
+            main_loss, z_loss = self._causal_loss(
+                hidden_states,
+                labels,
+                target_offset=main_target_offset,
+                compute_z_loss=z_loss_coefficient > 0.0,
+            )
+            if self.mtp_adapters and include_mtp_loss:
+                auxiliary_losses = [
+                    self._causal_loss(
+                        hidden_states,
+                        labels,
+                        target_offset=(
+                            head_index + 1
+                            if labels_are_shifted
+                            else head_index + 2
+                        ),
+                        adapter=adapter,
+                    )[0]
+                    for head_index, adapter in enumerate(self.mtp_adapters)
+                ]
+                mtp_loss = torch.stack(auxiliary_losses).mean()
+                total_loss = main_loss + effective_mtp_weight * mtp_loss
+            else:
+                total_loss = main_loss
+            if z_loss is not None:
+                total_loss = total_loss + z_loss_coefficient * z_loss
+
+        output_states = (
+            hidden_states
+            if logits_to_keep is None
+            else hidden_states[:, -logits_to_keep:, :]
+        )
+        logits = self.project_to_vocabulary(output_states) if return_logits else None
+        return TinyGDNOutput(
+            loss=total_loss,
+            logits=logits,
+            main_loss=main_loss,
+            mtp_loss=mtp_loss,
+            z_loss=z_loss,
+            hidden_states=hidden_states if return_hidden_states else None,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+    def parameter_report(self) -> dict[str, int]:
+        total = sum(parameter.numel() for parameter in self.parameters())
+        mtp = sum(parameter.numel() for parameter in self.mtp_adapters.parameters())
+        embeddings = self.embed_tokens.weight.numel()
+        return {
+            "deployable_core": total - mtp,
+            "training_total": total,
+            "embedding": embeddings,
+            "mtp_auxiliary": mtp,
+            "non_embedding_core": total - mtp - embeddings,
+        }
+
+    def save_checkpoint(self, output_dir: Path) -> None:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        self.config.save_json(output_dir / "config.json")
+        save_model(self, output_dir / "model.safetensors")
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        checkpoint_dir: Path,
+        *,
+        device: str | torch.device = "cpu",
+        dtype: torch.dtype | None = None,
+    ) -> TinyGDNForCausalLM:
+        config = TinyGDNConfig.from_json(checkpoint_dir / "config.json")
+        model = cls(config).to(device=device, dtype=dtype)
+        load_model(model, checkpoint_dir / "model.safetensors", device=str(device))
+        return model
+
+    def extra_repr(self) -> str:
+        report = self.parameter_report()
+        return (
+            f"core_parameters={report['deployable_core']:,}, "
+            f"training_parameters={report['training_total']:,}"
+        )
+
+    def get_architecture_metadata(self) -> dict[str, Any]:
+        return {
+            "architecture": self.config.architecture,
+            "layer_types": list(self.config.layer_types),
+            "effective_num_layers": self.config.effective_num_layers,
+            "shared_layer_indices": list(self.config.shared_layer_indices),
+            "parameter_report": self.parameter_report(),
+            "features": [
+                "32-layer deep-thin parameter allocation",
+                "Gated DeltaNet-2 recurrent memory",
+                "3:1 recurrent-to-full-attention hybrid",
+                "gated grouped-query attention",
+                "QK normalization",
+                "partial rotary embeddings",
+                "zero-centered RMSNorm",
+                "SwiGLU",
+                "tied input-output embeddings",
+                "optional multi-token prediction auxiliaries",
+            ],
+        }

tiny_gdn/smoltalk_chat.py

@@ -0,0 +1,545 @@
+"""SmolTalk2 / SmolLM3 chat-template contract for Tercet-R.
+
+The SFT subset stores the toggle and tools in `chat_template_kwargs`:
+
+- `enable_thinking`: True → assistant turns receive a zero-loss
+  `<|think|>\\n` control prefix; assistant targets may emit `<think>` naturally
+- `enable_thinking`: False → assistant turns receive a zero-loss
+  `<|no_think|>\\n` control prefix
+- `custom_instructions`: the system prompt, used verbatim. Only the SmolLM3
+  template directives `/think`, `/no_think`, and `/system_override` are
+  removed because they are control markers, not prompt text
+- `xml_tools` / `python_tools`: appended to the system prompt as a Hermes
+  `<tools>` block only when tools are present and the prompt does not already
+  carry one
+
+There is no default identity, metadata header, or wrapper: a conversation
+without a system prompt and without tools has no system turn at all.
+
+Tool results (`tool` / `function` roles, or text still wrapped in Nemotron
+`<tool_response>` XML) are rendered as a user turn where every result is
+prefixed with the `<|tool_response|>` special token.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, replace
+from typing import Any, Literal
+
+from tiny_gdn.tools import (
+    json_ready,
+    TOOL_RESPONSE_CLOSE,
+    TOOL_RESPONSE_OPEN,
+    TOOLS_CLOSE,
+    TOOLS_OPEN,
+    assistant_message_content,
+    coerce_tools,
+    message_text,
+    parse_arguments_payload,
+    strip_leaked_chatml,
+)
+
+
+ReasoningMode = Literal["/think", "/no_think"]
+
+NO_THINK_TOKEN = "<|no_think|>"
+NO_THINK_PREFIX = f"{NO_THINK_TOKEN}\n"
+THINK_CONTROL_TOKEN = "<|think|>"
+THINK_CONTROL_PREFIX = f"{THINK_CONTROL_TOKEN}\n"
+TOOL_RESPONSE_TOKEN = "<|tool_response|>"
+TOOL_RESPONSE_PREFIX = f"{TOOL_RESPONSE_TOKEN}\n"
+THINK_OPEN = "<think>"
+THINK_CLOSE = "</think>"
+TOOL_CALL_OPEN = "<tool_call>"
+TOOL_CALL_CLOSE = "</tool_call>"
+SYSTEM_OVERRIDE_MARK = "/system_override"
+THINK_MARK = "/think"
+NO_THINK_MARK = "/no_think"
+
+# SmolTalk2 SFT used this markdown web_search schema. Tercet-R stage-3
+# packs Nemotron Agentic-v2 search instead (`web-search` + Tavily JSON);
+# see tiny_gdn.web_search.NEMOTRON_WEB_SEARCH_TOOL.
+SMOLTALK_WEB_SEARCH_TOOL = {
+    "type": "function",
+    "function": {
+        "name": "web_search",
+        "description": (
+            "Performs a web search for a query and returns a string of the "
+            "top search results formatted as markdown with titles, links, "
+            "and descriptions."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "The search query to perform.",
+                }
+            },
+            "required": ["query"],
+        },
+    },
+}
+XML_TOOLS_PREAMBLE = (
+    "You may call one or more functions to assist with the user query.\n"
+    "You are provided with function signatures within <tools></tools> XML "
+    "tags:\n<tools>\n"
+)
+XML_TOOLS_EPILOGUE = (
+    "</tools>\n\n"
+    "For each function call, return a json object with function name and "
+    "arguments within <tool_call></tool_call> XML tags:\n"
+    "<tool_call>\n"
+    '{"name": <function-name>, "arguments": <args-json-object>}\n'
+    "</tool_call>"
+)
+PYTHON_TOOLS_PREAMBLE = (
+    "When you send a message containing Python code between '<code>' and "
+    "'</code>' tags, it will be executed in a stateful Jupyter notebook "
+    "environment, and you will then be given the output to continued "
+    "reasoning in an agentic loop.\n\n"
+    "You can use the following tools in your python code like regular "
+    "functions:\n<tools>\n"
+)
+PYTHON_TOOLS_EPILOGUE = (
+    "</tools>\n\n"
+    "The state persists between code executions: so variables that you "
+    "define in one step are still available thereafter."
+)
+
+
+@dataclass(frozen=True)
+class ChatTemplateKwargs:
+    custom_instructions: str = ""
+    enable_thinking: bool = True
+    xml_tools: tuple[str, ...] = ()
+    python_tools: tuple[str, ...] = ()
+
+
+def coerce_tool_strings(raw_tools: Any) -> tuple[str, ...]:
+    if raw_tools is None:
+        return ()
+    if isinstance(raw_tools, str):
+        text = raw_tools.strip()
+        return (text,) if text else ()
+    if isinstance(raw_tools, dict):
+        raw_tools = [raw_tools]
+    if not isinstance(raw_tools, list):
+        return ()
+    rendered: list[str] = []
+    for item in raw_tools:
+        if isinstance(item, str):
+            text = item.strip()
+            if text:
+                rendered.append(text)
+            continue
+        if isinstance(item, dict):
+            rendered.append(json.dumps(item, ensure_ascii=False))
+    return tuple(rendered)
+
+
+def parse_chat_template_kwargs(raw: Any) -> ChatTemplateKwargs:
+    if raw is None:
+        return ChatTemplateKwargs()
+    if isinstance(raw, str):
+        text = raw.strip()
+        if not text:
+            return ChatTemplateKwargs()
+        raw = json.loads(text)
+    if not isinstance(raw, dict):
+        return ChatTemplateKwargs()
+    enable = raw.get("enable_thinking")
+    if enable is None:
+        enable_thinking = True
+    elif isinstance(enable, bool):
+        enable_thinking = enable
+    else:
+        enable_thinking = str(enable).strip().lower() in {"1", "true", "yes", "on"}
+    custom = raw.get("custom_instructions")
+    return ChatTemplateKwargs(
+        custom_instructions=custom if isinstance(custom, str) else "",
+        enable_thinking=enable_thinking,
+        xml_tools=coerce_tool_strings(raw.get("xml_tools")),
+        python_tools=coerce_tool_strings(raw.get("python_tools")),
+    )
+
+
+def resolve_reasoning_mode(
+    enable_thinking: bool,
+    custom_instructions: str,
+) -> ReasoningMode:
+    if NO_THINK_MARK in custom_instructions:
+        return "/no_think"
+    if THINK_MARK in custom_instructions:
+        return "/think"
+    return "/think" if enable_thinking else "/no_think"
+
+
+def strip_template_markers(custom_instructions: str) -> str:
+    return (
+        custom_instructions.replace(NO_THINK_MARK, "")
+        .replace(THINK_MARK, "")
+        .replace(SYSTEM_OVERRIDE_MARK, "")
+        .strip()
+    )
+
+
+def strip_tools_wrapper(item: str) -> str:
+    """Drop a `<tools>` / `</tools>` wrapper some sources bake into each entry."""
+
+    text = item.strip()
+    if text.startswith(TOOLS_OPEN):
+        text = text[len(TOOLS_OPEN) :]
+    if text.endswith(TOOLS_CLOSE):
+        text = text[: -len(TOOLS_CLOSE)]
+    return text.strip()
+
+
+def render_xml_tools_block(xml_tools: tuple[str, ...]) -> str:
+    entries = [strip_tools_wrapper(item) for item in xml_tools]
+    entries = [item for item in entries if item]
+    if not entries:
+        return ""
+    body = "".join(f"{item}\n" for item in entries)
+    return f"{XML_TOOLS_PREAMBLE}{body}{XML_TOOLS_EPILOGUE}"
+
+
+def render_python_tools_block(python_tools: tuple[str, ...]) -> str:
+    entries = [strip_tools_wrapper(item) for item in python_tools]
+    entries = [item for item in entries if item]
+    if not entries:
+        return ""
+    body = "".join(f"{item}\n" for item in entries)
+    return f"{PYTHON_TOOLS_PREAMBLE}{body}{PYTHON_TOOLS_EPILOGUE}"
+
+
+def unwrap_tool_result(content: str) -> str:
+    """Return the bare tool output, without ChatML leaks or `<tool_response>` XML."""
+
+    text = strip_leaked_chatml(content).strip()
+    if text.startswith(TOOL_RESPONSE_TOKEN):
+        text = text[len(TOOL_RESPONSE_TOKEN) :].strip()
+    if TOOL_RESPONSE_OPEN in text:
+        chunks: list[str] = []
+        cursor = 0
+        while True:
+            start = text.find(TOOL_RESPONSE_OPEN, cursor)
+            if start < 0:
+                tail = text[cursor:].strip()
+                if tail:
+                    chunks.append(tail)
+                break
+            head = text[cursor:start].strip()
+            if head:
+                chunks.append(head)
+            start += len(TOOL_RESPONSE_OPEN)
+            end = text.find(TOOL_RESPONSE_CLOSE, start)
+            if end < 0:
+                inner = text[start:].strip()
+                cursor = len(text)
+            else:
+                inner = text[start:end].strip()
+                cursor = end + len(TOOL_RESPONSE_CLOSE)
+            if inner:
+                chunks.append(inner)
+            if cursor >= len(text):
+                break
+        text = "\n".join(chunks)
+    return text
+
+
+def wrap_smoltalk_tool_result(content: str) -> str:
+    text = unwrap_tool_result(content)
+    if not text:
+        return ""
+    return f"{TOOL_RESPONSE_PREFIX}{text}"
+
+
+def openai_tools_as_xml_strings(raw_tools: Any) -> tuple[str, ...]:
+    if raw_tools is None:
+        return ()
+    rendered: list[str] = []
+    for tool in coerce_tools(raw_tools):
+        rendered.append(
+            json.dumps(
+                json_ready({"type": "function", "function": tool}),
+                ensure_ascii=False,
+                default=str,
+            )
+        )
+    return tuple(rendered)
+
+
+def format_smoltalk_tool_call(name: str, arguments: dict[str, Any]) -> str:
+    payload = {"name": name, "arguments": json_ready(arguments)}
+    return (
+        f"{TOOL_CALL_OPEN}\n"
+        f"{json.dumps(payload, ensure_ascii=False, default=str)}\n"
+        f"{TOOL_CALL_CLOSE}"
+    )
+
+
+def format_smoltalk_tool_calls(raw_tool_calls: Any) -> str:
+    if not raw_tool_calls:
+        return ""
+    if not isinstance(raw_tool_calls, list):
+        raise ValueError("tool_calls must be a list")
+    chunks: list[str] = []
+    for raw_call in raw_tool_calls:
+        if not isinstance(raw_call, dict):
+            raise ValueError("Each tool_call must be an object")
+        payload = (
+            raw_call.get("function")
+            if isinstance(raw_call.get("function"), dict)
+            else raw_call
+        )
+        if not isinstance(payload, dict):
+            raise ValueError("tool_call is missing a function object")
+        name = payload.get("name")
+        if not isinstance(name, str) or not name.strip():
+            raise ValueError("tool_call is missing a function name")
+        arguments = parse_arguments_payload(payload.get("arguments"))
+        chunks.append(format_smoltalk_tool_call(name.strip(), arguments))
+    return "\n".join(chunks)
+
+
+def assistant_smoltalk_content(message: dict[str, Any]) -> str:
+    content = assistant_message_content(
+        {**message, "tool_calls": None}
+    )
+    if TOOL_CALL_OPEN in content:
+        return content
+    tool_xml = format_smoltalk_tool_calls(message.get("tool_calls"))
+    if not tool_xml:
+        return content
+    if content.strip():
+        return f"{content.rstrip()}\n{tool_xml}"
+    return tool_xml
+
+
+def render_system_prompt(kwargs: ChatTemplateKwargs) -> str:
+    """Render the system turn verbatim; empty string means no system turn."""
+
+    custom = strip_leaked_chatml(strip_template_markers(kwargs.custom_instructions))
+    tool_blocks: list[str] = []
+    if TOOLS_OPEN not in custom:
+        xml_block = render_xml_tools_block(kwargs.xml_tools)
+        if xml_block:
+            tool_blocks.append(xml_block)
+        python_block = render_python_tools_block(kwargs.python_tools)
+        if python_block:
+            tool_blocks.append(python_block)
+    sections = [custom] if custom else []
+    sections.extend(tool_blocks)
+    return "\n\n".join(sections).strip()
+
+
+def merge_kwargs_with_system_message(
+    messages: list[dict[str, Any]],
+    kwargs: ChatTemplateKwargs,
+) -> tuple[list[dict[str, Any]], ChatTemplateKwargs]:
+    if not messages:
+        return messages, kwargs
+    first = messages[0]
+    if first.get("role") != "system":
+        return messages, kwargs
+    system_text = message_text(first.get("content"))
+    if system_text.strip():
+        merged = (
+            system_text
+            if not kwargs.custom_instructions.strip()
+            else f"{kwargs.custom_instructions.rstrip()}\n\n{system_text.strip()}"
+        )
+        kwargs = replace(kwargs, custom_instructions=merged)
+    return messages[1:], kwargs
+
+
+def kwargs_from_inference_request(
+    *,
+    enable_thinking: bool,
+    custom_instructions: str = "",
+    xml_tools: Any = None,
+    python_tools: Any = None,
+    tools: Any = None,
+) -> ChatTemplateKwargs:
+    xml = coerce_tool_strings(xml_tools)
+    if not xml:
+        xml = openai_tools_as_xml_strings(tools)
+    return ChatTemplateKwargs(
+        custom_instructions=custom_instructions,
+        enable_thinking=enable_thinking,
+        xml_tools=xml,
+        python_tools=coerce_tool_strings(python_tools),
+    )
+
+
+def _normalize_role(raw_role: Any, index: int) -> str:
+    if not isinstance(raw_role, str):
+        raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
+    role = raw_role.strip().lower()
+    if role == "human":
+        return "user"
+    if role == "gpt":
+        return "assistant"
+    if role == "function":
+        return "tool"
+    if role not in {"system", "user", "assistant", "tool"}:
+        raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
+    return role
+
+
+def materialize_smoltalk_messages(
+    raw_messages: list[dict[str, Any]],
+    kwargs: ChatTemplateKwargs,
+) -> list[dict[str, str]]:
+    if not raw_messages:
+        raise ValueError("Chat history cannot be empty")
+    remaining, resolved = merge_kwargs_with_system_message(raw_messages, kwargs)
+    for index, raw_message in enumerate(remaining):
+        if not isinstance(raw_message, dict):
+            raise ValueError(f"Unsupported chat message at index {index}")
+        if _normalize_role(raw_message.get("role"), index) != "system":
+            continue
+        extra = strip_leaked_chatml(message_text(raw_message.get("content")))
+        if extra:
+            resolved = replace(
+                resolved,
+                custom_instructions=(
+                    f"{resolved.custom_instructions.rstrip()}\n\n{extra}".strip()
+                ),
+            )
+    mode = resolve_reasoning_mode(
+        resolved.enable_thinking,
+        resolved.custom_instructions,
+    )
+    masked_prefix = THINK_CONTROL_PREFIX if mode == "/think" else NO_THINK_PREFIX
+    prepared: list[dict[str, str]] = []
+    system_prompt = render_system_prompt(resolved)
+    if system_prompt:
+        prepared.append({"role": "system", "content": system_prompt})
+    pending_tool_results: list[str] = []
+    for index, raw_message in enumerate(remaining):
+        role = _normalize_role(raw_message.get("role"), index)
+        if role == "system":
+            continue
+        if role == "assistant":
+            content = assistant_smoltalk_content(raw_message)
+        else:
+            content = message_text(raw_message.get("content"))
+        if role == "tool":
+            content = wrap_smoltalk_tool_result(content)
+            if not content:
+                raise ValueError(f"Chat content at index {index} must be non-empty")
+            pending_tool_results.append(content)
+            continue
+        if pending_tool_results:
+            prepared.append(
+                {"role": "user", "content": "\n".join(pending_tool_results)}
+            )
+            pending_tool_results = []
+        if not content.strip():
+            raise ValueError(f"Chat content at index {index} must be non-empty")
+        prepared_message = {"role": role, "content": content}
+        if role == "assistant":
+            prepared_message["masked_prefix"] = masked_prefix
+        prepared.append(prepared_message)
+    if pending_tool_results:
+        prepared.append({"role": "user", "content": "\n".join(pending_tool_results)})
+    if not any(message["role"] != "system" for message in prepared):
+        raise ValueError("Chat history cannot be empty")
+    return prepared
+
+
+def format_think_block(reasoning: str, content: str) -> str:
+    """SmolTalk2 layout: `<think>\\n…\\n</think>\\n\\nanswer`."""
+
+    body = reasoning.strip()
+    if content:
+        return f"{THINK_OPEN}\n{body}\n{THINK_CLOSE}\n\n{content}"
+    return f"{THINK_OPEN}\n{body}\n{THINK_CLOSE}\n\n"
+
+
+def render_assistant_turn(
+    content: str,
+    *,
+    reasoning: str | None,
+    tool_calls: list[dict[str, Any]] | None,
+) -> dict[str, str]:
+    """Render one assistant turn with a per-turn thinking control prefix.
+
+    A turn with reasoning gets `<|think|>` and a `<think>` block; a turn
+    without reasoning gets `<|no_think|>`. Tool calls become SmolTalk JSON
+    `<tool_call>` blocks unless the content already inlines them.
+    """
+
+    text = content.strip()
+    if tool_calls and TOOL_CALL_OPEN not in text:
+        calls = "\n".join(
+            format_smoltalk_tool_call(call["name"], call["arguments"]) for call in tool_calls
+        )
+        text = f"{text}\n{calls}" if text else calls
+    if reasoning and reasoning.strip():
+        return {
+            "role": "assistant",
+            "content": format_think_block(reasoning, text),
+            "masked_prefix": THINK_CONTROL_PREFIX,
+        }
+    return {"role": "assistant", "content": text, "masked_prefix": NO_THINK_PREFIX}
+
+
+def materialize_canonical_conversation(
+    *,
+    system: str | None,
+    tools: Any,
+    messages: list[dict[str, Any]],
+) -> list[dict[str, str]]:
+    """Render a stage-3 canonical conversation into ChatML-ready turns.
+
+    `messages` entries are `{"role", "content", "reasoning"?, "tool_calls"?}`
+    with roles user / assistant / tool. The system prompt is used verbatim;
+    tools are appended only when the prompt lacks a `<tools>` block.
+    """
+
+    kwargs = ChatTemplateKwargs(
+        custom_instructions=system or "",
+        xml_tools=openai_tools_as_xml_strings(tools) if tools else (),
+    )
+    prepared: list[dict[str, str]] = []
+    system_prompt = render_system_prompt(kwargs)
+    if system_prompt:
+        prepared.append({"role": "system", "content": system_prompt})
+    pending_tool_results: list[str] = []
+    for index, message in enumerate(messages):
+        role = message["role"]
+        if role == "tool":
+            wrapped = wrap_smoltalk_tool_result(message_text(message.get("content")))
+            if not wrapped:
+                raise ValueError(f"Tool result at index {index} is empty")
+            pending_tool_results.append(wrapped)
+            continue
+        if pending_tool_results:
+            prepared.append({"role": "user", "content": "\n".join(pending_tool_results)})
+            pending_tool_results = []
+        if role == "user":
+            content = message_text(message.get("content"))
+            if not content.strip():
+                raise ValueError(f"User turn at index {index} is empty")
+            prepared.append({"role": "user", "content": content})
+            continue
+        if role == "assistant":
+            rendered = render_assistant_turn(
+                message_text(message.get("content")),
+                reasoning=message.get("reasoning"),
+                tool_calls=message.get("tool_calls"),
+            )
+            if not rendered["content"].strip():
+                raise ValueError(f"Assistant turn at index {index} is empty")
+            prepared.append(rendered)
+            continue
+        raise ValueError(f"Unsupported role at index {index}: {role!r}")
+    if pending_tool_results:
+        raise ValueError("Conversation ends with an unanswered tool result")
+    if not prepared or prepared[-1]["role"] != "assistant":
+        raise ValueError("Conversation must end with an assistant turn")
+    return prepared

tiny_gdn/tools.py

@@ -0,0 +1,579 @@
+"""Nemotron Cascade-2 tool protocol (XML in ChatML content).
+
+Tercet-R is trained on `nvidia/Nemotron-Cascade-2-SFT-Data`, which inlines
+available tools, calls, and results as `<tools>` / `<tool_call>` /
+`<tool_response>` in message text. That is not OpenAI `tool_calls` JSON and
+not this tokenizer's unused `<|tool_call|>` / `<|tool_response|>` specials.
+
+This module matches NVIDIA's Cascade-2 chat template:
+https://huggingface.co/nvidia/Nemotron-Cascade-2-30B-A3B/blob/main/chat_template.jinja
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import sys
+from dataclasses import dataclass
+from typing import Any
+
+
+TOOL_CALL_OPEN = "<tool_call>"
+TOOL_CALL_CLOSE = "</tool_call>"
+TOOL_RESPONSE_OPEN = "<tool_response>"
+TOOL_RESPONSE_CLOSE = "</tool_response>"
+TOOLS_OPEN = "<tools>"
+TOOLS_CLOSE = "</tools>"
+THINK_OPEN = "<think>"
+THINK_CLOSE = "</think>"
+
+CHATML_START_RE = re.compile(r"^<\|im_start\|>[A-Za-z]+\n")
+CHATML_END_RE = re.compile(r"\n?<\|im_end\|>\s*$")
+TOOL_CALL_BLOCK_RE = re.compile(
+    rf"{re.escape(TOOL_CALL_OPEN)}(.*?){re.escape(TOOL_CALL_CLOSE)}",
+    re.DOTALL,
+)
+FUNCTION_BLOCK_RE = re.compile(
+    r"<function=([^>\s]+)>(.*?)</function>",
+    re.DOTALL,
+)
+PARAMETER_BLOCK_RE = re.compile(
+    r"<parameter=([^>\s]+)>\n?(.*?)\n?</parameter>",
+    re.DOTALL,
+)
+FUNCTION_NAME_RE = re.compile(r"<function=([^>\s]+)")
+
+INFERENCE_ROLE_MAP = {
+    "system": "system",
+    "user": "user",
+    "human": "user",
+    "assistant": "assistant",
+    "gpt": "assistant",
+    "tool": "tool",
+    "function": "tool",
+}
+
+TOOLS_PREAMBLE = "# Tools\n\nYou have access to the following functions:\n\n"
+TOOL_CALL_INSTRUCTIONS = (
+    "\n\nIf you choose to call a function ONLY reply in the following format "
+    "with NO suffix:\n\n"
+    "<tool_call>\n"
+    "<function=example_function_name>\n"
+    "<parameter=example_parameter_1>\n"
+    "value_1\n"
+    "</parameter>\n"
+    "<parameter=example_parameter_2>\n"
+    "This is the value for the second parameter\n"
+    "that can span\n"
+    "multiple lines\n"
+    "</parameter>\n"
+    "</function>\n"
+    "</tool_call>\n\n"
+    "<IMPORTANT>\n"
+    "Reminder:\n"
+    "- Function calls MUST follow the specified format: an inner "
+    "<function=...></function> block must be nested within "
+    "<tool_call></tool_call> XML tags\n"
+    "- Required parameters MUST be specified\n"
+    "- You may provide optional reasoning for your function call in natural "
+    "language BEFORE the function call, but NOT after\n"
+    "- If there is no function call available, answer the question like "
+    "normal with your current knowledge and do not tell the user about "
+    "function calls\n"
+    "</IMPORTANT>"
+)
+
+
+@dataclass(frozen=True)
+class ParsedToolCall:
+    name: str
+    arguments: dict[str, Any]
+    raw: str
+
+
+def message_text(raw_content: Any) -> str:
+    if raw_content is None:
+        return ""
+    if isinstance(raw_content, str):
+        return raw_content
+    if isinstance(raw_content, list):
+        parts: list[str] = []
+        for item in raw_content:
+            if isinstance(item, str):
+                parts.append(item)
+                continue
+            if not isinstance(item, dict):
+                continue
+            part_type = item.get("type")
+            if part_type in {None, "text", "input_text", "output_text"}:
+                text = item.get("text")
+                if isinstance(text, str):
+                    parts.append(text)
+        return "".join(parts)
+    return str(raw_content)
+
+
+def strip_leaked_chatml(content: str) -> str:
+    text = content.strip()
+    while True:
+        match = CHATML_START_RE.match(text)
+        if match is None:
+            break
+        text = text[match.end() :]
+    text = CHATML_END_RE.sub("", text)
+    return text.strip()
+
+
+def has_tools_block(content: str) -> bool:
+    return TOOLS_OPEN in content
+
+
+def wrap_tool_response(content: str) -> str:
+    text = strip_leaked_chatml(content)
+    if TOOL_RESPONSE_OPEN in text:
+        return text
+    return f"{TOOL_RESPONSE_OPEN}\n\n\n{text}\n{TOOL_RESPONSE_CLOSE}"
+
+
+def _xml_value(value: Any) -> str:
+    if isinstance(value, dict) or (
+        isinstance(value, (list, tuple)) and not isinstance(value, (str, bytes))
+    ):
+        return json.dumps(value, ensure_ascii=False)
+    if value is True or value is False or value is None:
+        return str(value)
+    return str(value)
+
+
+def _render_extra_keys(payload: dict[str, Any], handled: set[str]) -> str:
+    chunks: list[str] = []
+    for key, value in payload.items():
+        if key in handled:
+            continue
+        chunks.append(f"\n<{key}>{_xml_value(value)}</{key}>")
+    return "".join(chunks)
+
+
+def _unwrap_tool(raw_tool: Any) -> dict[str, Any]:
+    if not isinstance(raw_tool, dict):
+        raise ValueError("Each tool must be an object")
+    if isinstance(raw_tool.get("function"), dict):
+        tool = dict(raw_tool["function"])
+    else:
+        tool = dict(raw_tool)
+    name = tool.get("name")
+    if not isinstance(name, str) or not name.strip():
+        raise ValueError("Tool is missing a function name")
+    tool["name"] = name.strip()
+    return tool
+
+
+def coerce_tools(raw_tools: Any) -> list[dict[str, Any]]:
+    if raw_tools is None:
+        return []
+    if isinstance(raw_tools, str):
+        text = raw_tools.strip()
+        if not text:
+            return []
+        raw_tools = loads_json(text)
+    if isinstance(raw_tools, dict):
+        raw_tools = [raw_tools]
+    if not isinstance(raw_tools, list):
+        raise ValueError("tools must be a list of function specs")
+    return [_unwrap_tool(item) for item in raw_tools]
+
+
+def render_function_schema(tool: dict[str, Any]) -> str:
+    chunks = [f"\n<function>\n<name>{tool['name']}</name>"]
+    description = tool.get("description")
+    if isinstance(description, str) and description.strip():
+        chunks.append(f"\n<description>{description.strip()}</description>")
+    chunks.append("\n<parameters>")
+    parameters = tool.get("parameters")
+    properties: dict[str, Any] = {}
+    if isinstance(parameters, dict):
+        raw_properties = parameters.get("properties")
+        if isinstance(raw_properties, dict):
+            properties = raw_properties
+        for param_name, raw_fields in properties.items():
+            fields = raw_fields if isinstance(raw_fields, dict) else {}
+            chunks.append("\n<parameter>")
+            chunks.append(f"\n<name>{param_name}</name>")
+            if "type" in fields:
+                chunks.append(f"\n<type>{_xml_value(fields['type'])}</type>")
+            if isinstance(fields.get("description"), str) and fields["description"].strip():
+                chunks.append(
+                    f"\n<description>{fields['description'].strip()}</description>"
+                )
+            if "enum" in fields:
+                chunks.append(f"\n<enum>{_xml_value(fields['enum'])}</enum>")
+            chunks.append(
+                _render_extra_keys(
+                    fields,
+                    {"name", "type", "description", "enum"},
+                )
+            )
+            chunks.append("\n</parameter>")
+        chunks.append(
+            _render_extra_keys(parameters, {"type", "properties", "required"})
+        )
+        if "required" in parameters:
+            chunks.append(f"\n<required>{_xml_value(parameters['required'])}</required>")
+    chunks.append("\n</parameters>")
+    chunks.append(
+        _render_extra_keys(
+            tool,
+            {"type", "name", "description", "parameters"},
+        )
+    )
+    chunks.append("\n</function>")
+    return "".join(chunks)
+
+
+def render_available_tools(raw_tools: Any) -> str:
+    tools = coerce_tools(raw_tools)
+    if not tools:
+        return ""
+    body = "".join(render_function_schema(tool) for tool in tools)
+    return (
+        f"{TOOLS_PREAMBLE}{TOOLS_OPEN}{body}\n{TOOLS_CLOSE}"
+        f"{TOOL_CALL_INSTRUCTIONS}"
+    )
+
+
+def inject_available_tools(system_content: str, raw_tools: Any) -> str:
+    block = render_available_tools(raw_tools)
+    if not block:
+        return system_content
+    if has_tools_block(system_content):
+        return system_content
+    if not system_content.strip():
+        return block
+    return f"{system_content.rstrip()}\n\n{block}"
+
+
+def parse_json_int(text: str) -> int | str:
+    """Keep oversized JSON integers as strings.
+
+    Python 3.12 refuses to convert integers longer than
+    ``sys.get_int_max_str_digits()`` (default 4300). Tool-call payloads
+    sometimes embed hashes or blobs as bare JSON numbers; those must not
+    abort packing.
+    """
+
+    digits = text.lstrip("+-")
+    limit = sys.get_int_max_str_digits()
+    if limit and len(digits) > limit:
+        return text
+    return int(text)
+
+
+def loads_json(text: str) -> Any:
+    return json.loads(text, parse_int=parse_json_int)
+
+
+# Models sometimes close a JSON string with `".}` instead of `"}`.
+_TRAILING_STRING_PERIOD_RE = re.compile(r'"\s*\.(?=\s*[}\],])')
+_TRAILING_COMMA_RE = re.compile(r",\s*(?=[}\]])")
+_JSON_NAME_RE = re.compile(r'"name"\s*:\s*"([^"]+)"')
+_JSON_QUERY_RE = re.compile(r'"query"\s*:\s*"([^"]*)"')
+
+
+def repair_jsonish(text: str) -> str:
+    repaired = _TRAILING_STRING_PERIOD_RE.sub('"', text.strip())
+    return _TRAILING_COMMA_RE.sub("", repaired)
+
+
+def loads_jsonish(text: str) -> Any:
+    candidates = (text.strip(), repair_jsonish(text), text.strip().replace("'", '"'))
+    seen: set[str] = set()
+    last_error: Exception | None = None
+    for candidate in candidates:
+        if not candidate or candidate in seen:
+            continue
+        seen.add(candidate)
+        try:
+            return loads_json(candidate)
+        except (json.JSONDecodeError, ValueError, RecursionError) as error:
+            last_error = error
+    if last_error is not None:
+        raise last_error
+    raise json.JSONDecodeError("Empty JSON", text, 0)
+
+
+def extract_json_tool_fields(text: str) -> tuple[str, dict[str, Any]] | None:
+    name_match = _JSON_NAME_RE.search(text)
+    if name_match is None:
+        return None
+    arguments: dict[str, Any] = {}
+    query_match = _JSON_QUERY_RE.search(text)
+    if query_match is not None:
+        arguments["query"] = query_match.group(1)
+    return name_match.group(1).strip(), arguments
+
+
+def _json_ready_int(value: int) -> int | str:
+    try:
+        json.dumps(value)
+    except ValueError:
+        previous = sys.get_int_max_str_digits()
+        sys.set_int_max_str_digits(0)
+        try:
+            return str(value)
+        finally:
+            sys.set_int_max_str_digits(previous)
+    return value
+
+
+def json_ready(value: Any) -> Any:
+    """Coerce Python literals into JSON-serialisable values.
+
+    Some SFT sources store tool arguments as Python literals. `ast.literal_eval`
+    turns `{1, 2}` into a `set` and `...` into `Ellipsis`, which later
+    `json.dumps` calls reject. Oversized ints are stored as decimal strings.
+    """
+
+    if value is None or isinstance(value, (bool, float, str)):
+        return value
+    if isinstance(value, int):
+        return _json_ready_int(value)
+    if value is Ellipsis:
+        return None
+    if isinstance(value, dict):
+        return {str(key): json_ready(item) for key, item in value.items()}
+    if isinstance(value, (set, frozenset)):
+        items = [json_ready(item) for item in value]
+        try:
+            return sorted(
+                items,
+                key=lambda item: json.dumps(item, sort_keys=True, default=str),
+            )
+        except TypeError:
+            return items
+    if isinstance(value, tuple):
+        return [json_ready(item) for item in value]
+    if isinstance(value, list):
+        return [json_ready(item) for item in value]
+    if isinstance(value, bytes):
+        try:
+            return value.decode("utf-8")
+        except UnicodeDecodeError:
+            return list(value)
+    try:
+        json.dumps(value)
+        return value
+    except (TypeError, ValueError):
+        return str(value)
+
+
+def parse_argument_value(raw: str) -> Any:
+    text = raw.strip()
+    if not text:
+        return ""
+    try:
+        return loads_json(text)
+    except (json.JSONDecodeError, ValueError, RecursionError):
+        return text
+
+
+def parse_arguments_payload(raw: Any) -> dict[str, Any]:
+    if raw is None:
+        return {}
+    ready = json_ready(raw)
+    if isinstance(ready, dict):
+        return ready
+    if isinstance(ready, str):
+        text = ready.strip()
+        if not text:
+            return {}
+        try:
+            loaded = loads_json(text)
+        except (json.JSONDecodeError, ValueError, RecursionError):
+            return {"value": ready}
+        loaded = json_ready(loaded)
+        if isinstance(loaded, dict):
+            return loaded
+        return {"value": loaded}
+    return {"value": ready}
+
+
+def format_tool_call_xml(name: str, arguments: dict[str, Any]) -> str:
+    chunks = [f"{TOOL_CALL_OPEN}\n<function={name}>\n"]
+    for key, value in arguments.items():
+        chunks.append(f"<parameter={key}>\n{_xml_value(value)}\n</parameter>\n")
+    chunks.append(f"</function>\n{TOOL_CALL_CLOSE}\n")
+    return "".join(chunks)
+
+
+def format_tool_calls_xml(raw_tool_calls: Any) -> str:
+    if not raw_tool_calls:
+        return ""
+    if not isinstance(raw_tool_calls, list):
+        raise ValueError("tool_calls must be a list")
+    chunks: list[str] = []
+    for raw_call in raw_tool_calls:
+        if not isinstance(raw_call, dict):
+            raise ValueError("Each tool_call must be an object")
+        payload = raw_call.get("function") if isinstance(raw_call.get("function"), dict) else raw_call
+        if not isinstance(payload, dict):
+            raise ValueError("tool_call is missing a function object")
+        name = payload.get("name")
+        if not isinstance(name, str) or not name.strip():
+            raise ValueError("tool_call is missing a function name")
+        arguments = parse_arguments_payload(payload.get("arguments"))
+        chunks.append(format_tool_call_xml(name.strip(), arguments))
+    return "".join(chunks)
+
+
+def parse_json_tool_inner(inner: str) -> ParsedToolCall | None:
+    text = inner.strip()
+    if not text:
+        return None
+    payload: Any | None
+    try:
+        payload = loads_jsonish(text)
+    except (json.JSONDecodeError, ValueError, RecursionError):
+        payload = None
+    if isinstance(payload, dict):
+        nested = payload.get("function")
+        source = nested if isinstance(nested, dict) else payload
+        name = source.get("name")
+        if not isinstance(name, str) or not name.strip():
+            name = payload.get("name")
+        if isinstance(name, str) and name.strip():
+            arguments = parse_arguments_payload(
+                source.get(
+                    "arguments",
+                    source.get("parameters", payload.get("arguments")),
+                )
+            )
+            return ParsedToolCall(name=name.strip(), arguments=arguments, raw="")
+    extracted = extract_json_tool_fields(repair_jsonish(text))
+    if extracted is None:
+        return None
+    name, arguments = extracted
+    return ParsedToolCall(name=name, arguments=arguments, raw="")
+
+
+def parse_tool_calls(text: str) -> list[ParsedToolCall]:
+    calls: list[ParsedToolCall] = []
+    for block in TOOL_CALL_BLOCK_RE.finditer(text):
+        inner = block.group(1)
+        raw = block.group(0).strip()
+        found_xml = False
+        for function in FUNCTION_BLOCK_RE.finditer(inner):
+            found_xml = True
+            name = function.group(1).strip()
+            arguments: dict[str, Any] = {}
+            for parameter in PARAMETER_BLOCK_RE.finditer(function.group(2)):
+                arguments[parameter.group(1).strip()] = parse_argument_value(
+                    parameter.group(2)
+                )
+            calls.append(
+                ParsedToolCall(
+                    name=name,
+                    arguments=json_ready(arguments),
+                    raw=raw,
+                )
+            )
+        if found_xml:
+            continue
+        parsed = parse_json_tool_inner(inner)
+        if parsed is not None:
+            calls.append(
+                ParsedToolCall(name=parsed.name, arguments=parsed.arguments, raw=raw)
+            )
+    return calls
+
+
+def openai_tool_calls_from_text(text: str) -> list[dict[str, Any]]:
+    encoded: list[dict[str, Any]] = []
+    for index, call in enumerate(parse_tool_calls(text)):
+        encoded.append(
+            {
+                "id": f"call_{index}_{call.name}",
+                "type": "function",
+                "function": {
+                    "name": call.name,
+                    "arguments": json.dumps(call.arguments, ensure_ascii=False),
+                },
+            }
+        )
+    return encoded
+
+
+def assistant_message_content(message: dict[str, Any]) -> str:
+    reasoning = message.get("reasoning_content")
+    content = message_text(message.get("content"))
+    if isinstance(reasoning, str) and reasoning.strip():
+        content = f"{THINK_OPEN}\n{reasoning.strip()}\n{THINK_CLOSE}\n{content}"
+    tool_xml = format_tool_calls_xml(message.get("tool_calls"))
+    if tool_xml:
+        if content.strip():
+            return f"{content.rstrip()}\n{tool_xml}"
+        return tool_xml
+    return content
+
+
+def _flush_tool_group(
+    group: list[str],
+    messages: list[dict[str, str]],
+) -> None:
+    if not group:
+        return
+    messages.append({"role": "user", "content": "\n".join(group)})
+    group.clear()
+
+
+def prepare_inference_messages(
+    raw_messages: list[dict[str, Any]],
+    *,
+    tools: Any | None = None,
+) -> list[dict[str, str]]:
+    if not raw_messages:
+        raise ValueError("Chat history cannot be empty")
+
+    prepared: list[dict[str, str]] = []
+    pending_tool_results: list[str] = []
+    for index, raw_message in enumerate(raw_messages):
+        if not isinstance(raw_message, dict):
+            raise ValueError(f"Unsupported chat message at index {index}")
+        raw_role = raw_message.get("role")
+        if not isinstance(raw_role, str):
+            raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
+        role = INFERENCE_ROLE_MAP.get(raw_role.strip().lower())
+        if role is None:
+            raise ValueError(f"Unsupported chat role at index {index}: {raw_role!r}")
+        if role == "assistant":
+            content = assistant_message_content(raw_message)
+        else:
+            content = message_text(raw_message.get("content"))
+        if role == "tool":
+            content = wrap_tool_response(content)
+            if not content.strip():
+                raise ValueError(f"Chat content at index {index} must be non-empty")
+            pending_tool_results.append(content)
+            continue
+        _flush_tool_group(pending_tool_results, prepared)
+        if not content.strip():
+            if role == "system":
+                continue
+            raise ValueError(f"Chat content at index {index} must be non-empty")
+        if role == "system":
+            content = strip_leaked_chatml(content)
+        prepared.append({"role": role, "content": content})
+    _flush_tool_group(pending_tool_results, prepared)
+
+    tools_block = render_available_tools(tools)
+    if tools_block:
+        if prepared and prepared[0]["role"] == "system":
+            prepared[0] = {
+                "role": "system",
+                "content": inject_available_tools(prepared[0]["content"], tools),
+            }
+        else:
+            prepared.insert(0, {"role": "system", "content": tools_block})
+
+    if not prepared:
+        raise ValueError("Chat history cannot be empty")
+    return prepared

tiny_gdn/web_search.py

@@ -0,0 +1,475 @@
+"""Execute Nemotron Agentic-v2 `web-search` the way Tercet-R stage-3 SFT saw it.
+
+`nvidia/Nemotron-SFT-Agentic-v2` `data/search.jsonl` (packed as
+`agentic_v2_search`) declares a single tool named `web-search` whose
+observations are compact Tavily Search JSON:
+
+    {"query", "follow_up_questions", "answer", "images", "results",
+     "response_time", "request_id"}
+
+Each `results[]` item is `{url, title, content, score, raw_content}`
+(about ten hits, with an LLM `answer` when Tavily supplies one). The chat
+template then prefixes that JSON with `<|tool_response|>`.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from collections.abc import Mapping, Sequence
+from html import unescape
+from html.parser import HTMLParser
+from typing import Any
+from urllib.error import HTTPError, URLError
+from urllib.parse import parse_qs, quote, unquote, urlencode, urlparse
+from urllib.request import Request, urlopen
+
+
+NEMOTRON_WEB_SEARCH_TOOL: dict[str, Any] = {
+    "type": "function",
+    "function": {
+        "name": "web-search",
+        "description": "Search the web for a query.",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "query": {
+                    "type": "string",
+                    "description": "Search query.",
+                }
+            },
+            "required": ["query"],
+        },
+    },
+}
+
+WEB_SEARCH_TOOL_NAME = "web-search"
+DEFAULT_MAX_RESULTS = 10
+TAVILY_SEARCH_URL = "https://api.tavily.com/search"
+DUCKDUCKGO_HTML_URL = "https://html.duckduckgo.com/html/"
+WIKIPEDIA_API_URL = "https://en.wikipedia.org/w/api.php"
+USER_AGENT = (
+    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
+    "(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
+)
+REQUEST_TIMEOUT_SECONDS = 30.0
+
+
+def is_web_search_tool_name(name: str) -> bool:
+    return name.strip().lower().replace("_", "-") == WEB_SEARCH_TOOL_NAME
+
+
+def query_from_arguments(arguments: Mapping[str, Any] | None) -> str:
+    if not arguments:
+        return ""
+    raw = arguments.get("query")
+    if isinstance(raw, str):
+        return raw.strip()
+    if raw is None:
+        return ""
+    return str(raw).strip()
+
+
+def dumps_search_payload(payload: Mapping[str, Any]) -> str:
+    return json.dumps(canonical_search_payload(payload), ensure_ascii=False)
+
+
+def canonical_search_payload(payload: Mapping[str, Any]) -> dict[str, Any]:
+    results_raw = payload.get("results")
+    results: list[dict[str, Any]] = []
+    if isinstance(results_raw, Sequence) and not isinstance(results_raw, (str, bytes)):
+        for item in results_raw:
+            if isinstance(item, Mapping):
+                results.append(_canonical_result(item))
+    images_raw = payload.get("images")
+    images: list[Any] = []
+    if isinstance(images_raw, Sequence) and not isinstance(images_raw, (str, bytes)):
+        images = list(images_raw)
+    return {
+        "query": str(payload.get("query") or ""),
+        "follow_up_questions": payload.get("follow_up_questions"),
+        "answer": payload.get("answer"),
+        "images": images,
+        "results": results,
+        "response_time": payload.get("response_time"),
+        "request_id": payload.get("request_id"),
+    }
+
+
+def _canonical_result(item: Mapping[str, Any]) -> dict[str, Any]:
+    if item.get("error") and not item.get("url") and not item.get("title"):
+        return {"error": str(item["error"])}
+    result: dict[str, Any] = {
+        "url": str(item.get("url") or ""),
+        "title": str(item.get("title") or ""),
+        "content": str(item.get("content") or ""),
+        "score": _as_score(item.get("score")),
+        "raw_content": item.get("raw_content"),
+    }
+    if item.get("error"):
+        result["error"] = str(item["error"])
+    return result
+
+
+def _as_score(value: Any) -> float | None:
+    if value is None or value == "":
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def search_web(
+    query: str,
+    *,
+    max_results: int = DEFAULT_MAX_RESULTS,
+    api_key: str | None = None,
+) -> str:
+    """Return a Tavily-shaped JSON string for one `web-search` observation."""
+
+    cleaned = query.strip()
+    started = time.perf_counter()
+    if not cleaned:
+        return dumps_search_payload(
+            _empty_payload(
+                query,
+                error="Search query was empty",
+                response_time=time.perf_counter() - started,
+            )
+        )
+    key = (api_key if api_key is not None else os.environ.get("TAVILY_API_KEY", "")).strip()
+    if key:
+        try:
+            return dumps_search_payload(
+                _search_tavily(cleaned, api_key=key, max_results=max_results)
+            )
+        except Exception as error:
+            fallback_error = f"Tavily search failed: {error}"
+            try:
+                return dumps_search_payload(
+                    _search_without_tavily(
+                        cleaned,
+                        max_results=max_results,
+                        fallback_error=fallback_error,
+                    )
+                )
+            except Exception:
+                return dumps_search_payload(
+                    _empty_payload(
+                        cleaned,
+                        error=fallback_error,
+                        response_time=time.perf_counter() - started,
+                    )
+                )
+    try:
+        return dumps_search_payload(
+            _search_without_tavily(cleaned, max_results=max_results)
+        )
+    except Exception as error:
+        return dumps_search_payload(
+            _empty_payload(
+                cleaned,
+                error=str(error),
+                response_time=time.perf_counter() - started,
+            )
+        )
+
+
+def _empty_payload(
+    query: str,
+    *,
+    error: str,
+    response_time: float | None,
+) -> dict[str, Any]:
+    return {
+        "query": query,
+        "follow_up_questions": None,
+        "answer": None,
+        "images": [],
+        "results": [{"error": error}],
+        "response_time": response_time,
+        "request_id": None,
+    }
+
+
+def _search_tavily(
+    query: str,
+    *,
+    api_key: str,
+    max_results: int,
+) -> dict[str, Any]:
+    body = json.dumps(
+        {
+            "query": query,
+            "max_results": max(1, min(int(max_results), 20)),
+            "include_answer": True,
+            "include_raw_content": False,
+            "include_images": False,
+            "search_depth": "basic",
+        }
+    ).encode("utf-8")
+    payload = _read_json(
+        TAVILY_SEARCH_URL,
+        data=body,
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+            "Accept": "application/json",
+        },
+        method="POST",
+    )
+    if not isinstance(payload, dict):
+        raise RuntimeError("Tavily returned a non-object payload")
+    payload.setdefault("query", query)
+    return payload
+
+
+def _search_without_tavily(
+    query: str,
+    *,
+    max_results: int,
+    fallback_error: str | None = None,
+) -> dict[str, Any]:
+    started = time.perf_counter()
+    results: list[dict[str, Any]] = []
+    errors: list[str] = []
+    if fallback_error:
+        errors.append(fallback_error)
+    try:
+        results.extend(_search_duckduckgo_html(query, max_results=max_results))
+    except Exception as error:
+        errors.append(f"DuckDuckGo search failed: {error}")
+    if len(results) < max_results:
+        try:
+            wiki = _search_wikipedia(query, max_results=max_results)
+            seen = {item["url"] for item in results}
+            for item in wiki:
+                if item["url"] in seen:
+                    continue
+                results.append(item)
+                if len(results) >= max_results:
+                    break
+        except Exception as error:
+            errors.append(f"Wikipedia search failed: {error}")
+    elapsed = time.perf_counter() - started
+    if not results:
+        message = "; ".join(errors) if errors else "Search returned no results"
+        return _empty_payload(query, error=message, response_time=elapsed)
+    clipped = results[: max(1, min(int(max_results), 20))]
+    for index, item in enumerate(clipped):
+        if item.get("score") is None:
+            item["score"] = round(max(0.05, 1.0 - (index * 0.07)), 8)
+        item.setdefault("raw_content", None)
+    answer = clipped[0].get("content") if clipped else None
+    if isinstance(answer, str) and not answer.strip():
+        answer = None
+    return {
+        "query": query,
+        "follow_up_questions": None,
+        "answer": answer,
+        "images": [],
+        "results": clipped,
+        "response_time": elapsed,
+        "request_id": None,
+    }
+
+
+def _search_duckduckgo_html(query: str, *, max_results: int) -> list[dict[str, Any]]:
+    encoded = urlencode({"q": query, "kl": "us-en"})
+    html = _read_text(
+        DUCKDUCKGO_HTML_URL,
+        data=encoded.encode("utf-8"),
+        headers={
+            "Content-Type": "application/x-www-form-urlencoded",
+            "Accept": "text/html",
+        },
+        method="POST",
+    )
+    parser = _DuckDuckGoParser()
+    parser.feed(html)
+    parser.close()
+    results: list[dict[str, Any]] = []
+    seen: set[str] = set()
+    for item in parser.results:
+        url = item["url"]
+        if not url or url in seen:
+            continue
+        seen.add(url)
+        results.append(
+            {
+                "url": url,
+                "title": item["title"],
+                "content": item["content"],
+                "score": None,
+                "raw_content": None,
+            }
+        )
+        if len(results) >= max_results:
+            break
+    return results
+
+
+def _search_wikipedia(query: str, *, max_results: int) -> list[dict[str, Any]]:
+    params = urlencode(
+        {
+            "action": "query",
+            "list": "search",
+            "srsearch": query,
+            "srlimit": max(1, min(int(max_results), 10)),
+            "srprop": "snippet",
+            "format": "json",
+            "utf8": "1",
+        }
+    )
+    payload = _read_json(
+        f"{WIKIPEDIA_API_URL}?{params}",
+        headers={"Accept": "application/json"},
+        method="GET",
+    )
+    if not isinstance(payload, dict):
+        return []
+    query_block = payload.get("query")
+    hits = query_block.get("search") if isinstance(query_block, dict) else None
+    if not isinstance(hits, list):
+        return []
+    results: list[dict[str, Any]] = []
+    for hit in hits:
+        if not isinstance(hit, dict):
+            continue
+        title = str(hit.get("title") or "").strip()
+        if not title:
+            continue
+        snippet = unescape(_strip_tags(str(hit.get("snippet") or ""))).strip()
+        results.append(
+            {
+                "url": f"https://en.wikipedia.org/wiki/{quote(title.replace(' ', '_'))}",
+                "title": title,
+                "content": snippet,
+                "score": None,
+                "raw_content": None,
+            }
+        )
+    return results
+
+
+def unwrap_ddg_url(href: str) -> str:
+    raw = href.strip()
+    if raw.startswith("//"):
+        raw = f"https:{raw}"
+    parsed = urlparse(raw)
+    host = parsed.netloc.lower()
+    if "duckduckgo.com" in host:
+        values = parse_qs(parsed.query).get("uddg") or []
+        if values:
+            return unquote(values[0])
+    return raw
+
+
+def _strip_tags(value: str) -> str:
+    return _TagStripper.strip(value)
+
+
+class _TagStripper(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__()
+        self.chunks: list[str] = []
+
+    def handle_data(self, data: str) -> None:
+        self.chunks.append(data)
+
+    @classmethod
+    def strip(cls, value: str) -> str:
+        parser = cls()
+        parser.feed(value)
+        parser.close()
+        return "".join(parser.chunks)
+
+
+class _DuckDuckGoParser(HTMLParser):
+    def __init__(self) -> None:
+        super().__init__()
+        self.results: list[dict[str, str]] = []
+        self._in_title = False
+        self._in_snippet = False
+        self._href = ""
+        self._buffer: list[str] = []
+        self._pending: dict[str, str] | None = None
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        classes = {item for item in (dict(attrs).get("class") or "").split() if item}
+        if tag == "a" and "result__a" in classes:
+            self._flush_pending()
+            self._in_title = True
+            self._href = dict(attrs).get("href") or ""
+            self._buffer = []
+            return
+        if tag in {"a", "div"} and "result__snippet" in classes:
+            self._in_snippet = True
+            self._buffer = []
+
+    def handle_endtag(self, tag: str) -> None:
+        if tag == "a" and self._in_title:
+            title = unescape("".join(self._buffer)).strip()
+            url = unwrap_ddg_url(self._href)
+            self._pending = {"title": title, "url": url, "content": ""}
+            self._in_title = False
+            self._buffer = []
+            return
+        if self._in_snippet and tag in {"a", "div"}:
+            snippet = unescape("".join(self._buffer)).strip()
+            if self._pending is not None:
+                self._pending["content"] = snippet
+                self.results.append(self._pending)
+                self._pending = None
+            self._in_snippet = False
+            self._buffer = []
+
+    def handle_data(self, data: str) -> None:
+        if self._in_title or self._in_snippet:
+            self._buffer.append(data)
+
+    def close(self) -> None:
+        self._flush_pending()
+        super().close()
+
+    def _flush_pending(self) -> None:
+        if self._pending is None:
+            return
+        self.results.append(self._pending)
+        self._pending = None
+
+
+def _read_json(
+    url: str,
+    *,
+    data: bytes | None = None,
+    headers: Mapping[str, str],
+    method: str,
+) -> Any:
+    text = _read_text(url, data=data, headers=headers, method=method)
+    return json.loads(text)
+
+
+def _read_text(
+    url: str,
+    *,
+    data: bytes | None = None,
+    headers: Mapping[str, str],
+    method: str,
+) -> str:
+    merged = {"User-Agent": USER_AGENT, **dict(headers)}
+    request = Request(url, data=data, headers=merged, method=method)
+    try:
+        with urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
+            raw = response.read()
+            charset = response.headers.get_content_charset() or "utf-8"
+    except HTTPError as error:
+        detail = error.read().decode("utf-8", errors="replace")
+        raise RuntimeError(
+            f"{error.code} {error.reason} from {url}: {detail[:500]}"
+        ) from error
+    except URLError as error:
+        raise RuntimeError(f"Could not reach {url}: {error.reason}") from error
+    return raw.decode(charset, errors="replace")

tokenizer_config.json

@@ -0,0 +1,74 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "bos_token": "<|begin_of_text|>",
+  "eos_token": "<|end_of_text|>",
+  "pad_token": "<|padding|>",
+  "unk_token": "<|unknown|>",
+  "model_max_length": 2048,
+  "clean_up_tokenization_spaces": false,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "chat_template": "{%- set ns = namespace(xml_tools=none, tools_emitted=false) -%}\n{%- if xml_tools is defined and xml_tools -%}\n{%- set ns.xml_tools = xml_tools -%}\n{%- elif tools is defined and tools -%}\n{%- set ns.xml_tools = tools -%}\n{%- endif -%}\n{%- set tools_preamble = 'You may call one or more functions to assist with the user query.\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\\n' -%}\n{%- set tools_epilogue = '</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\\n</tool_call>' -%}\n{%- for message in messages -%}\n{%- if loop.first -%}{{- bos_token -}}{%- endif -%}\n{%- if loop.first and ns.xml_tools and message['role'] != 'system' -%}\n{{- '<|im_start|>system\\n' + tools_preamble -}}\n{%- for tool in ns.xml_tools -%}\n{%- if tool is string -%}{{- (tool | replace('<tools>', '') | replace('</tools>', '') | trim) + '\\n' -}}\n{%- else -%}{{- tool | tojson + '\\n' -}}\n{%- endif -%}\n{%- endfor -%}\n{{- tools_epilogue + '<|im_end|>\\n' -}}\n{%- set ns.tools_emitted = true -%}\n{%- endif -%}\n{%- set raw_role = message['role'] -%}\n{%- set role = 'user' if raw_role == 'tool' or raw_role == 'function' else raw_role -%}\n{%- set content_text = namespace(value='') -%}\n{%- if message['content'] is string -%}\n{%- set content_text.value = message['content'] -%}\n{%- elif message['content'] is iterable -%}\n{%- for item in message['content'] -%}\n{%- if item['type'] == 'text' -%}{%- set content_text.value = content_text.value + item['text'] -%}{%- endif -%}\n{%- endfor -%}\n{%- endif -%}\n{%- set is_tool = raw_role == 'tool' or raw_role == 'function' -%}\n{%- set prev_is_tool = loop.previtem is defined and (loop.previtem['role'] == 'tool' or loop.previtem['role'] == 'function') -%}\n{%- set next_is_tool = loop.nextitem is defined and (loop.nextitem['role'] == 'tool' or loop.nextitem['role'] == 'function') -%}\n{%- if raw_role == 'system' and not (content_text.value | trim) and not (ns.xml_tools and not ns.tools_emitted) -%}\n{%- else -%}\n{%- if is_tool and prev_is_tool -%}\n{{- '\\n' -}}\n{%- else -%}\n{{- '<|im_start|>' + role + '\\n' -}}\n{%- endif -%}\n{%- if raw_role == 'assistant' and enable_thinking is defined -%}\n{{- ('<|think|>\\n' if enable_thinking else '<|no_think|>\\n') -}}\n{%- endif -%}\n{%- if is_tool -%}\n{{- '<|tool_response|>\\n' + (content_text.value | replace('<tool_response>', '') | replace('</tool_response>', '') | trim) -}}\n{%- elif raw_role == 'system' -%}\n{{- content_text.value | replace('/system_override', '') | replace('/no_think', '') | replace('/think', '') | trim -}}\n{%- else -%}\n{{- content_text.value -}}\n{%- endif -%}\n{%- if raw_role == 'system' and ns.xml_tools and not ns.tools_emitted and '<tools>' not in content_text.value -%}\n{%- if content_text.value | trim -%}{{- '\\n\\n' -}}{%- endif -%}\n{{- tools_preamble -}}\n{%- for tool in ns.xml_tools -%}\n{%- if tool is string -%}{{- (tool | replace('<tools>', '') | replace('</tools>', '') | trim) + '\\n' -}}\n{%- else -%}{{- tool | tojson + '\\n' -}}\n{%- endif -%}\n{%- endfor -%}\n{{- tools_epilogue -}}\n{%- set ns.tools_emitted = true -%}\n{%- endif -%}\n{%- if raw_role == 'assistant' and message['tool_calls'] is defined and message['tool_calls'] -%}\n{%- if '<tool_call>' not in content_text.value -%}\n{%- for tool_call in message['tool_calls'] -%}\n{%- set fn = tool_call['function'] if tool_call['function'] is defined else tool_call -%}\n{%- if loop.first and not (content_text.value | trim) -%}\n{{- '<tool_call>\\n{\"name\": \"' + fn['name'] + '\", \"arguments\": ' -}}\n{%- else -%}\n{{- '\\n<tool_call>\\n{\"name\": \"' + fn['name'] + '\", \"arguments\": ' -}}\n{%- endif -%}\n{%- if fn['arguments'] is string -%}{{- fn['arguments'] -}}\n{%- else -%}{{- fn['arguments'] | tojson -}}\n{%- endif -%}\n{{- '}\\n</tool_call>' -}}\n{%- endfor -%}\n{%- endif -%}\n{%- endif -%}\n{%- if not (is_tool and next_is_tool) -%}\n{{- '<|im_end|>\\n' -}}\n{%- endif -%}\n{%- endif -%}\n{%- endfor -%}\n{%- if add_generation_prompt -%}\n{{- '<|im_start|>assistant\\n' -}}\n{%- if enable_thinking is defined -%}{{- ('<|think|>\\n' if enable_thinking else '<|no_think|>\\n') -}}{%- endif -%}\n{%- endif -%}",
+  "extra_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|tool_call|>",
+    "<|tool_response|>",
+    "<think>",
+    "</think>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|fim_pad|>",
+    "<|no_think|>",
+    "<|think|>",
+    "<|reserved_002|>",
+    "<|reserved_003|>",
+    "<|reserved_004|>",
+    "<|reserved_005|>",
+    "<|reserved_006|>",
+    "<|reserved_007|>",
+    "<|reserved_008|>",
+    "<|reserved_009|>",
+    "<|reserved_010|>",
+    "<|reserved_011|>",
+    "<|reserved_012|>",
+    "<|reserved_013|>",
+    "<|reserved_014|>",
+    "<|reserved_015|>",
+    "<|reserved_016|>",
+    "<|reserved_017|>",
+    "<|reserved_018|>",
+    "<|reserved_019|>",
+    "<|reserved_020|>",
+    "<|reserved_021|>",
+    "<|reserved_022|>",
+    "<|reserved_023|>",
+    "<|reserved_024|>",
+    "<|reserved_025|>",
+    "<|reserved_026|>",
+    "<|reserved_027|>",
+    "<|reserved_028|>",
+    "<|reserved_029|>",
+    "<|reserved_030|>",
+    "<|reserved_031|>",
+    "<|reserved_032|>",
+    "<|reserved_033|>",
+    "<|reserved_034|>",
+    "<|reserved_035|>",
+    "<|reserved_036|>",
+    "<|reserved_037|>",
+    "<|reserved_038|>",
+    "<|reserved_039|>",
+    "<|reserved_040|>",
+    "<|reserved_041|>",
+    "<|reserved_042|>",
+    "<|reserved_043|>",
+    "<|reserved_044|>",
+    "<|reserved_045|>",
+    "<|reserved_046|>",
+    "<|reserved_047|>",
+    "<|reserved_048|>",
+    "<|reserved_049|>"
+  ]
+}

validation.json

@@ -0,0 +1,21 @@
+{
+  "checkpoint": "checkpoint-00005300",
+  "ema": {
+    "assistant_targets": 97349,
+    "batches": 16,
+    "elapsed_seconds": 2.1927712520046043,
+    "loss": 2.0271428378822587,
+    "perplexity": 7.592362728419498,
+    "targets_per_second": 44395.419682287764
+  },
+  "normal": {
+    "assistant_targets": 97349,
+    "batches": 16,
+    "elapsed_seconds": 6.0049016539996956,
+    "loss": 2.0189123026430678,
+    "perplexity": 7.530129975181167,
+    "targets_per_second": 16211.589399662953
+  },
+  "optimizer_step": 5300,
+  "type": "sft_validation"
+}

windows_fla_patches/fla/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Keep package import light. Eagerly importing fla.layers pulls every Triton
+# kernel and breaks on Windows + Triton 3.7.
+
+from pkgutil import extend_path
+
+__path__ = extend_path(__path__, __name__)
+__version__ = "0.5.2"
+__all__: list[str] = []

windows_fla_patches/fla/layers/__init__.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Lazy layer exports — avoid compiling every Triton kernel at import time.
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_EXPORTS: dict[str, tuple[str, str]] = {
+    "ABCAttention": (".abc", "ABCAttention"),
+    "Attention": (".attn", "Attention"),
+    "BasedLinearAttention": (".based", "BasedLinearAttention"),
+    "BitAttention": (".bitattn", "BitAttention"),
+    "Comba": (".comba", "Comba"),
+    "DeltaNet": (".delta_net", "DeltaNet"),
+    "DeltaFormerAttention": (".deltaformer", "DeltaFormerAttention"),
+    "ForgettingAttention": (".forgetting_attn", "ForgettingAttention"),
+    "GatedDeltaNet": (".gated_deltanet", "GatedDeltaNet"),
+    "GatedDeltaProduct": (".gated_deltaproduct", "GatedDeltaProduct"),
+    "GatedDeltaNet2": (".gdn2", "GatedDeltaNet2"),
+    "GatedLinearAttention": (".gla", "GatedLinearAttention"),
+    "GatedSlotAttention": (".gsa", "GatedSlotAttention"),
+    "HGRNAttention": (".hgrn", "HGRNAttention"),
+    "HGRN2Attention": (".hgrn2", "HGRN2Attention"),
+    "KimiDeltaAttention": (".kda", "KimiDeltaAttention"),
+    "LightNetAttention": (".lightnet", "LightNetAttention"),
+    "LinearAttention": (".linear_attn", "LinearAttention"),
+    "LogLinearMamba2": (".log_linear_mamba2", "LogLinearMamba2"),
+    "Mamba": (".mamba", "Mamba"),
+    "Mamba2": (".mamba2", "Mamba2"),
+    "Mamba3": (".mamba3", "Mamba3"),
+    "MesaNet": (".mesa_net", "MesaNet"),
+    "MultiheadLatentAttention": (".mla", "MultiheadLatentAttention"),
+    "MoBA": (".moba", "MoBA"),
+    "MomAttention": (".mom", "MomAttention"),
+    "MultiScaleRetention": (".multiscale_retention", "MultiScaleRetention"),
+    "NativeSparseAttention": (".nsa", "NativeSparseAttention"),
+    "Parallax": (".parallax", "Parallax"),
+    "PaTHAttention": (".path_attn", "PaTHAttention"),
+    "Raven": (".raven", "Raven"),
+    "ReBasedLinearAttention": (".rebased", "ReBasedLinearAttention"),
+    "RodimusAttention": (".rodimus", "RodimusAttention"),
+    "SlidingWindowSharedKeyAttention": (".rodimus", "SlidingWindowSharedKeyAttention"),
+    "RWKV6Attention": (".rwkv6", "RWKV6Attention"),
+    "RWKV7Attention": (".rwkv7", "RWKV7Attention"),
+    "WallAttention": (".wall_attn", "WallAttention"),
+    "YOCOCrossAttention": (".yoco", "YOCOCrossAttention"),
+    "YOCOGatedRetention": (".yoco", "YOCOGatedRetention"),
+    "YOCOSharedKVBuilder": (".yoco", "YOCOSharedKVBuilder"),
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    spec = _EXPORTS.get(name)
+    if spec is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module_name, attr = spec
+    value = getattr(importlib.import_module(module_name, __name__), attr)
+    globals()[name] = value
+    return value

windows_fla_patches/fla/ops/__init__.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# Lazy public exports so importing fla.ops.utils / fla.ops.gdn2 does not
+# eagerly compile every Triton kernel (needed on Windows + Triton 3.7).
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_EXPORTS: dict[str, str] = {
+    "chunk_abc": "fla.ops.abc",
+    "parallel_attn": "fla.ops.attn",
+    "fused_attnres": "fla.ops.attnres",
+    "fused_chunk_based": "fla.ops.based",
+    "parallel_based": "fla.ops.based",
+    "chunk_comba": "fla.ops.comba",
+    "fused_recurrent_comba": "fla.ops.comba",
+    "chunk_delta_rule": "fla.ops.delta_rule",
+    "fused_chunk_delta_rule": "fla.ops.delta_rule",
+    "fused_recurrent_delta_rule": "fla.ops.delta_rule",
+    "parallel_forgetting_attn": "fla.ops.forgetting_attn",
+    "chunk_gated_delta_rule": "fla.ops.gated_delta_rule",
+    "chunk_gdn": "fla.ops.gated_delta_rule",
+    "fused_recurrent_gated_delta_rule": "fla.ops.gated_delta_rule",
+    "fused_recurrent_gdn": "fla.ops.gated_delta_rule",
+    "chunk_dplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "chunk_iplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "fused_recurrent_dplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "fused_recurrent_iplr_delta_rule": "fla.ops.generalized_delta_rule",
+    "chunk_gla": "fla.ops.gla",
+    "fused_chunk_gla": "fla.ops.gla",
+    "fused_recurrent_gla": "fla.ops.gla",
+    "chunk_gsa": "fla.ops.gsa",
+    "fused_recurrent_gsa": "fla.ops.gsa",
+    "fused_recurrent_hgrn": "fla.ops.hgrn",
+    "chunk_kda": "fla.ops.kda",
+    "fused_recurrent_kda": "fla.ops.kda",
+    "chunk_lightning_attn": "fla.ops.lightning_attn",
+    "fused_recurrent_lightning_attn": "fla.ops.lightning_attn",
+    "chunk_linear_attn": "fla.ops.linear_attn",
+    "fused_chunk_linear_attn": "fla.ops.linear_attn",
+    "fused_recurrent_linear_attn": "fla.ops.linear_attn",
+    "chunk_log_linear_attn": "fla.ops.log_linear_attn",
+    "chunk_mesa_net": "fla.ops.mesa_net",
+    "parallel_nsa": "fla.ops.nsa",
+    "parallel_parallax": "fla.ops.parallax",
+    "parallel_path_attn": "fla.ops.path_attn",
+    "chunk_retention": "fla.ops.retention",
+    "fused_chunk_retention": "fla.ops.retention",
+    "fused_recurrent_retention": "fla.ops.retention",
+    "parallel_retention": "fla.ops.retention",
+    "chunk_rwkv6": "fla.ops.rwkv6",
+    "fused_recurrent_rwkv6": "fla.ops.rwkv6",
+    "chunk_rwkv7": "fla.ops.rwkv7",
+    "fused_recurrent_rwkv7": "fla.ops.rwkv7",
+    "chunk_simple_gla": "fla.ops.simple_gla",
+    "fused_chunk_simple_gla": "fla.ops.simple_gla",
+    "fused_recurrent_simple_gla": "fla.ops.simple_gla",
+    "parallel_simple_gla": "fla.ops.simple_gla",
+    "parallel_wall_attn": "fla.ops.wall_attn",
+    "parallel_wall_attn_decode": "fla.ops.wall_attn",
+    "chunk_gdn2": "fla.ops.gdn2",
+    "fused_recurrent_gdn2": "fla.ops.gdn2",
+}
+
+__all__ = list(_EXPORTS)
+
+
+def __getattr__(name: str) -> Any:
+    module_name = _EXPORTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(importlib.import_module(module_name), name)
+    globals()[name] = value
+    return value

windows_fla_patches/fla/ops/simple_gla/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# For a list of all contributors, visit:
+#   https://github.com/fla-org/flash-linear-attention/graphs/contributors
+
+from .chunk import chunk_simple_gla
+from .fused_chunk import fused_chunk_simple_gla
+from .fused_recurrent import fused_recurrent_simple_gla
+
+# Triton 3.7 on Windows can fail while decorating parallel kernels at import time.
+try:
+    from .parallel import parallel_simple_gla
+except Exception:  # noqa: BLE001
+    parallel_simple_gla = None
+
+__all__ = [
+    'chunk_simple_gla',
+    'fused_chunk_simple_gla',
+    'fused_recurrent_simple_gla',
+    'parallel_simple_gla',
+]