@relay-harness/coding-agent 0.0.0-stage → 1.0.3

package/dist/extensions/laya/telemetry.js

@@ -0,0 +1,107 @@
+import { appendFileSync, existsSync, mkdirSync, readFileSync } from "node:fs";
+import { dirname } from "node:path";
+import { PerformanceHistory } from "./policy.js";
+const MAX_REQUEST_CHARS = 2000;
+const TEST_COMMAND = /\b(test|tests|vitest|jest|mocha|pytest|go test|cargo test|mvn|gradle|tsc|check|lint|build)\b|test\.sh/;
+/** Reads `[laya:*]` and `[harness:*]` messages as harness text, not as the user's words. */
+export const HARNESS_MESSAGE = /^\[(laya|harness):[a-z-]+\]/;
+/** Outcome of the latest check (test, type check, build, lint) run through bash, or null when none ran. */
+export function lastCheckPassed(messages) {
+    const commands = new Map();
+    let passed = null;
+    for (const message of messages) {
+        if (message.role === "assistant") {
+            for (const block of message.content) {
+                if (block.type === "toolCall" && typeof block.arguments.command === "string") {
+                    commands.set(block.id, block.arguments.command);
+                }
+            }
+        }
+        else if (message.role === "toolResult" && message.toolName === "bash") {
+            const command = commands.get(message.toolCallId);
+            if (command && TEST_COMMAND.test(command))
+                passed = !message.isError;
+        }
+    }
+    return passed;
+}
+/** Outcome, usage and checks of one agent run. */
+export function summarizeRun(messages) {
+    let tokens = 0;
+    let cost = 0;
+    let toolFailures = 0;
+    const models = new Set();
+    let outcome = "completed";
+    for (const message of messages) {
+        if (message.role === "assistant") {
+            tokens += message.usage?.totalTokens ?? 0;
+            cost += message.usage?.cost.total ?? 0;
+            models.add(`${message.provider}/${message.model}`);
+            outcome =
+                message.stopReason === "error" ? "error" : message.stopReason === "aborted" ? "aborted" : "completed";
+        }
+        else if (message.role === "toolResult" && message.isError) {
+            toolFailures++;
+        }
+    }
+    return { outcome, testsPassed: lastCheckPassed(messages), toolFailures, tokens, cost, models: [...models] };
+}
+export function truncateRequest(request) {
+    return request.length > MAX_REQUEST_CHARS ? `${request.slice(0, MAX_REQUEST_CHARS)}…` : request;
+}
+/** Append-only JSONL store. Unreadable lines are skipped. */
+export class TelemetryStore {
+    path;
+    constructor(path) {
+        this.path = path;
+    }
+    append(record) {
+        mkdirSync(dirname(this.path), { recursive: true });
+        appendFileSync(this.path, `${JSON.stringify(record)}\n`, "utf8");
+    }
+    /** The newest `limit` records, oldest first. */
+    read(limit = 5000) {
+        if (!existsSync(this.path))
+            return [];
+        const records = [];
+        for (const line of readFileSync(this.path, "utf8").split("\n").slice(-limit)) {
+            if (!line.trim())
+                continue;
+            try {
+                records.push(JSON.parse(line));
+            }
+            catch {
+                // A partially written line from a crashed process.
+            }
+        }
+        return records;
+    }
+}
+export function historyKey(record) {
+    return `${record.classification.type}:${Math.round(record.classification.complexity * 4)}`;
+}
+/** Aborted runs say nothing about the model, so they are not counted. */
+export function historyFromTelemetry(records) {
+    const history = new PerformanceHistory();
+    for (const record of records) {
+        if (record.result.outcome === "aborted")
+            continue;
+        history.record(historyKey(record), `${record.selected.provider}/${record.selected.model}`, record.result.success);
+    }
+    return history;
+}
+/**
+ * Training exercises with evidence-based tier labels: for each successful run, the tier that
+ * finished it. Only `capability_tier` is labeled; the other answers came from Laya itself and would
+ * teach it nothing new. Review the rows before adding them (`/laya-data`).
+ */
+export function datasetFromTelemetry(records) {
+    return records
+        .filter((record) => record.result.success && record.request.trim())
+        .map((record) => ({
+        state: { request: record.request },
+        expected: { capability_tier: record.selected.tier },
+        source: "harness",
+    }));
+}
+//# sourceMappingURL=telemetry.js.map

package/dist/extensions/laya/telemetry.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"telemetry.js","sourceRoot":"","sources":["../../../src/extensions/laya/telemetry.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,cAAc,EAAE,UAAU,EAAE,SAAS,EAAE,YAAY,EAAE,MAAM,SAAS,CAAC;AAC9E,OAAO,EAAE,OAAO,EAAE,MAAM,WAAW,CAAC;AAGpC,OAAO,EAAyB,kBAAkB,EAAsB,MAAM,aAAa,CAAC;AAoC5F,MAAM,iBAAiB,GAAG,IAAI,CAAC;AAC/B,MAAM,YAAY,GACjB,uGAAuG,CAAC;AAEzG,4FAA4F;AAC5F,MAAM,CAAC,MAAM,eAAe,GAAG,6BAA6B,CAAC;AAE7D,2GAA2G;AAC3G,MAAM,UAAU,eAAe,CAAC,QAAiC;IAChE,MAAM,QAAQ,GAAG,IAAI,GAAG,EAAkB,CAAC;IAC3C,IAAI,MAAM,GAAmB,IAAI,CAAC;IAClC,KAAK,MAAM,OAAO,IAAI,QAAQ,EAAE,CAAC;QAChC,IAAI,OAAO,CAAC,IAAI,KAAK,WAAW,EAAE,CAAC;YAClC,KAAK,MAAM,KAAK,IAAI,OAAO,CAAC,OAAO,EAAE,CAAC;gBACrC,IAAI,KAAK,CAAC,IAAI,KAAK,UAAU,IAAI,OAAO,KAAK,CAAC,SAAS,CAAC,OAAO,KAAK,QAAQ,EAAE,CAAC;oBAC9E,QAAQ,CAAC,GAAG,CAAC,KAAK,CAAC,EAAE,EAAE,KAAK,CAAC,SAAS,CAAC,OAAO,CAAC,CAAC;gBACjD,CAAC;YACF,CAAC;QACF,CAAC;aAAM,IAAI,OAAO,CAAC,IAAI,KAAK,YAAY,IAAI,OAAO,CAAC,QAAQ,KAAK,MAAM,EAAE,CAAC;YACzE,MAAM,OAAO,GAAG,QAAQ,CAAC,GAAG,CAAC,OAAO,CAAC,UAAU,CAAC,CAAC;YACjD,IAAI,OAAO,IAAI,YAAY,CAAC,IAAI,CAAC,OAAO,CAAC;gBAAE,MAAM,GAAG,CAAC,OAAO,CAAC,OAAO,CAAC;QACtE,CAAC;IACF,CAAC;IACD,OAAO,MAAM,CAAC;AACf,CAAC;AAED,kDAAkD;AAClD,MAAM,UAAU,YAAY,CAAC,QAAiC;IAQ7D,IAAI,MAAM,GAAG,CAAC,CAAC;IACf,IAAI,IAAI,GAAG,CAAC,CAAC;IACb,IAAI,YAAY,GAAG,CAAC,CAAC;IACrB,MAAM,MAAM,GAAG,IAAI,GAAG,EAAU,CAAC;IACjC,IAAI,OAAO,GAAyC,WAAW,CAAC;IAChE,KAAK,MAAM,OAAO,IAAI,QAAQ,EAAE,CAAC;QAChC,IAAI,OAAO,CAAC,IAAI,KAAK,WAAW,EAAE,CAAC;YAClC,MAAM,IAAI,OAAO,CAAC,KAAK,EAAE,WAAW,IAAI,CAAC,CAAC;YAC1C,IAAI,IAAI,OAAO,CAAC,KAAK,EAAE,IAAI,CAAC,KAAK,IAAI,CAAC,CAAC;YACvC,MAAM,CAAC,GAAG,CAAC,GAAG,OAAO,CAAC,QAAQ,IAAI,OAAO,CAAC,KAAK,EAAE,CAAC,CAAC;YACnD,OAAO;gBACN,OAAO,CAAC,UAAU,KAAK,OAAO,CAAC,CAAC,CAAC,OAAO,CAAC,CAAC,CAAC,OAAO,CAAC,UAAU,KAAK,SAAS,CAAC,CAAC,CAAC,SAAS,CAAC,CAAC,CAAC,WAAW,CAAC;QACxG,CAAC;aAAM,IAAI,OAAO,CAAC,IAAI,KAAK,YAAY,IAAI,OAAO,CAAC,OAAO,EAAE,CAAC;YAC7D,YAAY,EAAE,CAAC;QAChB,CAAC;IACF,CAAC;IACD,OAAO,EAAE,OAAO,EAAE,WAAW,EAAE,eAAe,CAAC,QAAQ,CAAC,EAAE,YAAY,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,CAAC,GAAG,MAAM,CAAC,EAAE,CAAC;AAC7G,CAAC;AAED,MAAM,UAAU,eAAe,CAAC,OAAe;IAC9C,OAAO,OAAO,CAAC,MAAM,GAAG,iBAAiB,CAAC,CAAC,CAAC,GAAG,OAAO,CAAC,KAAK,CAAC,CAAC,EAAE,iBAAiB,CAAC,GAAG,CAAC,CAAC,CAAC,OAAO,CAAC;AACjG,CAAC;AAED,6DAA6D;AAC7D,MAAM,OAAO,cAAc;IACjB,IAAI,CAAS;IAEtB,YAAY,IAAY;QACvB,IAAI,CAAC,IAAI,GAAG,IAAI,CAAC;IAClB,CAAC;IAED,MAAM,CAAC,MAAuB;QAC7B,SAAS,CAAC,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,EAAE,EAAE,SAAS,EAAE,IAAI,EAAE,CAAC,CAAC;QACnD,cAAc,CAAC,IAAI,CAAC,IAAI,EAAE,GAAG,IAAI,CAAC,SAAS,CAAC,MAAM,CAAC,IAAI,EAAE,MAAM,CAAC,CAAC;IAClE,CAAC;IAED,gDAAgD;IAChD,IAAI,CAAC,KAAK,GAAG,IAAI;QAChB,IAAI,CAAC,UAAU,CAAC,IAAI,CAAC,IAAI,CAAC;YAAE,OAAO,EAAE,CAAC;QACtC,MAAM,OAAO,GAAsB,EAAE,CAAC;QACtC,KAAK,MAAM,IAAI,IAAI,YAAY,CAAC,IAAI,CAAC,IAAI,EAAE,MAAM,CAAC,CAAC,KAAK,CAAC,IAAI,CAAC,CAAC,KAAK,CAAC,CAAC,KAAK,CAAC,EAAE,CAAC;YAC9E,IAAI,CAAC,IAAI,CAAC,IAAI,EAAE;gBAAE,SAAS;YAC3B,IAAI,CAAC;gBACJ,OAAO,CAAC,IAAI,CAAC,IAAI,CAAC,KAAK,CAAC,IAAI,CAAoB,CAAC,CAAC;YACnD,CAAC;YAAC,MAAM,CAAC;gBACR,mDAAmD;YACpD,CAAC;QACF,CAAC;QACD,OAAO,OAAO,CAAC;IAChB,CAAC;CACD;AAED,MAAM,UAAU,UAAU,CAAC,MAA+C;IACzE,OAAO,GAAG,MAAM,CAAC,cAAc,CAAC,IAAI,IAAI,IAAI,CAAC,KAAK,CAAC,MAAM,CAAC,cAAc,CAAC,UAAU,GAAG,CAAC,CAAC,EAAE,CAAC;AAC5F,CAAC;AAED,yEAAyE;AACzE,MAAM,UAAU,oBAAoB,CAAC,OAAmC;IACvE,MAAM,OAAO,GAAG,IAAI,kBAAkB,EAAE,CAAC;IACzC,KAAK,MAAM,MAAM,IAAI,OAAO,EAAE,CAAC;QAC9B,IAAI,MAAM,CAAC,MAAM,CAAC,OAAO,KAAK,SAAS;YAAE,SAAS;QAClD,OAAO,CAAC,MAAM,CAAC,UAAU,CAAC,MAAM,CAAC,EAAE,GAAG,MAAM,CAAC,QAAQ,CAAC,QAAQ,IAAI,MAAM,CAAC,QAAQ,CAAC,KAAK,EAAE,EAAE,MAAM,CAAC,MAAM,CAAC,OAAO,CAAC,CAAC;IACnH,CAAC;IACD,OAAO,OAAO,CAAC;AAChB,CAAC;AAED;;;;GAIG;AACH,MAAM,UAAU,oBAAoB,CAAC,OAAmC;IAKvE,OAAO,OAAO;SACZ,MAAM,CAAC,CAAC,MAAM,EAAE,EAAE,CAAC,MAAM,CAAC,MAAM,CAAC,OAAO,IAAI,MAAM,CAAC,OAAO,CAAC,IAAI,EAAE,CAAC;SAClE,GAAG,CAAC,CAAC,MAAM,EAAE,EAAE,CAAC,CAAC;QACjB,KAAK,EAAE,EAAE,OAAO,EAAE,MAAM,CAAC,OAAO,EAAE;QAClC,QAAQ,EAAE,EAAE,eAAe,EAAE,MAAM,CAAC,QAAQ,CAAC,IAAI,EAAE;QACnD,MAAM,EAAE,SAAS;KACjB,CAAC,CAAC,CAAC;AACN,CAAC","sourcesContent":["import { appendFileSync, existsSync, mkdirSync, readFileSync } from \"node:fs\";\nimport { dirname } from \"node:path\";\nimport type { AgentMessage } from \"@relay-harness/agent-core\";\nimport type { TaskAssessment } from \"./assessment.ts\";\nimport { type EscalationReason, PerformanceHistory, type PolicyProfile } from \"./policy.ts\";\nimport type { CapabilityTier, TaskScope, TaskType } from \"./questions.ts\";\n\n/** One routed request and its outcome, as written to `telemetry.jsonl`. */\nexport interface TelemetryRecord {\n\ttask_id: string;\n\ttimestamp: string;\n\t/** The user's request, truncated. Kept locally to export training exercises. */\n\trequest: string;\n\tclassification: {\n\t\tsource: TaskAssessment[\"source\"];\n\t\ttype: TaskType;\n\t\tcomplexity: number;\n\t\trisk: number;\n\t\tambiguity: number;\n\t\treasoning: number;\n\t\tscope: TaskScope;\n\t\tagent: string;\n\t\tvalidation: string;\n\t\tconfidence: number;\n\t\trecommended_tier: CapabilityTier;\n\t\trecommended_effort: string;\n\t};\n\tpolicy: { profile: PolicyProfile; required_tier: CapabilityTier; rules: string[] };\n\tselected: { provider: string; model: string; tier: CapabilityTier; thinking_level: string };\n\tresult: {\n\t\tsuccess: boolean;\n\t\toutcome: \"completed\" | \"error\" | \"aborted\";\n\t\tattempts: number;\n\t\tescalations: Array<{ reason: EscalationReason; from: CapabilityTier; to: CapabilityTier }>;\n\t\ttests_passed: boolean | null;\n\t\ttool_failures: number;\n\t};\n\tusage: { tokens: number; cost: number; duration_ms: number; models: string[] };\n}\n\nconst MAX_REQUEST_CHARS = 2000;\nconst TEST_COMMAND =\n\t/\\b(test|tests|vitest|jest|mocha|pytest|go test|cargo test|mvn|gradle|tsc|check|lint|build)\\b|test\\.sh/;\n\n/** Reads `[laya:*]` and `[harness:*]` messages as harness text, not as the user's words. */\nexport const HARNESS_MESSAGE = /^\\[(laya|harness):[a-z-]+\\]/;\n\n/** Outcome of the latest check (test, type check, build, lint) run through bash, or null when none ran. */\nexport function lastCheckPassed(messages: readonly AgentMessage[]): boolean | null {\n\tconst commands = new Map<string, string>();\n\tlet passed: boolean | null = null;\n\tfor (const message of messages) {\n\t\tif (message.role === \"assistant\") {\n\t\t\tfor (const block of message.content) {\n\t\t\t\tif (block.type === \"toolCall\" && typeof block.arguments.command === \"string\") {\n\t\t\t\t\tcommands.set(block.id, block.arguments.command);\n\t\t\t\t}\n\t\t\t}\n\t\t} else if (message.role === \"toolResult\" && message.toolName === \"bash\") {\n\t\t\tconst command = commands.get(message.toolCallId);\n\t\t\tif (command && TEST_COMMAND.test(command)) passed = !message.isError;\n\t\t}\n\t}\n\treturn passed;\n}\n\n/** Outcome, usage and checks of one agent run. */\nexport function summarizeRun(messages: readonly AgentMessage[]): {\n\toutcome: TelemetryRecord[\"result\"][\"outcome\"];\n\ttestsPassed: boolean | null;\n\ttoolFailures: number;\n\ttokens: number;\n\tcost: number;\n\tmodels: string[];\n} {\n\tlet tokens = 0;\n\tlet cost = 0;\n\tlet toolFailures = 0;\n\tconst models = new Set<string>();\n\tlet outcome: TelemetryRecord[\"result\"][\"outcome\"] = \"completed\";\n\tfor (const message of messages) {\n\t\tif (message.role === \"assistant\") {\n\t\t\ttokens += message.usage?.totalTokens ?? 0;\n\t\t\tcost += message.usage?.cost.total ?? 0;\n\t\t\tmodels.add(`${message.provider}/${message.model}`);\n\t\t\toutcome =\n\t\t\t\tmessage.stopReason === \"error\" ? \"error\" : message.stopReason === \"aborted\" ? \"aborted\" : \"completed\";\n\t\t} else if (message.role === \"toolResult\" && message.isError) {\n\t\t\ttoolFailures++;\n\t\t}\n\t}\n\treturn { outcome, testsPassed: lastCheckPassed(messages), toolFailures, tokens, cost, models: [...models] };\n}\n\nexport function truncateRequest(request: string): string {\n\treturn request.length > MAX_REQUEST_CHARS ? `${request.slice(0, MAX_REQUEST_CHARS)}…` : request;\n}\n\n/** Append-only JSONL store. Unreadable lines are skipped. */\nexport class TelemetryStore {\n\treadonly path: string;\n\n\tconstructor(path: string) {\n\t\tthis.path = path;\n\t}\n\n\tappend(record: TelemetryRecord): void {\n\t\tmkdirSync(dirname(this.path), { recursive: true });\n\t\tappendFileSync(this.path, `${JSON.stringify(record)}\\n`, \"utf8\");\n\t}\n\n\t/** The newest `limit` records, oldest first. */\n\tread(limit = 5000): TelemetryRecord[] {\n\t\tif (!existsSync(this.path)) return [];\n\t\tconst records: TelemetryRecord[] = [];\n\t\tfor (const line of readFileSync(this.path, \"utf8\").split(\"\\n\").slice(-limit)) {\n\t\t\tif (!line.trim()) continue;\n\t\t\ttry {\n\t\t\t\trecords.push(JSON.parse(line) as TelemetryRecord);\n\t\t\t} catch {\n\t\t\t\t// A partially written line from a crashed process.\n\t\t\t}\n\t\t}\n\t\treturn records;\n\t}\n}\n\nexport function historyKey(record: Pick<TelemetryRecord, \"classification\">): string {\n\treturn `${record.classification.type}:${Math.round(record.classification.complexity * 4)}`;\n}\n\n/** Aborted runs say nothing about the model, so they are not counted. */\nexport function historyFromTelemetry(records: readonly TelemetryRecord[]): PerformanceHistory {\n\tconst history = new PerformanceHistory();\n\tfor (const record of records) {\n\t\tif (record.result.outcome === \"aborted\") continue;\n\t\thistory.record(historyKey(record), `${record.selected.provider}/${record.selected.model}`, record.result.success);\n\t}\n\treturn history;\n}\n\n/**\n * Training exercises with evidence-based tier labels: for each successful run, the tier that\n * finished it. Only `capability_tier` is labeled; the other answers came from Laya itself and would\n * teach it nothing new. Review the rows before adding them (`/laya-data`).\n */\nexport function datasetFromTelemetry(records: readonly TelemetryRecord[]): Array<{\n\tstate: { request: string };\n\texpected: { capability_tier: CapabilityTier };\n\tsource: \"harness\";\n}> {\n\treturn records\n\t\t.filter((record) => record.result.success && record.request.trim())\n\t\t.map((record) => ({\n\t\t\tstate: { request: record.request },\n\t\t\texpected: { capability_tier: record.selected.tier },\n\t\t\tsource: \"harness\",\n\t\t}));\n}\n"]}

package/dist/extensions/laya/train-script.d.ts

@@ -0,0 +1,8 @@
+/**
+ * Training script for the routing model. It is built into the Laya Docker image as `train.py`, and
+ * `LayaTrainer` runs it in a training container. It is adapted from laya-trainer's
+ * `laya_ml.py`; the recipe and the hash split match, so its test split is the one the shipped model
+ * never saw.
+ */
+export declare const TRAIN_SCRIPT: string;
+//# sourceMappingURL=train-script.d.ts.map

package/dist/extensions/laya/train-script.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"train-script.d.ts","sourceRoot":"","sources":["../../../src/extensions/laya/train-script.ts"],"names":[],"mappings":"AAAA;;;;;GAKG;AACH,eAAO,MAAM,YAAY,QA8ZxB,CAAC"}

package/dist/extensions/laya/train-script.js

@@ -0,0 +1,422 @@
+/**
+ * Training script for the routing model. It is built into the Laya Docker image as `train.py`, and
+ * `LayaTrainer` runs it in a training container. It is adapted from laya-trainer's
+ * `laya_ml.py`; the recipe and the hash split match, so its test split is the one the shipped model
+ * never saw.
+ */
+export const TRAIN_SCRIPT = String.raw `"""Teaches Relay's Laya routing model the exercises of the training workspace.
+
+Relay writes this script to the Laya home and runs it in the Laya Python environment. The recipe is
+laya-trainer's: soft targets, a policy-gradient term on noisy logits and soft cross-entropy, then
+temperature calibration on the validation split.
+
+Training is incremental. It starts from the model that routes today and mixes the focus exercises
+(the ones labeled from sessions) with a replay sample of the other exercises, so the model learns
+the new tasks without forgetting the rest. Afterwards the new and the current model answer the
+held-out test split, and Relay activates the new model only when it is not worse.
+
+Progress goes to stderr as "PROGRESS <json>" lines. The result is one JSON object on stdout.
+"""
+import os
+
+os.environ.setdefault("USE_TF", "0")
+
+import argparse
+import gc
+import hashlib
+import json
+import math
+import random
+import re
+import sys
+import time
+from pathlib import Path
+
+# Some libraries replace sys.stdout later; writing bytes keeps the output UTF-8 on every platform.
+_OUT = getattr(sys.stdout, "buffer", None)
+_ERR = getattr(sys.stderr, "buffer", None)
+
+
+def _write(buf, fallback, text):
+    if buf is not None:
+        buf.write(text.encode("utf-8"))
+        buf.flush()
+    else:
+        fallback.write(text)
+        fallback.flush()
+
+
+def emit(obj, code=0):
+    _write(_OUT, sys.stdout, json.dumps(obj, ensure_ascii=False) + "\n")
+    sys.exit(code)
+
+
+def fail(error, message, **details):
+    emit({"ok": False, "error": error, "message": message, **details}, 1)
+
+
+def progress(**fields):
+    _write(_ERR, sys.stderr, "PROGRESS " + json.dumps(fields) + "\n")
+
+
+# ---------------------------------------------------------------- exercises
+def norm(s):
+    return re.sub(r"\s+", " ", str(s).strip().lower())
+
+
+def state_text(row):
+    s = row.get("state")
+    return s if isinstance(s, str) else json.dumps(s, ensure_ascii=False, sort_keys=True)
+
+
+def split_of(row):
+    """The row's own split, or laya-trainer's fixed 80/10/10 split by text hash."""
+    if row.get("split") in ("train", "val", "test"):
+        return row["split"]
+    h = int(hashlib.sha1(norm(state_text(row)).encode()).hexdigest(), 16) % 100
+    return "test" if h < 10 else "val" if h < 20 else "train"
+
+
+def read_jsonl(path):
+    rows = []
+    for line in Path(path).read_text(encoding="utf-8").splitlines():
+        if line.strip():
+            try:
+                rows.append(json.loads(line))
+            except json.JSONDecodeError:
+                continue
+    return rows
+
+
+def question(spec):
+    q = {"type": spec["type"], "instructions": spec["instructions"]}
+    if spec["type"] != "noul" or spec.get("criteria"):
+        q["criteria"] = spec["criteria"]
+    return q
+
+
+def target_for(spec, v, smooth=0.02):
+    t = spec["type"]
+    if t == "choice":
+        keys = list(spec["criteria"])
+        return [(1 - smooth) * (k == v) + smooth / len(keys) for k in keys]
+    if t == "noul":
+        y = 1.0 if v else 0.0
+        return [(1 - smooth) * (1 - y) + smooth / 2, (1 - smooth) * y + smooth / 2]
+    n = len(spec["criteria"])  # score: ordinal target, neighbors get some mass
+    p = [0.0] * n
+    p[v] = 0.8
+    near = [i for i in (v - 1, v + 1) if 0 <= i < n]
+    for i in near:
+        p[i] += 0.2 / len(near)
+    return p
+
+
+def build_items(tok, cfg, rows, decs):
+    from laya.common import QTYPES, build_sequence, render_options
+
+    items, skipped = [], 0
+    for r in rows:
+        for qid, v in (r.get("expected") or {}).items():
+            spec = decs.get(qid)
+            if not spec:
+                continue
+            crit = spec.get("criteria", {} if spec["type"] == "choice" else None)
+            q = {"t": spec["type"], "ins": spec["instructions"], "crit": crit}
+            target = target_for(spec, v)
+            seq, markers = build_sequence(tok, r["state"], q, cfg["max_len"], cfg["head_max_len"])
+            if len(markers) != len(render_options({"t": spec["type"], "crit": crit})):
+                skipped += 1
+                continue
+            items.append({"ids": seq, "markers": markers, "qtype": QTYPES[spec["type"]], "target": target,
+                          "label": target.index(max(target))})
+    return items, skipped
+
+
+def collate(items, pad_id):
+    import torch
+
+    b, length = len(items), max(len(i["ids"]) for i in items)
+    k_max = max(len(i["markers"]) for i in items)
+    ids = torch.full((b, length), pad_id, dtype=torch.long)
+    att = torch.zeros((b, length), dtype=torch.long)
+    pos = torch.zeros((b, k_max), dtype=torch.long)
+    mask = torch.zeros((b, k_max), dtype=torch.bool)
+    tgt = torch.zeros((b, k_max))
+    for n, it in enumerate(items):
+        ids[n, :len(it["ids"])] = torch.tensor(it["ids"])
+        att[n, :len(it["ids"])] = 1
+        k = len(it["markers"])
+        pos[n, :k] = torch.tensor(it["markers"])
+        mask[n, :k] = True
+        tgt[n, :len(it["target"])] = torch.tensor(it["target"])
+    return ids, att, pos, mask, tgt, torch.tensor([i["qtype"] for i in items])
+
+
+def fit_temperature(samples):
+    import torch
+    from laya.common import TEMP_MAX, TEMP_MIN
+
+    k_max = max(len(l) for l, _ in samples)
+    lg = torch.full((len(samples), k_max), -1e4)
+    tg = torch.zeros((len(samples), k_max))
+    for i, (l, t) in enumerate(samples):
+        lg[i, :len(l)] = torch.as_tensor(l)
+        tg[i, :len(t)] = torch.as_tensor(t, dtype=torch.float32)
+    lt = torch.zeros(1, requires_grad=True)
+    opt = torch.optim.LBFGS([lt], lr=0.1, max_iter=100)
+
+    def closure():
+        opt.zero_grad()
+        loss = -(tg * torch.log_softmax(lg / lt.exp(), -1)).sum(-1).mean()
+        loss.backward()
+        return loss
+
+    opt.step(closure)
+    return float(torch.clamp(lt.exp(), TEMP_MIN, TEMP_MAX).item())
+
+
+def save_model(model, tok, cfg, path, meta):
+    from safetensors.torch import save_file
+
+    path.mkdir(parents=True, exist_ok=True)
+    save_file({k: v.detach().half().cpu().contiguous() for k, v in model.state_dict().items()},
+              str(path / "model.safetensors"))
+    model.encoder.config.save_pretrained(path / "encoder")
+    tok.save_pretrained(path / "tokenizer")
+    (path / "rl_agent_config.json").write_text(json.dumps(cfg, indent=2), encoding="utf-8")
+    (path / "laya_trainer_meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def pick_device(requested):
+    import torch
+
+    if requested and requested != "auto":
+        return torch.device(requested)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+# ---------------------------------------------------------------- scoring
+def predicted_label(spec, ans):
+    if spec["type"] == "choice":
+        return ans["choice"]
+    if spec["type"] == "noul":
+        return ans["noul"] >= 0.5
+    probs = ans["probabilities"]
+    return int(max(probs, key=lambda k: probs[k]))
+
+
+def score(agent, rows, decs):
+    """Correct answers per question, over the rows that label it."""
+    per = {}
+    for qid, spec in decs.items():
+        sel = [r for r in rows if qid in (r.get("expected") or {})]
+        if not sel:
+            continue
+        outs = agent.predict_batch([r["state"] for r in sel], {qid: question(spec)}, batch_size=16)
+        correct = sum(1 for r, o in zip(sel, outs) if predicted_label(spec, o["answers"][qid]) == r["expected"][qid])
+        per[qid] = {"n": len(sel), "correct": correct}
+    return {"n": sum(m["n"] for m in per.values()), "correct": sum(m["correct"] for m in per.values()),
+            "per_question": per}
+
+
+# ---------------------------------------------------------------- learn
+def cmd_learn(a):
+    import torch
+    from safetensors.torch import load_file
+    from transformers import AutoTokenizer
+    from laya.common import build_model, proper_reward
+
+    t0 = time.time()
+    random.seed(a.seed)
+    torch.manual_seed(a.seed)
+    workspace = Path(a.workspace)
+    decs = json.loads((workspace / "decisions.json").read_text(encoding="utf-8"))["decisions"]
+    rows = [r for r in read_jsonl(workspace / "data" / "dataset.jsonl") if r.get("expected")]
+    focus = [r for r in rows if r.get("source") == a.focus_source and split_of(r) == "train"]
+    rest = [r for r in rows if r.get("source") != a.focus_source]
+    if not focus:
+        fail("no_focus", "There are no '%s' exercises to learn." % a.focus_source)
+    pool = [r for r in rest if split_of(r) == "train"]
+    val = [r for r in rest if split_of(r) == "val"][:a.val_rows]
+    test = [r for r in rows if split_of(r) == "test"]
+    replay_n = min(len(pool), max(a.replay_min, min(a.replay_max, a.replay_ratio * len(focus))))
+    replay = random.Random(a.seed).sample(pool, replay_n)
+
+    init = Path(a.init)
+    out = Path(a.out)
+    device = pick_device(a.device)
+    low_memory = a.low_memory == "on" or (
+        a.low_memory == "auto" and device.type == "cuda"
+        and torch.cuda.get_device_properties(device).total_memory < 12 * 1024 ** 3)
+    cfg = json.loads((init / "rl_agent_config.json").read_text(encoding="utf-8"))
+    cfg["gradient_checkpointing"] = True
+    tok = AutoTokenizer.from_pretrained(str(init / "tokenizer"))
+    train_items, skipped = build_items(tok, cfg, focus * a.repeat_focus + replay, decs)
+    val_items, _ = build_items(tok, cfg, val, decs)
+    if not train_items:
+        fail("no_items", "No exercise fits the model input.")
+
+    model = build_model(cfg, encoder_dir=str(init / "encoder"))
+    model.load_state_dict(load_file(str(init / "model.safetensors")), strict=True)
+    model.float()
+    model.encoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    model.head_checkpointing = True
+    if low_memory:  # freeze the embeddings and the lower half of the encoder: less optimizer memory
+        layers = max((int(m.group(1)) for n, _ in model.named_parameters()
+                      if (m := re.search(r"layers\.(\d+)\.", n)) and "encoder." in n), default=-1) + 1
+        for n, p in model.named_parameters():
+            m = re.search(r"layers\.(\d+)\.", n)
+            if "encoder." in n and ("embeddings" in n or (m and int(m.group(1)) < layers // 2)):
+                p.requires_grad_(False)
+    model.to(device).train()
+    use_amp = device.type == "cuda" and torch.cuda.is_bf16_supported()
+
+    enc = [p for n, p in model.named_parameters() if "encoder." in n and p.requires_grad]
+    head = [p for n, p in model.named_parameters() if "encoder." not in n]
+    opt = torch.optim.AdamW([{"params": enc, "lr": a.lr_encoder}, {"params": head, "lr": a.lr_head}],
+                            weight_decay=0.01)
+    updates = max(1, math.ceil(len(train_items) / a.micro_batch / a.grad_accum) * a.epochs)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=updates, eta_min=1e-6)
+
+    def forward(chunk):
+        ids, att, pos, mask, tgt, qt = [t.to(device) for t in collate(chunk, tok.pad_token_id)]
+        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=use_amp):
+            logits, act = model(ids, att, pos, mask, qt)
+        return logits.float(), act, mask, tgt, qt
+
+    def evaluate(items):
+        model.eval()
+        loss = correct = 0
+        with torch.no_grad():
+            for s in range(0, len(items), a.micro_batch):
+                chunk = items[s:s + a.micro_batch]
+                lg, _, mask, tgt, _ = forward(chunk)
+                lp = torch.log_softmax(lg.masked_fill(~mask, -1e4), -1)
+                loss += -(tgt * lp).sum(-1).sum().item()
+                correct += (lp.argmax(-1).cpu() == torch.tensor([i["label"] for i in chunk])).sum().item()
+        model.train()
+        return loss / max(1, len(items)), correct / max(1, len(items))
+
+    steps = math.ceil(len(train_items) / a.micro_batch)
+    every = max(1, steps // 10)
+    history = []
+    progress(phase="train", epoch=0, epochs=a.epochs, done=0, device=device.type, low_memory=low_memory,
+             items=len(train_items))
+    try:
+        for ep in range(a.epochs):
+            random.Random(a.seed + ep).shuffle(train_items)
+            opt.zero_grad(set_to_none=True)
+            sigma = 0.4 + (0.1 - 0.4) * ep / max(1, a.epochs - 1)
+            total = batches = 0
+            started = time.time()
+            for s in range(0, len(train_items), a.micro_batch):
+                lg, act, mask, tgt, qt = forward(train_items[s:s + a.micro_batch])
+                k = mask.sum(-1, keepdim=True).float()
+                eps = torch.randn((4,) + lg.shape, device=device) * sigma * mask
+                eps = (eps - eps.sum(-1, keepdim=True) / k) * mask
+                noisy = lg.detach().unsqueeze(0) + eps
+                probs = torch.softmax(noisy.masked_fill(~mask, -1e4), -1)
+                with torch.no_grad():
+                    reward = proper_reward(probs, tgt.unsqueeze(0), qt, mask, w_sph=0.75, w_rps=1.0)
+                    adv = reward - reward.mean(0, keepdim=True)
+                    adv = adv / (adv.std() + 1e-6)
+                logp = -(((noisy - lg.unsqueeze(0)) ** 2) * mask).sum(-1) / (2 * sigma ** 2)
+                loss_rl = -(adv * logp).mean()
+                loss_ce = -(tgt * torch.log_softmax(lg.masked_fill(~mask, -1e4), -1)).sum(-1).mean()
+                loss = (loss_rl + loss_ce + 0.0 * act.sum()) / a.grad_accum
+                loss.backward()
+                batches += 1
+                total += loss.item() * a.grad_accum
+                if batches % a.grad_accum == 0 or s + a.micro_batch >= len(train_items):
+                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                    opt.step()
+                    sched.step()
+                    opt.zero_grad(set_to_none=True)
+                if batches % every == 0 and batches < steps:
+                    done = batches / steps
+                    elapsed = time.time() - started
+                    progress(phase="train", epoch=ep + 1, epochs=a.epochs, done=round(done, 3),
+                             eta_s=round(elapsed / done * (1 - done) + (a.epochs - ep - 1) * elapsed / done))
+            val_loss, val_acc = evaluate(val_items) if val_items else (total / batches, None)
+            history.append({"epoch": ep + 1, "train_loss": round(total / batches, 4), "val_loss": round(val_loss, 4),
+                            "val_acc": None if val_acc is None else round(val_acc, 4)})
+            progress(phase="train", epoch=ep + 1, epochs=a.epochs, done=1, val_acc=history[-1]["val_acc"])
+    except torch.cuda.OutOfMemoryError:
+        fail("out_of_memory", "The GPU ran out of memory while training.")
+
+    model.eval()
+    samples = [[], [], []]
+    with torch.no_grad():
+        for s in range(0, len(val_items), a.micro_batch):
+            chunk = val_items[s:s + a.micro_batch]
+            lg, _, _, _, _ = forward(chunk)
+            for i, it in enumerate(chunk):
+                samples[it["qtype"]].append((lg[i, :len(it["markers"])].cpu(), it["target"]))
+    previous = cfg.get("temperature") or [1.0, 1.0, 1.0]
+    cfg["temperature"] = [fit_temperature(g) if len(g) >= 10 else previous[i] for i, g in enumerate(samples)]
+    cfg.pop("temperature_by_options", None)
+    cfg["fine_tuned"] = True
+    meta = {"version": out.name, "init": str(init), "epochs": a.epochs, "history": history,
+            "temperatures": cfg["temperature"],
+            "rows": {"focus": len(focus), "replay": len(replay), "val": len(val), "test": len(test)},
+            "items": {"train": len(train_items), "val": len(val_items), "skipped": skipped},
+            "device": device.type, "low_memory": low_memory, "created": time.strftime("%Y-%m-%d %H:%M:%S")}
+    save_model(model, tok, cfg, out, meta)
+    del model, opt, sched
+    gc.collect()
+    if device.type == "cuda":
+        torch.cuda.empty_cache()
+
+    import laya
+
+    progress(phase="evaluate", model="new")
+    agent = laya.load(str(out), device=device.type)
+    candidate = score(agent, test, decs)
+    learned = score(agent, focus, decs)
+    del agent
+    gc.collect()
+    progress(phase="evaluate", model="current")
+    agent = laya.load(str(init), device=device.type)
+    current = score(agent, test, decs)
+    before = score(agent, focus, decs)
+    emit({"ok": True, "model": str(out), "seconds": round(time.time() - t0), "device": device.type,
+          "low_memory": low_memory, "rows": meta["rows"], "history": history,
+          "test": {"candidate": candidate, "current": current},
+          "focus": {"candidate": learned, "current": before}})
+
+
+def main():
+    ap = argparse.ArgumentParser(description="Teach Relay's Laya routing model")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    p = sub.add_parser("learn")
+    p.add_argument("--workspace", required=True, help="folder with decisions.json and data/dataset.jsonl")
+    p.add_argument("--init", required=True, help="model to start from")
+    p.add_argument("--out", required=True, help="folder of the new model")
+    p.add_argument("--focus-source", default="session")
+    p.add_argument("--repeat-focus", type=int, default=4)
+    p.add_argument("--replay-min", type=int, default=100)
+    p.add_argument("--replay-ratio", type=int, default=4)
+    p.add_argument("--replay-max", type=int, default=600)
+    p.add_argument("--val-rows", type=int, default=100)
+    p.add_argument("--epochs", type=int, default=3)
+    p.add_argument("--micro-batch", type=int, default=2)
+    p.add_argument("--grad-accum", type=int, default=8)
+    p.add_argument("--lr-encoder", type=float, default=2e-5)
+    p.add_argument("--lr-head", type=float, default=1e-4)
+    p.add_argument("--low-memory", default="auto", choices=["auto", "on", "off"])
+    p.add_argument("--device", default="auto")
+    p.add_argument("--seed", type=int, default=7)
+    p.set_defaults(f=cmd_learn)
+    a = ap.parse_args()
+    a.f(a)
+
+
+if __name__ == "__main__":
+    main()
+`;
+//# sourceMappingURL=train-script.js.map

package/dist/extensions/laya/train-script.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"train-script.js","sourceRoot":"","sources":["../../../src/extensions/laya/train-script.ts"],"names":[],"mappings":"AAAA;;;;;GAKG;AACH,MAAM,CAAC,MAAM,YAAY,GAAG,MAAM,CAAC,GAAG,CAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CA8ZrC,CAAC","sourcesContent":["/**\n * Training script for the routing model. It is built into the Laya Docker image as `train.py`, and\n * `LayaTrainer` runs it in a training container. It is adapted from laya-trainer's\n * `laya_ml.py`; the recipe and the hash split match, so its test split is the one the shipped model\n * never saw.\n */\nexport const TRAIN_SCRIPT = String.raw`\"\"\"Teaches Relay's Laya routing model the exercises of the training workspace.\n\nRelay writes this script to the Laya home and runs it in the Laya Python environment. The recipe is\nlaya-trainer's: soft targets, a policy-gradient term on noisy logits and soft cross-entropy, then\ntemperature calibration on the validation split.\n\nTraining is incremental. It starts from the model that routes today and mixes the focus exercises\n(the ones labeled from sessions) with a replay sample of the other exercises, so the model learns\nthe new tasks without forgetting the rest. Afterwards the new and the current model answer the\nheld-out test split, and Relay activates the new model only when it is not worse.\n\nProgress goes to stderr as \"PROGRESS <json>\" lines. The result is one JSON object on stdout.\n\"\"\"\nimport os\n\nos.environ.setdefault(\"USE_TF\", \"0\")\n\nimport argparse\nimport gc\nimport hashlib\nimport json\nimport math\nimport random\nimport re\nimport sys\nimport time\nfrom pathlib import Path\n\n# Some libraries replace sys.stdout later; writing bytes keeps the output UTF-8 on every platform.\n_OUT = getattr(sys.stdout, \"buffer\", None)\n_ERR = getattr(sys.stderr, \"buffer\", None)\n\n\ndef _write(buf, fallback, text):\n    if buf is not None:\n        buf.write(text.encode(\"utf-8\"))\n        buf.flush()\n    else:\n        fallback.write(text)\n        fallback.flush()\n\n\ndef emit(obj, code=0):\n    _write(_OUT, sys.stdout, json.dumps(obj, ensure_ascii=False) + \"\\n\")\n    sys.exit(code)\n\n\ndef fail(error, message, **details):\n    emit({\"ok\": False, \"error\": error, \"message\": message, **details}, 1)\n\n\ndef progress(**fields):\n    _write(_ERR, sys.stderr, \"PROGRESS \" + json.dumps(fields) + \"\\n\")\n\n\n# ---------------------------------------------------------------- exercises\ndef norm(s):\n    return re.sub(r\"\\s+\", \" \", str(s).strip().lower())\n\n\ndef state_text(row):\n    s = row.get(\"state\")\n    return s if isinstance(s, str) else json.dumps(s, ensure_ascii=False, sort_keys=True)\n\n\ndef split_of(row):\n    \"\"\"The row's own split, or laya-trainer's fixed 80/10/10 split by text hash.\"\"\"\n    if row.get(\"split\") in (\"train\", \"val\", \"test\"):\n        return row[\"split\"]\n    h = int(hashlib.sha1(norm(state_text(row)).encode()).hexdigest(), 16) % 100\n    return \"test\" if h < 10 else \"val\" if h < 20 else \"train\"\n\n\ndef read_jsonl(path):\n    rows = []\n    for line in Path(path).read_text(encoding=\"utf-8\").splitlines():\n        if line.strip():\n            try:\n                rows.append(json.loads(line))\n            except json.JSONDecodeError:\n                continue\n    return rows\n\n\ndef question(spec):\n    q = {\"type\": spec[\"type\"], \"instructions\": spec[\"instructions\"]}\n    if spec[\"type\"] != \"noul\" or spec.get(\"criteria\"):\n        q[\"criteria\"] = spec[\"criteria\"]\n    return q\n\n\ndef target_for(spec, v, smooth=0.02):\n    t = spec[\"type\"]\n    if t == \"choice\":\n        keys = list(spec[\"criteria\"])\n        return [(1 - smooth) * (k == v) + smooth / len(keys) for k in keys]\n    if t == \"noul\":\n        y = 1.0 if v else 0.0\n        return [(1 - smooth) * (1 - y) + smooth / 2, (1 - smooth) * y + smooth / 2]\n    n = len(spec[\"criteria\"])  # score: ordinal target, neighbors get some mass\n    p = [0.0] * n\n    p[v] = 0.8\n    near = [i for i in (v - 1, v + 1) if 0 <= i < n]\n    for i in near:\n        p[i] += 0.2 / len(near)\n    return p\n\n\ndef build_items(tok, cfg, rows, decs):\n    from laya.common import QTYPES, build_sequence, render_options\n\n    items, skipped = [], 0\n    for r in rows:\n        for qid, v in (r.get(\"expected\") or {}).items():\n            spec = decs.get(qid)\n            if not spec:\n                continue\n            crit = spec.get(\"criteria\", {} if spec[\"type\"] == \"choice\" else None)\n            q = {\"t\": spec[\"type\"], \"ins\": spec[\"instructions\"], \"crit\": crit}\n            target = target_for(spec, v)\n            seq, markers = build_sequence(tok, r[\"state\"], q, cfg[\"max_len\"], cfg[\"head_max_len\"])\n            if len(markers) != len(render_options({\"t\": spec[\"type\"], \"crit\": crit})):\n                skipped += 1\n                continue\n            items.append({\"ids\": seq, \"markers\": markers, \"qtype\": QTYPES[spec[\"type\"]], \"target\": target,\n                          \"label\": target.index(max(target))})\n    return items, skipped\n\n\ndef collate(items, pad_id):\n    import torch\n\n    b, length = len(items), max(len(i[\"ids\"]) for i in items)\n    k_max = max(len(i[\"markers\"]) for i in items)\n    ids = torch.full((b, length), pad_id, dtype=torch.long)\n    att = torch.zeros((b, length), dtype=torch.long)\n    pos = torch.zeros((b, k_max), dtype=torch.long)\n    mask = torch.zeros((b, k_max), dtype=torch.bool)\n    tgt = torch.zeros((b, k_max))\n    for n, it in enumerate(items):\n        ids[n, :len(it[\"ids\"])] = torch.tensor(it[\"ids\"])\n        att[n, :len(it[\"ids\"])] = 1\n        k = len(it[\"markers\"])\n        pos[n, :k] = torch.tensor(it[\"markers\"])\n        mask[n, :k] = True\n        tgt[n, :len(it[\"target\"])] = torch.tensor(it[\"target\"])\n    return ids, att, pos, mask, tgt, torch.tensor([i[\"qtype\"] for i in items])\n\n\ndef fit_temperature(samples):\n    import torch\n    from laya.common import TEMP_MAX, TEMP_MIN\n\n    k_max = max(len(l) for l, _ in samples)\n    lg = torch.full((len(samples), k_max), -1e4)\n    tg = torch.zeros((len(samples), k_max))\n    for i, (l, t) in enumerate(samples):\n        lg[i, :len(l)] = torch.as_tensor(l)\n        tg[i, :len(t)] = torch.as_tensor(t, dtype=torch.float32)\n    lt = torch.zeros(1, requires_grad=True)\n    opt = torch.optim.LBFGS([lt], lr=0.1, max_iter=100)\n\n    def closure():\n        opt.zero_grad()\n        loss = -(tg * torch.log_softmax(lg / lt.exp(), -1)).sum(-1).mean()\n        loss.backward()\n        return loss\n\n    opt.step(closure)\n    return float(torch.clamp(lt.exp(), TEMP_MIN, TEMP_MAX).item())\n\n\ndef save_model(model, tok, cfg, path, meta):\n    from safetensors.torch import save_file\n\n    path.mkdir(parents=True, exist_ok=True)\n    save_file({k: v.detach().half().cpu().contiguous() for k, v in model.state_dict().items()},\n              str(path / \"model.safetensors\"))\n    model.encoder.config.save_pretrained(path / \"encoder\")\n    tok.save_pretrained(path / \"tokenizer\")\n    (path / \"rl_agent_config.json\").write_text(json.dumps(cfg, indent=2), encoding=\"utf-8\")\n    (path / \"laya_trainer_meta.json\").write_text(json.dumps(meta, ensure_ascii=False, indent=2), encoding=\"utf-8\")\n\n\ndef pick_device(requested):\n    import torch\n\n    if requested and requested != \"auto\":\n        return torch.device(requested)\n    if torch.cuda.is_available():\n        return torch.device(\"cuda\")\n    if getattr(torch.backends, \"mps\", None) and torch.backends.mps.is_available():\n        return torch.device(\"mps\")\n    return torch.device(\"cpu\")\n\n\n# ---------------------------------------------------------------- scoring\ndef predicted_label(spec, ans):\n    if spec[\"type\"] == \"choice\":\n        return ans[\"choice\"]\n    if spec[\"type\"] == \"noul\":\n        return ans[\"noul\"] >= 0.5\n    probs = ans[\"probabilities\"]\n    return int(max(probs, key=lambda k: probs[k]))\n\n\ndef score(agent, rows, decs):\n    \"\"\"Correct answers per question, over the rows that label it.\"\"\"\n    per = {}\n    for qid, spec in decs.items():\n        sel = [r for r in rows if qid in (r.get(\"expected\") or {})]\n        if not sel:\n            continue\n        outs = agent.predict_batch([r[\"state\"] for r in sel], {qid: question(spec)}, batch_size=16)\n        correct = sum(1 for r, o in zip(sel, outs) if predicted_label(spec, o[\"answers\"][qid]) == r[\"expected\"][qid])\n        per[qid] = {\"n\": len(sel), \"correct\": correct}\n    return {\"n\": sum(m[\"n\"] for m in per.values()), \"correct\": sum(m[\"correct\"] for m in per.values()),\n            \"per_question\": per}\n\n\n# ---------------------------------------------------------------- learn\ndef cmd_learn(a):\n    import torch\n    from safetensors.torch import load_file\n    from transformers import AutoTokenizer\n    from laya.common import build_model, proper_reward\n\n    t0 = time.time()\n    random.seed(a.seed)\n    torch.manual_seed(a.seed)\n    workspace = Path(a.workspace)\n    decs = json.loads((workspace / \"decisions.json\").read_text(encoding=\"utf-8\"))[\"decisions\"]\n    rows = [r for r in read_jsonl(workspace / \"data\" / \"dataset.jsonl\") if r.get(\"expected\")]\n    focus = [r for r in rows if r.get(\"source\") == a.focus_source and split_of(r) == \"train\"]\n    rest = [r for r in rows if r.get(\"source\") != a.focus_source]\n    if not focus:\n        fail(\"no_focus\", \"There are no '%s' exercises to learn.\" % a.focus_source)\n    pool = [r for r in rest if split_of(r) == \"train\"]\n    val = [r for r in rest if split_of(r) == \"val\"][:a.val_rows]\n    test = [r for r in rows if split_of(r) == \"test\"]\n    replay_n = min(len(pool), max(a.replay_min, min(a.replay_max, a.replay_ratio * len(focus))))\n    replay = random.Random(a.seed).sample(pool, replay_n)\n\n    init = Path(a.init)\n    out = Path(a.out)\n    device = pick_device(a.device)\n    low_memory = a.low_memory == \"on\" or (\n        a.low_memory == \"auto\" and device.type == \"cuda\"\n        and torch.cuda.get_device_properties(device).total_memory < 12 * 1024 ** 3)\n    cfg = json.loads((init / \"rl_agent_config.json\").read_text(encoding=\"utf-8\"))\n    cfg[\"gradient_checkpointing\"] = True\n    tok = AutoTokenizer.from_pretrained(str(init / \"tokenizer\"))\n    train_items, skipped = build_items(tok, cfg, focus * a.repeat_focus + replay, decs)\n    val_items, _ = build_items(tok, cfg, val, decs)\n    if not train_items:\n        fail(\"no_items\", \"No exercise fits the model input.\")\n\n    model = build_model(cfg, encoder_dir=str(init / \"encoder\"))\n    model.load_state_dict(load_file(str(init / \"model.safetensors\")), strict=True)\n    model.float()\n    model.encoder.gradient_checkpointing_enable(gradient_checkpointing_kwargs={\"use_reentrant\": False})\n    model.head_checkpointing = True\n    if low_memory:  # freeze the embeddings and the lower half of the encoder: less optimizer memory\n        layers = max((int(m.group(1)) for n, _ in model.named_parameters()\n                      if (m := re.search(r\"layers\\.(\\d+)\\.\", n)) and \"encoder.\" in n), default=-1) + 1\n        for n, p in model.named_parameters():\n            m = re.search(r\"layers\\.(\\d+)\\.\", n)\n            if \"encoder.\" in n and (\"embeddings\" in n or (m and int(m.group(1)) < layers // 2)):\n                p.requires_grad_(False)\n    model.to(device).train()\n    use_amp = device.type == \"cuda\" and torch.cuda.is_bf16_supported()\n\n    enc = [p for n, p in model.named_parameters() if \"encoder.\" in n and p.requires_grad]\n    head = [p for n, p in model.named_parameters() if \"encoder.\" not in n]\n    opt = torch.optim.AdamW([{\"params\": enc, \"lr\": a.lr_encoder}, {\"params\": head, \"lr\": a.lr_head}],\n                            weight_decay=0.01)\n    updates = max(1, math.ceil(len(train_items) / a.micro_batch / a.grad_accum) * a.epochs)\n    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=updates, eta_min=1e-6)\n\n    def forward(chunk):\n        ids, att, pos, mask, tgt, qt = [t.to(device) for t in collate(chunk, tok.pad_token_id)]\n        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=use_amp):\n            logits, act = model(ids, att, pos, mask, qt)\n        return logits.float(), act, mask, tgt, qt\n\n    def evaluate(items):\n        model.eval()\n        loss = correct = 0\n        with torch.no_grad():\n            for s in range(0, len(items), a.micro_batch):\n                chunk = items[s:s + a.micro_batch]\n                lg, _, mask, tgt, _ = forward(chunk)\n                lp = torch.log_softmax(lg.masked_fill(~mask, -1e4), -1)\n                loss += -(tgt * lp).sum(-1).sum().item()\n                correct += (lp.argmax(-1).cpu() == torch.tensor([i[\"label\"] for i in chunk])).sum().item()\n        model.train()\n        return loss / max(1, len(items)), correct / max(1, len(items))\n\n    steps = math.ceil(len(train_items) / a.micro_batch)\n    every = max(1, steps // 10)\n    history = []\n    progress(phase=\"train\", epoch=0, epochs=a.epochs, done=0, device=device.type, low_memory=low_memory,\n             items=len(train_items))\n    try:\n        for ep in range(a.epochs):\n            random.Random(a.seed + ep).shuffle(train_items)\n            opt.zero_grad(set_to_none=True)\n            sigma = 0.4 + (0.1 - 0.4) * ep / max(1, a.epochs - 1)\n            total = batches = 0\n            started = time.time()\n            for s in range(0, len(train_items), a.micro_batch):\n                lg, act, mask, tgt, qt = forward(train_items[s:s + a.micro_batch])\n                k = mask.sum(-1, keepdim=True).float()\n                eps = torch.randn((4,) + lg.shape, device=device) * sigma * mask\n                eps = (eps - eps.sum(-1, keepdim=True) / k) * mask\n                noisy = lg.detach().unsqueeze(0) + eps\n                probs = torch.softmax(noisy.masked_fill(~mask, -1e4), -1)\n                with torch.no_grad():\n                    reward = proper_reward(probs, tgt.unsqueeze(0), qt, mask, w_sph=0.75, w_rps=1.0)\n                    adv = reward - reward.mean(0, keepdim=True)\n                    adv = adv / (adv.std() + 1e-6)\n                logp = -(((noisy - lg.unsqueeze(0)) ** 2) * mask).sum(-1) / (2 * sigma ** 2)\n                loss_rl = -(adv * logp).mean()\n                loss_ce = -(tgt * torch.log_softmax(lg.masked_fill(~mask, -1e4), -1)).sum(-1).mean()\n                loss = (loss_rl + loss_ce + 0.0 * act.sum()) / a.grad_accum\n                loss.backward()\n                batches += 1\n                total += loss.item() * a.grad_accum\n                if batches % a.grad_accum == 0 or s + a.micro_batch >= len(train_items):\n                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n                    opt.step()\n                    sched.step()\n                    opt.zero_grad(set_to_none=True)\n                if batches % every == 0 and batches < steps:\n                    done = batches / steps\n                    elapsed = time.time() - started\n                    progress(phase=\"train\", epoch=ep + 1, epochs=a.epochs, done=round(done, 3),\n                             eta_s=round(elapsed / done * (1 - done) + (a.epochs - ep - 1) * elapsed / done))\n            val_loss, val_acc = evaluate(val_items) if val_items else (total / batches, None)\n            history.append({\"epoch\": ep + 1, \"train_loss\": round(total / batches, 4), \"val_loss\": round(val_loss, 4),\n                            \"val_acc\": None if val_acc is None else round(val_acc, 4)})\n            progress(phase=\"train\", epoch=ep + 1, epochs=a.epochs, done=1, val_acc=history[-1][\"val_acc\"])\n    except torch.cuda.OutOfMemoryError:\n        fail(\"out_of_memory\", \"The GPU ran out of memory while training.\")\n\n    model.eval()\n    samples = [[], [], []]\n    with torch.no_grad():\n        for s in range(0, len(val_items), a.micro_batch):\n            chunk = val_items[s:s + a.micro_batch]\n            lg, _, _, _, _ = forward(chunk)\n            for i, it in enumerate(chunk):\n                samples[it[\"qtype\"]].append((lg[i, :len(it[\"markers\"])].cpu(), it[\"target\"]))\n    previous = cfg.get(\"temperature\") or [1.0, 1.0, 1.0]\n    cfg[\"temperature\"] = [fit_temperature(g) if len(g) >= 10 else previous[i] for i, g in enumerate(samples)]\n    cfg.pop(\"temperature_by_options\", None)\n    cfg[\"fine_tuned\"] = True\n    meta = {\"version\": out.name, \"init\": str(init), \"epochs\": a.epochs, \"history\": history,\n            \"temperatures\": cfg[\"temperature\"],\n            \"rows\": {\"focus\": len(focus), \"replay\": len(replay), \"val\": len(val), \"test\": len(test)},\n            \"items\": {\"train\": len(train_items), \"val\": len(val_items), \"skipped\": skipped},\n            \"device\": device.type, \"low_memory\": low_memory, \"created\": time.strftime(\"%Y-%m-%d %H:%M:%S\")}\n    save_model(model, tok, cfg, out, meta)\n    del model, opt, sched\n    gc.collect()\n    if device.type == \"cuda\":\n        torch.cuda.empty_cache()\n\n    import laya\n\n    progress(phase=\"evaluate\", model=\"new\")\n    agent = laya.load(str(out), device=device.type)\n    candidate = score(agent, test, decs)\n    learned = score(agent, focus, decs)\n    del agent\n    gc.collect()\n    progress(phase=\"evaluate\", model=\"current\")\n    agent = laya.load(str(init), device=device.type)\n    current = score(agent, test, decs)\n    before = score(agent, focus, decs)\n    emit({\"ok\": True, \"model\": str(out), \"seconds\": round(time.time() - t0), \"device\": device.type,\n          \"low_memory\": low_memory, \"rows\": meta[\"rows\"], \"history\": history,\n          \"test\": {\"candidate\": candidate, \"current\": current},\n          \"focus\": {\"candidate\": learned, \"current\": before}})\n\n\ndef main():\n    ap = argparse.ArgumentParser(description=\"Teach Relay's Laya routing model\")\n    sub = ap.add_subparsers(dest=\"cmd\", required=True)\n    p = sub.add_parser(\"learn\")\n    p.add_argument(\"--workspace\", required=True, help=\"folder with decisions.json and data/dataset.jsonl\")\n    p.add_argument(\"--init\", required=True, help=\"model to start from\")\n    p.add_argument(\"--out\", required=True, help=\"folder of the new model\")\n    p.add_argument(\"--focus-source\", default=\"session\")\n    p.add_argument(\"--repeat-focus\", type=int, default=4)\n    p.add_argument(\"--replay-min\", type=int, default=100)\n    p.add_argument(\"--replay-ratio\", type=int, default=4)\n    p.add_argument(\"--replay-max\", type=int, default=600)\n    p.add_argument(\"--val-rows\", type=int, default=100)\n    p.add_argument(\"--epochs\", type=int, default=3)\n    p.add_argument(\"--micro-batch\", type=int, default=2)\n    p.add_argument(\"--grad-accum\", type=int, default=8)\n    p.add_argument(\"--lr-encoder\", type=float, default=2e-5)\n    p.add_argument(\"--lr-head\", type=float, default=1e-4)\n    p.add_argument(\"--low-memory\", default=\"auto\", choices=[\"auto\", \"on\", \"off\"])\n    p.add_argument(\"--device\", default=\"auto\")\n    p.add_argument(\"--seed\", type=int, default=7)\n    p.set_defaults(f=cmd_learn)\n    a = ap.parse_args()\n    a.f(a)\n\n\nif __name__ == \"__main__\":\n    main()\n`;\n"]}