Compare commits
28
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
cdcd5bb092
|
||
|
|
c7505a9344 | ||
|
|
8ed8d04d3a | ||
|
|
6a0d4fbb0c | ||
|
|
73a376e57c
|
||
|
|
06cf081c47
|
||
|
|
0a122fb26f
|
||
|
|
e3248cc186
|
||
|
|
057cb72806
|
||
|
|
5264f78b2d
|
||
|
|
2cbfd06e40 | ||
|
|
816c8dc96f
|
||
|
|
042d708da3 | ||
|
|
421bf55697
|
||
|
|
e242514121
|
||
|
|
1ba160491c
|
||
|
|
047005203e
|
||
|
|
86a6f25a6b
|
||
|
|
ec315a7f58
|
||
|
|
7c17a9fcd4 | ||
|
|
c6502e1519
|
||
|
|
21a1083308 | ||
|
|
83cd621eee
|
||
|
|
c68d1e225b
|
||
|
|
65dc1103f1 | ||
|
|
ac150fa61c
|
||
|
|
4466d5cbfc
|
||
|
|
ec0dd319e8
|
@@ -525,7 +525,14 @@ jobs:
|
||||
cuda_home: /usr/local/cuda-13.0
|
||||
build_jobs: 8
|
||||
nvcc_threads: 4
|
||||
cargo_features: "cuda cudnn"
|
||||
# flash-attn on ampere (#337): Laya's encoder runs every
|
||||
# attention through flash_attn_varlen_windowed, matching the
|
||||
# PyTorch reference on a 3060 (36 vs 39 ms at 512 tokens) where
|
||||
# the eager path is 1.4-1.7x slower. This also moves the
|
||||
# flavour's other flash-capable paths onto flash -- qwen3_5
|
||||
# full attention and Z-Image -- which validate-neuron.sh and
|
||||
# validate-image.sh gate. NEURON_FLASH_ATTN=0 rolls back.
|
||||
cargo_features: "cuda cudnn flash-attn"
|
||||
- flavour: ada
|
||||
compute_cap: "89"
|
||||
runner: cuda-13.0
|
||||
@@ -540,8 +547,8 @@ jobs:
|
||||
build_jobs: 8
|
||||
nvcc_threads: 4
|
||||
# flash-attn on blackwell first (#95): beast carries the
|
||||
# agentic prefill pain; ada/ampere follow once the win is
|
||||
# measured. NEURON_FLASH_ATTN=0 is the runtime rollback.
|
||||
# agentic prefill pain. NEURON_FLASH_ATTN=0 is the runtime
|
||||
# rollback.
|
||||
cargo_features: "cuda cudnn flash-attn"
|
||||
runs-on: ${{ matrix.runner }}
|
||||
env:
|
||||
|
||||
@@ -244,10 +244,14 @@ jobs:
|
||||
flavour: blackwell
|
||||
load_timeout: 900
|
||||
max_prompt_tokens: 131072
|
||||
- host: benjy.hanzalova.internal
|
||||
flavour: ada
|
||||
load_timeout: 300
|
||||
max_prompt_tokens: 16384
|
||||
# benjy is out of the matrix while its boot NVMe is replaced
|
||||
# (dead 2026-09-29): every deploy failed at SSH init against a
|
||||
# host that is not there. Restore these four lines when it is
|
||||
# back on the mesh.
|
||||
# - host: benjy.hanzalova.internal
|
||||
# flavour: ada
|
||||
# load_timeout: 300
|
||||
# max_prompt_tokens: 16384
|
||||
- host: quadbrat.hanzalova.internal
|
||||
flavour: ampere
|
||||
load_timeout: 300
|
||||
|
||||
Generated
+2
@@ -3166,6 +3166,8 @@ dependencies = [
|
||||
"toml",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
"unicode-normalization",
|
||||
"unicode-properties",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -7,9 +7,10 @@
|
||||
# ─────────────────────────────────────────────────────────────────────
|
||||
# neuron.toml for quadbrat.hanzalova.internal
|
||||
#
|
||||
# 1x RTX 3060 (12 GB) — small / quantised tier. Pre-warms Qwen3-1.7B
|
||||
# (bf16, ~4 GB), leaving ~7 GB for KV cache so long contexts on a small
|
||||
# model still have plenty of room.
|
||||
# 1x RTX 3060 (12 GB). Resident: Laya, the System-One decision model
|
||||
# (#333, ~2.3 GB bf16 for its three checkpoints). It replaced the
|
||||
# Qwen3-1.7B pre-warm (operator decision 2026-09-29, #341); the 1.7B is
|
||||
# still cold-loadable beside it for helexa/small.
|
||||
#
|
||||
# Synced to /etc/neuron/neuron.toml by script/infra-setup.sh.
|
||||
|
||||
@@ -21,10 +22,21 @@ name = "candle"
|
||||
[harness.candle]
|
||||
|
||||
[[default_models]]
|
||||
model_id = "Qwen/Qwen3-1.7B"
|
||||
model_id = "convaiinnovations/laya"
|
||||
harness = "candle"
|
||||
devices = [0]
|
||||
|
||||
# A decision request is one forward of tens of milliseconds (36 ms for
|
||||
# one question over a 512-token state, 162 ms for five), so a one-slot
|
||||
# queue of 8 rejects ordinary bursts. Admit
|
||||
# several at once -- the device worker still runs forwards one at a
|
||||
# time -- and let a deep, short-deadline queue absorb the rest. Applies
|
||||
# to every model on this host, including a cold-loaded helexa/small.
|
||||
[harness.candle.admission]
|
||||
max_in_flight = 4
|
||||
max_queue_depth = 64
|
||||
max_wait_secs = 10
|
||||
|
||||
# Text-to-image (#204): the 12 GB tier serves the q8_0-quantized DiT
|
||||
# (8.7 GB resident incl. VAE). Ceiling is 768² — naive f32 attention
|
||||
# transients OOM 1024² beside the resident model; lifts to 1024² if
|
||||
|
||||
@@ -107,6 +107,27 @@ pub const DEFAULT_RESIDENCY_PRIORITY: u32 = 100;
|
||||
/// grant on its own.
|
||||
pub const PINNED_RESIDENCY_PRIORITY: u32 = 1000;
|
||||
|
||||
/// The alias a `/v1/systemone` decision request routes to when its
|
||||
/// `model` is absent or names nothing this fleet serves.
|
||||
///
|
||||
/// Decision clients written against TypeSafe's Jev API send a `model`
|
||||
/// naming *their* service (e.g. `jev-1`), and Laya's own server treats
|
||||
/// any unrecognised value as "let the router choose". Every hop of the
|
||||
/// serving chain applies the same rule, scoped to the decision endpoint
|
||||
/// only — a chat request for an unknown model is still a 404. The
|
||||
/// operator points the alias at a decision model in `[aliases]`.
|
||||
pub const DECISION_DEFAULT_ALIAS: &str = "helexa/one";
|
||||
|
||||
/// Header carrying the concrete decision model a request was routed to.
|
||||
///
|
||||
/// A decision request's `model` field can name a checkpoint *inside* a
|
||||
/// served model family (`multilingual`, `typed-decisions`) or a foreign
|
||||
/// service, so the gateway forwards the body untouched and states the
|
||||
/// routing outcome here instead of rewriting `model`. The serving node
|
||||
/// uses it to pick the model and still reads `model` for a checkpoint
|
||||
/// pin.
|
||||
pub const DECISION_MODEL_HEADER: &str = "x-helexa-decision-model";
|
||||
|
||||
/// The full model catalogue.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||
pub struct ModelCatalogue {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -2,6 +2,7 @@ pub mod anthropic;
|
||||
pub mod build_info;
|
||||
pub mod catalogue;
|
||||
pub mod config;
|
||||
pub mod decisions;
|
||||
pub mod discovery;
|
||||
pub mod entitlements;
|
||||
pub mod error_envelope;
|
||||
|
||||
@@ -24,6 +24,7 @@ pub fn api_routes() -> Router<Arc<CortexState>> {
|
||||
.route("/v1/completions", post(completions))
|
||||
.route("/v1/responses", post(responses))
|
||||
.route("/v1/images/generations", post(images_generations))
|
||||
.route("/v1/systemone", post(systemone))
|
||||
.route("/v1/models", get(list_models))
|
||||
.route("/v1/messages", post(anthropic_messages))
|
||||
.route("/health", get(health))
|
||||
@@ -349,6 +350,192 @@ async fn images_generations(
|
||||
})
|
||||
}
|
||||
|
||||
/// `POST /v1/systemone` — typed decisions (TypeSafe Jev wire protocol)
|
||||
/// from a decision model such as Laya (#338).
|
||||
///
|
||||
/// Routing differs from the other inference paths in one way: a
|
||||
/// missing or unknown `model` resolves to
|
||||
/// [`DECISION_DEFAULT_ALIAS`](cortex_core::catalogue::DECISION_DEFAULT_ALIAS)
|
||||
/// instead of 404ing, because Jev clients name their own service there.
|
||||
/// A `model` that *does* resolve — an alias, or a model this fleet
|
||||
/// knows — routes as usual, so asking a text model for a decision
|
||||
/// reaches neuron and gets its `wrong_modality` answer rather than a
|
||||
/// silent substitution.
|
||||
///
|
||||
/// The body is forwarded untouched (its `model` may pin a checkpoint
|
||||
/// inside the served family); the routed model travels in
|
||||
/// [`DECISION_MODEL_HEADER`](cortex_core::catalogue::DECISION_MODEL_HEADER).
|
||||
///
|
||||
/// Metered like Jev: input tokens only. The reservation is the exact
|
||||
/// upper bound — questions × the sequence cap — and settles to
|
||||
/// `usage.input_tokens` from the response.
|
||||
async fn systemone(
|
||||
State(fleet): State<Arc<CortexState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
use cortex_core::catalogue::{DECISION_DEFAULT_ALIAS, DECISION_MODEL_HEADER};
|
||||
|
||||
log_inbound("systemone", "/v1/systemone", &body);
|
||||
let requested = extract_model(&body);
|
||||
let route = match requested.as_deref() {
|
||||
Some(model) => match router::resolve(&fleet, model).await {
|
||||
Err(router::RouteError::ModelNotFound(_)) => {
|
||||
tracing::debug!(
|
||||
requested = model,
|
||||
default = DECISION_DEFAULT_ALIAS,
|
||||
"unknown decision model; using the default"
|
||||
);
|
||||
router::resolve(&fleet, DECISION_DEFAULT_ALIAS).await
|
||||
}
|
||||
other => other,
|
||||
},
|
||||
None => router::resolve(&fleet, DECISION_DEFAULT_ALIAS).await,
|
||||
};
|
||||
let route = match route {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
handler = "systemone",
|
||||
requested = requested.as_deref().unwrap_or("<none>"),
|
||||
error = %e,
|
||||
"route resolve failed"
|
||||
);
|
||||
return route_error_response(&e);
|
||||
}
|
||||
};
|
||||
|
||||
touch_model(&fleet, &route.node_name, &route.resolved_model_id).await;
|
||||
|
||||
let labels = [
|
||||
("model", route.resolved_model_id.clone()),
|
||||
("node", route.node_name.clone()),
|
||||
];
|
||||
metrics::counter!("cortex_requests_total", &labels).increment(1);
|
||||
metrics::counter!("cortex_decision_requests_total", &labels).increment(1);
|
||||
if route.cold_start {
|
||||
metrics::counter!("cortex_cold_starts_total", &labels).increment(1);
|
||||
}
|
||||
|
||||
let principal = crate::metering::principal_from_headers(&headers);
|
||||
let bound = crate::metering::decision_reservation_tokens(&body);
|
||||
let guard = match &principal {
|
||||
Some(principal) => {
|
||||
match crate::metering::reserve_or_reject(
|
||||
Arc::clone(&fleet.entitlements),
|
||||
principal,
|
||||
bound,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(guard) => Some(guard),
|
||||
Err(env) => return crate::error::envelope_response(env),
|
||||
}
|
||||
}
|
||||
None => None,
|
||||
};
|
||||
|
||||
let mut forward_headers = strip_hop_headers(&headers);
|
||||
match axum::http::HeaderValue::from_str(&route.resolved_model_id) {
|
||||
Ok(v) => {
|
||||
forward_headers.insert(DECISION_MODEL_HEADER, v);
|
||||
}
|
||||
Err(_) => {
|
||||
// A model id that is not a valid header value cannot have
|
||||
// come from the catalogue or a neuron's model list.
|
||||
return error_response(500, "api_error", "invalid_model_id", "internal error");
|
||||
}
|
||||
}
|
||||
|
||||
let url = format!("{}/v1/systemone", route.endpoint);
|
||||
let start = Instant::now();
|
||||
let upstream = fleet
|
||||
.http_client
|
||||
.post(&url)
|
||||
.headers(forward_headers)
|
||||
.header("content-type", "application/json")
|
||||
.body(body)
|
||||
.send()
|
||||
.await;
|
||||
let upstream = match upstream {
|
||||
Ok(r) => r,
|
||||
Err(e) => {
|
||||
metrics::counter!("cortex_request_errors_total", &labels).increment(1);
|
||||
tracing::warn!(node = %route.node_name, url = %url, error = %e, "systemone proxy failed");
|
||||
return error_response(
|
||||
502,
|
||||
"api_error",
|
||||
"upstream_unreachable",
|
||||
"failed to reach the serving node",
|
||||
);
|
||||
}
|
||||
};
|
||||
let status = upstream.status();
|
||||
let passthrough: Vec<(axum::http::HeaderName, axum::http::HeaderValue)> = upstream
|
||||
.headers()
|
||||
.iter()
|
||||
.filter(|(n, _)| {
|
||||
matches!(
|
||||
n.as_str(),
|
||||
"server-timing" | "x-inference-time-ms" | "retry-after"
|
||||
)
|
||||
})
|
||||
.map(|(n, v)| (n.clone(), v.clone()))
|
||||
.collect();
|
||||
let resp_bytes = match upstream.bytes().await {
|
||||
Ok(b) => b,
|
||||
Err(e) => {
|
||||
metrics::counter!("cortex_request_errors_total", &labels).increment(1);
|
||||
tracing::warn!(node = %route.node_name, error = %e, "systemone response read failed");
|
||||
return error_response(
|
||||
502,
|
||||
"api_error",
|
||||
"upstream_read_failed",
|
||||
"failed to read the serving node's response",
|
||||
);
|
||||
}
|
||||
};
|
||||
metrics::histogram!("cortex_request_duration_seconds", &labels)
|
||||
.record(start.elapsed().as_secs_f64());
|
||||
|
||||
if status.is_success() {
|
||||
let input_tokens = serde_json::from_slice::<Value>(&resp_bytes)
|
||||
.ok()
|
||||
.and_then(|v| v.get("usage")?.get("input_tokens")?.as_u64());
|
||||
if let Some(n) = input_tokens {
|
||||
metrics::counter!("cortex_decision_input_tokens_total", &labels).increment(n);
|
||||
}
|
||||
if let (Some(principal), Some(guard)) = (principal, guard) {
|
||||
// No usage block means the node answered outside the
|
||||
// contract; settle at the reserved bound rather than free.
|
||||
let spent = input_tokens.unwrap_or_else(|| {
|
||||
tracing::warn!(node = %route.node_name, "systemone response carried no usage");
|
||||
bound
|
||||
});
|
||||
crate::metering::record_spend(&principal, spent, 0);
|
||||
fleet
|
||||
.served_usage
|
||||
.add(&principal.account_id, &principal.key_id, spent);
|
||||
guard.settle(spent);
|
||||
}
|
||||
} else {
|
||||
metrics::counter!("cortex_request_errors_total", &labels).increment(1);
|
||||
// guard drops here -> reservation released, no spend recorded.
|
||||
}
|
||||
|
||||
let mut response = axum::response::Response::builder()
|
||||
.status(status)
|
||||
.header(axum::http::header::CONTENT_TYPE, "application/json");
|
||||
for (name, value) in passthrough {
|
||||
response = response.header(name, value);
|
||||
}
|
||||
response
|
||||
.body(axum::body::Body::from(resp_bytes))
|
||||
.unwrap_or_else(|_| {
|
||||
error_response(500, "api_error", "response_build_failed", "internal error")
|
||||
})
|
||||
}
|
||||
|
||||
/// Forward client headers minus hop-by-hop and host headers, keeping
|
||||
/// the cortex-stamped principal headers intact for neuron's fair-share
|
||||
/// admission (#54).
|
||||
|
||||
@@ -239,6 +239,35 @@ fn budget_error_to_envelope(err: BudgetError) -> OpenAiError {
|
||||
}
|
||||
}
|
||||
|
||||
/// Sequence cap assumed for a decision request that does not set
|
||||
/// `max_len`: the largest default any shipped Laya checkpoint uses.
|
||||
const DECISION_DEFAULT_MAX_LEN: u64 = 1024;
|
||||
|
||||
/// The ceiling a decision request may raise `max_len` to (Laya's own
|
||||
/// server caps the budget at the same value).
|
||||
const DECISION_MAX_LEN_CAP: u64 = 8192;
|
||||
|
||||
/// Upper-bound input tokens for a `/v1/systemone` request (#338/#340).
|
||||
///
|
||||
/// A decision model encodes one sequence per question, each at most
|
||||
/// `max_len` tokens, and bills input tokens only — so questions ×
|
||||
/// `max_len` is a true ceiling, not an estimate. The node reports the
|
||||
/// exact count in `usage.input_tokens` and settle corrects to it.
|
||||
pub fn decision_reservation_tokens(body: &[u8]) -> u64 {
|
||||
let v: serde_json::Value = serde_json::from_slice(body).unwrap_or_default();
|
||||
let questions = match v.get("questions") {
|
||||
Some(serde_json::Value::Object(m)) => m.len() as u64,
|
||||
Some(serde_json::Value::Array(a)) => a.len() as u64,
|
||||
_ => 1,
|
||||
};
|
||||
let max_len = v
|
||||
.get("max_len")
|
||||
.and_then(serde_json::Value::as_u64)
|
||||
.unwrap_or(DECISION_DEFAULT_MAX_LEN)
|
||||
.min(DECISION_MAX_LEN_CAP);
|
||||
questions.max(1).saturating_mul(max_len.max(1))
|
||||
}
|
||||
|
||||
/// Upper-bound tokens to reserve for a request (#52): an over-estimate of
|
||||
/// the prompt plus the maximum output. `advertised_output` is the model's
|
||||
/// `limit.output` (#62), used when the request omits `max_(completion_)tokens`.
|
||||
@@ -271,6 +300,33 @@ fn estimate_prompt_tokens(body: &[u8]) -> u64 {
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn decision_reservation_is_questions_times_sequence_cap() {
|
||||
let body = br#"{"state":"x","questions":{"a":{},"b":{},"c":{}}}"#;
|
||||
assert_eq!(
|
||||
decision_reservation_tokens(body),
|
||||
3 * DECISION_DEFAULT_MAX_LEN
|
||||
);
|
||||
let body = br#"{"state":"x","questions":{"a":{},"b":{}},"max_len":256}"#;
|
||||
assert_eq!(decision_reservation_tokens(body), 512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decision_reservation_caps_a_raised_max_len() {
|
||||
let body = br#"{"questions":{"a":{}},"max_len":1000000}"#;
|
||||
assert_eq!(decision_reservation_tokens(body), DECISION_MAX_LEN_CAP);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn decision_reservation_never_zero() {
|
||||
assert_eq!(
|
||||
decision_reservation_tokens(b"not json"),
|
||||
DECISION_DEFAULT_MAX_LEN
|
||||
);
|
||||
let body = br#"{"questions":{},"max_len":0}"#;
|
||||
assert_eq!(decision_reservation_tokens(body), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn requested_max_output_prefers_max_completion_tokens() {
|
||||
let body = br#"{"model":"m","max_completion_tokens":256,"max_tokens":99}"#;
|
||||
|
||||
@@ -0,0 +1,274 @@
|
||||
//! `/v1/systemone` decision proxy tests (#338).
|
||||
//!
|
||||
//! The mock neuron echoes what it received — the `model` field and the
|
||||
//! routed-model header — so each test can assert both halves of the
|
||||
//! contract: *where* cortex sent the request, and that it did not
|
||||
//! rewrite the body on the way.
|
||||
|
||||
use axum::Router;
|
||||
use axum::http::HeaderMap;
|
||||
use axum::response::{IntoResponse, Json};
|
||||
use axum::routing::{get, post};
|
||||
use cortex_core::catalogue::{DECISION_DEFAULT_ALIAS, DECISION_MODEL_HEADER};
|
||||
use cortex_core::config::{
|
||||
ApiKeyConfig, EntitlementsConfig, EvictionSettings, EvictionStrategy, GatewayConfig,
|
||||
GatewaySettings, NeuronEndpoint,
|
||||
};
|
||||
use cortex_core::entitlements::CapWindow;
|
||||
use cortex_core::node::{ModelEntry, ModelStatus};
|
||||
use cortex_gateway::state::CortexState;
|
||||
use serde_json::{Value, json};
|
||||
use std::sync::Arc;
|
||||
use tokio::net::TcpListener;
|
||||
|
||||
const DECISION_MODEL: &str = "convaiinnovations/laya";
|
||||
const TEXT_MODEL: &str = "Qwen/Qwen3-1.7B";
|
||||
|
||||
/// Mock neuron: `/v1/systemone` answers with a Jev-shaped body that
|
||||
/// echoes the forwarded `model` and routed-model header, and reports
|
||||
/// `usage.input_tokens` = 41 per question.
|
||||
async fn spawn_decision_mock_neuron() -> String {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let base_url = format!("http://{addr}");
|
||||
let inference_url = base_url.clone();
|
||||
let app = Router::new()
|
||||
.route(
|
||||
"/models/{*rest}",
|
||||
get(move || {
|
||||
let url = inference_url.clone();
|
||||
async move { Json(json!({"url": url})) }
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/v1/systemone",
|
||||
post(|headers: HeaderMap, Json(body): Json<Value>| async move {
|
||||
let routed = headers
|
||||
.get(DECISION_MODEL_HEADER)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
let questions = body
|
||||
.get("questions")
|
||||
.and_then(Value::as_object)
|
||||
.map(|m| m.len() as u64)
|
||||
.unwrap_or(0);
|
||||
(
|
||||
[("server-timing", "inference;dur=12.50")],
|
||||
Json(json!({
|
||||
"model": "laya-rl-agent",
|
||||
"answers": {},
|
||||
"usage": {"input_tokens": 41 * questions, "output_tokens": 0},
|
||||
"echo_model": body.get("model").cloned().unwrap_or(Value::Null),
|
||||
"echo_routed": routed,
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
}),
|
||||
);
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
base_url
|
||||
}
|
||||
|
||||
fn entry(id: &str, capabilities: &[&str]) -> ModelEntry {
|
||||
ModelEntry {
|
||||
id: id.into(),
|
||||
status: ModelStatus::Loaded,
|
||||
last_accessed: None,
|
||||
vram_estimate_mb: Some(1000),
|
||||
capabilities: capabilities.iter().map(|c| c.to_string()).collect(),
|
||||
tool_call: false,
|
||||
reasoning: false,
|
||||
limit: None,
|
||||
servable: None,
|
||||
reasoning_budget: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// A gateway fronting the mock with one decision and one text model
|
||||
/// loaded. `default_alias` controls whether `helexa/one` is configured;
|
||||
/// `hard_cap` switches on keyed auth with that balance.
|
||||
async fn spawn_gateway(neuron_url: &str, default_alias: bool, hard_cap: Option<u64>) -> String {
|
||||
let entitlements = match hard_cap {
|
||||
Some(cap) => EntitlementsConfig {
|
||||
require_auth: true,
|
||||
keys: vec![ApiKeyConfig {
|
||||
key: "sk-dec".into(),
|
||||
account_id: "acct-dec".into(),
|
||||
key_id: Some("key-dec".into()),
|
||||
hard_cap: Some(cap),
|
||||
window: CapWindow::Balance,
|
||||
}],
|
||||
},
|
||||
None => EntitlementsConfig::default(),
|
||||
};
|
||||
let config = GatewayConfig {
|
||||
gateway: GatewaySettings {
|
||||
listen: "127.0.0.1:0".into(),
|
||||
metrics_listen: "127.0.0.1:0".into(),
|
||||
},
|
||||
eviction: EvictionSettings {
|
||||
strategy: EvictionStrategy::Lru,
|
||||
defrag_after_cycles: 0,
|
||||
},
|
||||
neurons: vec![NeuronEndpoint {
|
||||
name: "mock-node".into(),
|
||||
endpoint: neuron_url.to_string(),
|
||||
}],
|
||||
models_config: "/dev/null".into(),
|
||||
entitlements,
|
||||
upstream: Default::default(),
|
||||
};
|
||||
let mut state = CortexState::from_config(&config);
|
||||
if default_alias {
|
||||
state
|
||||
.catalogue
|
||||
.aliases
|
||||
.insert(DECISION_DEFAULT_ALIAS.into(), DECISION_MODEL.into());
|
||||
}
|
||||
let fleet = Arc::new(state);
|
||||
{
|
||||
let mut nodes = fleet.nodes.write().await;
|
||||
let node = nodes.get_mut("mock-node").unwrap();
|
||||
node.healthy = true;
|
||||
node.models
|
||||
.insert(DECISION_MODEL.into(), entry(DECISION_MODEL, &["decision"]));
|
||||
node.models
|
||||
.insert(TEXT_MODEL.into(), entry(TEXT_MODEL, &[]));
|
||||
}
|
||||
let app = cortex_gateway::build_app(Arc::clone(&fleet));
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
async fn send(gateway: &str, body: Value) -> (u16, HeaderMap, Value) {
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{gateway}/v1/systemone"))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.expect("request should reach the gateway");
|
||||
let status = resp.status().as_u16();
|
||||
let mut headers = HeaderMap::new();
|
||||
for (k, v) in resp.headers() {
|
||||
headers.insert(k.clone(), v.clone());
|
||||
}
|
||||
(status, headers, resp.json().await.unwrap())
|
||||
}
|
||||
|
||||
fn request(model: Option<&str>) -> Value {
|
||||
let mut body = json!({
|
||||
"state": "I was charged twice this month",
|
||||
"questions": {"refund": {"type": "noul", "instructions": "Refund request?"}},
|
||||
});
|
||||
if let Some(m) = model {
|
||||
body["model"] = json!(m);
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_explicit_decision_model_routes_and_body_is_untouched() {
|
||||
let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
|
||||
let (status, headers, body) = send(&gateway, request(Some(DECISION_MODEL))).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
assert_eq!(body["echo_routed"], DECISION_MODEL);
|
||||
assert_eq!(body["echo_model"], DECISION_MODEL);
|
||||
assert_eq!(body["usage"]["input_tokens"], 41);
|
||||
// Upstream timing headers survive the hop.
|
||||
assert_eq!(
|
||||
headers.get("server-timing").and_then(|v| v.to_str().ok()),
|
||||
Some("inference;dur=12.50")
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_foreign_model_name_falls_back_to_default_alias() {
|
||||
// A Jev client names its own service. The request must still be
|
||||
// served — routed via helexa/one — and the client's value must reach
|
||||
// the node unchanged (it may be a checkpoint pin the node honours).
|
||||
let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
|
||||
let (status, _, body) = send(&gateway, request(Some("jev-1"))).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
assert_eq!(body["echo_routed"], DECISION_MODEL);
|
||||
assert_eq!(body["echo_model"], "jev-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_missing_model_falls_back_to_default_alias() {
|
||||
let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
|
||||
let (status, _, body) = send(&gateway, request(None)).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
assert_eq!(body["echo_routed"], DECISION_MODEL);
|
||||
assert_eq!(body["echo_model"], Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_default_alias_itself_resolves() {
|
||||
let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
|
||||
let (status, _, body) = send(&gateway, request(Some(DECISION_DEFAULT_ALIAS))).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
assert_eq!(body["echo_routed"], DECISION_MODEL);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_known_non_decision_model_is_not_substituted() {
|
||||
// A model the fleet does serve routes to itself; the node owns the
|
||||
// wrong_modality answer. Silently swapping in the decision model
|
||||
// would hide the client's mistake.
|
||||
let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, None).await;
|
||||
let (status, _, body) = send(&gateway, request(Some(TEXT_MODEL))).await;
|
||||
assert_eq!(status, 200, "{body}");
|
||||
assert_eq!(body["echo_routed"], TEXT_MODEL);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_unknown_model_without_default_alias_is_404() {
|
||||
let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, false, None).await;
|
||||
let (status, _, body) = send(&gateway, request(Some("jev-1"))).await;
|
||||
assert_eq!(status, 404, "{body}");
|
||||
assert_eq!(body["error"]["code"], "model_not_found");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_budget_below_reservation_rejected_before_dispatch() {
|
||||
// One question reserves one full sequence (1024 tokens by default);
|
||||
// a 100-token balance cannot cover it.
|
||||
let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, Some(100)).await;
|
||||
let resp = reqwest::Client::new()
|
||||
.post(format!("{gateway}/v1/systemone"))
|
||||
.bearer_auth("sk-dec")
|
||||
.json(&request(Some(DECISION_MODEL)))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 429);
|
||||
let body: Value = resp.json().await.unwrap();
|
||||
assert_eq!(body["error"]["code"], "insufficient_quota");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_spend_settles_to_reported_input_tokens() {
|
||||
// Balance for exactly two one-question requests at the reserved
|
||||
// bound (2 × 1024). Settling to the reported 41 tokens leaves room
|
||||
// for many more; settling at the bound would exhaust it on the third.
|
||||
let gateway = spawn_gateway(&spawn_decision_mock_neuron().await, true, Some(2048)).await;
|
||||
let client = reqwest::Client::new();
|
||||
for i in 0..5 {
|
||||
let resp = client
|
||||
.post(format!("{gateway}/v1/systemone"))
|
||||
.bearer_auth("sk-dec")
|
||||
.json(&request(Some(DECISION_MODEL)))
|
||||
.send()
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(resp.status(), 200, "request {i} should be affordable");
|
||||
// Settlement is spawned; let it land before the next reservation.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
|
||||
}
|
||||
}
|
||||
@@ -153,17 +153,29 @@ pub trait Scenario: Send + Sync {
|
||||
fn prompt_size(&self) -> u32;
|
||||
|
||||
/// Whether this scenario should run against the given model. The
|
||||
/// default runs against everything except image-generation models
|
||||
/// (a chat request at an image model is a 422 `wrong_modality`);
|
||||
/// modality-specific scenarios override this.
|
||||
/// default runs against chat models only: an image-generation model
|
||||
/// or a decision model (#333) answers a chat request with a 422
|
||||
/// `wrong_modality`. Modality-specific scenarios override this.
|
||||
fn applies_to(&self, model: &ModelInfo) -> bool {
|
||||
!model.capabilities.iter().any(|c| c == "image")
|
||||
serves_chat(model)
|
||||
}
|
||||
|
||||
/// Issue one shaped request and measure it.
|
||||
async fn run(&self, ctx: &RunCtx) -> Result<ScenarioMetrics>;
|
||||
}
|
||||
|
||||
/// Capabilities that mark a model as serving an endpoint other than
|
||||
/// chat completions.
|
||||
const NON_CHAT_MODALITIES: [&str; 2] = ["image", "decision"];
|
||||
|
||||
/// Whether a model answers `/v1/chat/completions`.
|
||||
pub fn serves_chat(model: &ModelInfo) -> bool {
|
||||
!model
|
||||
.capabilities
|
||||
.iter()
|
||||
.any(|c| NON_CHAT_MODALITIES.contains(&c.as_str()))
|
||||
}
|
||||
|
||||
/// Fixed-seed image-latency scenario (#203): one 9-step generation at a
|
||||
/// configured square size against `/v1/images/generations`, recording
|
||||
/// the server-side phase timing (`usage.helexa_timing`) and the metered
|
||||
@@ -883,6 +895,34 @@ async fn stream_and_measure_inner(
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn model(capabilities: &[&str]) -> ModelInfo {
|
||||
serde_json::from_value(serde_json::json!({
|
||||
"id": "m",
|
||||
"harness": "candle",
|
||||
"status": "loaded",
|
||||
"devices": [0],
|
||||
"vram_used_mb": null,
|
||||
"capabilities": capabilities,
|
||||
}))
|
||||
.expect("model info")
|
||||
}
|
||||
|
||||
/// Chat scenarios skip every model that answers chat with
|
||||
/// `wrong_modality`: image generators and decision models alike.
|
||||
#[test]
|
||||
fn chat_scenarios_skip_image_and_decision_models() {
|
||||
assert!(serves_chat(&model(&[])));
|
||||
assert!(serves_chat(&model(&["text", "vision"])));
|
||||
assert!(!serves_chat(&model(&["image"])));
|
||||
assert!(!serves_chat(&model(&["decision"])));
|
||||
let image = ImageLatencyScenario {
|
||||
id: "image:512".into(),
|
||||
side_px: 512,
|
||||
};
|
||||
assert!(image.applies_to(&model(&["image"])));
|
||||
assert!(!image.applies_to(&model(&["decision"])));
|
||||
}
|
||||
|
||||
/// A burst must report server-side phase timing, or `queue_wait_ms`
|
||||
/// (TTFT minus server prefill, #85) cannot be computed for the very
|
||||
/// cells it was designed for, and the prefill-tok/s panel stays
|
||||
|
||||
@@ -149,25 +149,81 @@ pub async fn dispatch(
|
||||
|
||||
let candidates = match select_cortexes(state, &model).await {
|
||||
Selection::Candidates(c) => c,
|
||||
Selection::UnknownModel => {
|
||||
return envelope_response(
|
||||
OpenAiError::new(
|
||||
404,
|
||||
"invalid_request_error",
|
||||
"model_not_found",
|
||||
format!("no operator serves model '{model}'"),
|
||||
)
|
||||
.with_param("model"),
|
||||
);
|
||||
}
|
||||
Selection::NoReachableCapacity => {
|
||||
return envelope_response(OpenAiError::service_unavailable(
|
||||
format!("model '{model}' is temporarily unavailable on all operators"),
|
||||
Some(RETRY_AFTER_SECS),
|
||||
));
|
||||
}
|
||||
other => return selection_error(other, &model),
|
||||
};
|
||||
forward(state, path, headers, body, &model, candidates).await
|
||||
}
|
||||
|
||||
/// Proxy a `/v1/systemone` decision request (#338).
|
||||
///
|
||||
/// Same as [`dispatch`], except that a missing `model`, or one no
|
||||
/// operator serves, selects capacity for
|
||||
/// [`DECISION_DEFAULT_ALIAS`](cortex_core::catalogue::DECISION_DEFAULT_ALIAS)
|
||||
/// instead of 404ing: decision clients written against the Jev API name
|
||||
/// their own service in `model`. The body is forwarded with the client's
|
||||
/// `model` intact — cortex applies the same fallback and the serving node
|
||||
/// may read it as a checkpoint pin.
|
||||
pub async fn dispatch_decision(
|
||||
state: &RouterState,
|
||||
path: &str,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
use cortex_core::catalogue::DECISION_DEFAULT_ALIAS;
|
||||
|
||||
let requested = extract_model(&body);
|
||||
let (mut model, body) = match requested.as_deref() {
|
||||
Some(m) => match state.aliases.get(m) {
|
||||
Some(real) => {
|
||||
tracing::debug!(alias = m, model = %real, "resolving tier alias");
|
||||
(real.clone(), rewrite_model_in_body(&body, real))
|
||||
}
|
||||
None => (m.to_string(), body),
|
||||
},
|
||||
None => (DECISION_DEFAULT_ALIAS.to_string(), body),
|
||||
};
|
||||
let mut selection = select_cortexes(state, &model).await;
|
||||
if selection == Selection::UnknownModel
|
||||
&& let Some(default) = state.aliases.get(DECISION_DEFAULT_ALIAS)
|
||||
{
|
||||
tracing::debug!(requested = %model, default = %default, "unknown decision model; using the default");
|
||||
model = default.clone();
|
||||
selection = select_cortexes(state, &model).await;
|
||||
}
|
||||
match selection {
|
||||
Selection::Candidates(c) => forward(state, path, headers, body, &model, c).await,
|
||||
other => selection_error(other, &model),
|
||||
}
|
||||
}
|
||||
|
||||
/// The #63 envelope for a selection that produced no candidates.
|
||||
fn selection_error(selection: Selection, model: &str) -> Response {
|
||||
match selection {
|
||||
Selection::Candidates(_) | Selection::UnknownModel => envelope_response(
|
||||
OpenAiError::new(
|
||||
404,
|
||||
"invalid_request_error",
|
||||
"model_not_found",
|
||||
format!("no operator serves model '{model}'"),
|
||||
)
|
||||
.with_param("model"),
|
||||
),
|
||||
Selection::NoReachableCapacity => envelope_response(OpenAiError::service_unavailable(
|
||||
format!("model '{model}' is temporarily unavailable on all operators"),
|
||||
Some(RETRY_AFTER_SECS),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Try `candidates` in order, failing over only on transport errors.
|
||||
async fn forward(
|
||||
state: &RouterState,
|
||||
path: &str,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
model: &str,
|
||||
candidates: Vec<CortexEndpoint>,
|
||||
) -> Response {
|
||||
// Try candidates in order, failing over only on transport errors. A
|
||||
// genuine HTTP response (any status — including cortex's #63 429/400)
|
||||
// is returned verbatim and never retried away.
|
||||
|
||||
@@ -14,6 +14,7 @@ pub fn api_routes() -> Router<Arc<RouterState>> {
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/completions", post(completions))
|
||||
.route("/v1/images/generations", post(images_generations))
|
||||
.route("/v1/systemone", post(systemone))
|
||||
.route("/v1/responses", post(responses))
|
||||
.route("/v1/messages", post(messages))
|
||||
.route("/v1/models", get(list_models))
|
||||
@@ -63,6 +64,17 @@ async fn images_generations(
|
||||
dispatch::dispatch(&state, "/v1/images/generations", headers, body).await
|
||||
}
|
||||
|
||||
/// `POST /v1/systemone` (#338): typed decisions (TypeSafe Jev wire
|
||||
/// protocol) from a decision model. Unlike the other paths an unknown or
|
||||
/// missing `model` is not a 404 — see [`dispatch::dispatch_decision`].
|
||||
async fn systemone(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> Response {
|
||||
dispatch::dispatch_decision(&state, "/v1/systemone", headers, body).await
|
||||
}
|
||||
|
||||
async fn messages(
|
||||
State(state): State<Arc<RouterState>>,
|
||||
headers: HeaderMap,
|
||||
|
||||
@@ -14,7 +14,7 @@ use axum::routing::post;
|
||||
use axum::{Json, Router};
|
||||
use cortex_core::node::CortexModelEntry;
|
||||
use helexa_router::config::{CortexEndpoint, RouterConfig};
|
||||
use helexa_router::dispatch::{Selection, dispatch, select_cortexes};
|
||||
use helexa_router::dispatch::{Selection, dispatch, dispatch_decision, select_cortexes};
|
||||
use helexa_router::state::{CortexTopology, RouterState};
|
||||
use serde_json::{Value, json};
|
||||
use std::collections::HashMap;
|
||||
@@ -87,6 +87,7 @@ async fn mock_handler(State(m): State<MockCortex>, headers: HeaderMap, body: Byt
|
||||
async fn spawn_cortex(mock: MockCortex) -> String {
|
||||
let app = Router::new()
|
||||
.route("/v1/chat/completions", post(mock_handler))
|
||||
.route("/v1/systemone", post(mock_handler))
|
||||
.with_state(mock);
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
@@ -349,3 +350,75 @@ async fn ranking_prefers_loaded_then_region_then_headroom() {
|
||||
let names: Vec<&str> = order.iter().map(|e| e.name.as_str()).collect();
|
||||
assert_eq!(names, vec!["warm-eu", "warm-us", "cold-eu"]);
|
||||
}
|
||||
|
||||
// ── /v1/systemone default-model fallback (#338) ──────────────────────
|
||||
|
||||
/// Router state with c1 serving MODEL, and `helexa/one` pointing at it
|
||||
/// when `with_default` is set.
|
||||
async fn decision_state(with_default: bool) -> RouterState {
|
||||
let url = spawn_cortex(ok_cortex("c1")).await;
|
||||
let mut state = state_with(vec![ep("c1", &url, None)], None);
|
||||
if with_default {
|
||||
state.aliases.insert(
|
||||
cortex_core::catalogue::DECISION_DEFAULT_ALIAS.to_string(),
|
||||
MODEL.to_string(),
|
||||
);
|
||||
}
|
||||
set_topology(&state, "c1", true, true, true, 2).await;
|
||||
state
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decision_foreign_model_selects_default_and_keeps_body_model() {
|
||||
// A Jev client names its own service. Capacity is chosen for the
|
||||
// default decision model, and the client's `model` is forwarded
|
||||
// untouched for cortex (and the node) to interpret.
|
||||
let state = decision_state(true).await;
|
||||
let body = Bytes::from_static(b"{\"model\":\"jev-1\",\"questions\":{}}");
|
||||
let resp = dispatch_decision(&state, "/v1/systemone", HeaderMap::new(), body).await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(body["served_by"], "c1");
|
||||
assert_eq!(body["model_seen"], "jev-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decision_missing_model_selects_default() {
|
||||
let state = decision_state(true).await;
|
||||
let body = Bytes::from_static(b"{\"questions\":{}}");
|
||||
let resp = dispatch_decision(&state, "/v1/systemone", HeaderMap::new(), body).await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(body["served_by"], "c1");
|
||||
assert_eq!(body["model_seen"], "");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decision_default_alias_is_rewritten_like_any_alias() {
|
||||
let state = decision_state(true).await;
|
||||
let body = Bytes::from_static(b"{\"model\":\"helexa/one\",\"questions\":{}}");
|
||||
let resp = dispatch_decision(&state, "/v1/systemone", HeaderMap::new(), body).await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
assert_eq!(status, StatusCode::OK);
|
||||
assert_eq!(body["model_seen"], MODEL);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn decision_unknown_model_without_default_is_404() {
|
||||
let state = decision_state(false).await;
|
||||
let body = Bytes::from_static(b"{\"model\":\"jev-1\",\"questions\":{}}");
|
||||
let resp = dispatch_decision(&state, "/v1/systemone", HeaderMap::new(), body).await;
|
||||
let (status, body) = body_json(resp).await;
|
||||
assert_eq!(status, StatusCode::NOT_FOUND);
|
||||
assert_eq!(body["error"]["code"], "model_not_found");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn chat_unknown_model_does_not_take_the_decision_fallback() {
|
||||
// The fallback is scoped to the decision endpoint.
|
||||
let state = decision_state(true).await;
|
||||
let body = Bytes::from_static(b"{\"model\":\"jev-1\",\"stream\":false}");
|
||||
let resp = dispatch(&state, "/v1/chat/completions", HeaderMap::new(), body).await;
|
||||
let (status, _) = body_json(resp).await;
|
||||
assert_eq!(status, StatusCode::NOT_FOUND);
|
||||
}
|
||||
|
||||
@@ -76,6 +76,12 @@ toml.workspace = true
|
||||
# transitively via candle-core.
|
||||
rayon = "1"
|
||||
|
||||
# Unicode tables for the decision-model language detector (#339), which
|
||||
# must classify characters exactly as Python's `str.isalpha` / `re` and
|
||||
# `unicodedata.combining` do. Both already in the tree via tokenizers.
|
||||
unicode-normalization = "0.1"
|
||||
unicode-properties = "0.1"
|
||||
|
||||
# candle for in-process inference. CUDA support is gated behind the
|
||||
# crate's `cuda` feature (default off) so the workspace builds on
|
||||
# non-CUDA hosts and CI runners.
|
||||
|
||||
@@ -55,6 +55,7 @@ pub fn neuron_routes() -> Router<Arc<NeuronState>> {
|
||||
.route("/models/{model_id}/capabilities", get(model_capabilities))
|
||||
.route("/v1/chat/completions", post(chat_completions))
|
||||
.route("/v1/images/generations", post(images_generations))
|
||||
.route("/v1/systemone", post(systemone))
|
||||
.route("/v1/responses", post(responses))
|
||||
}
|
||||
|
||||
@@ -189,6 +190,7 @@ fn preflight_kind(err: &PreflightError) -> &'static str {
|
||||
PreflightError::EmptyRepo { .. } => "empty_repo",
|
||||
PreflightError::TpRequiresSafetensors { .. } => "tp_requires_safetensors",
|
||||
PreflightError::QuantNotFound { .. } => "quant_not_found",
|
||||
PreflightError::Unsupported { .. } => "unsupported_load_option",
|
||||
}
|
||||
}
|
||||
|
||||
@@ -668,6 +670,83 @@ async fn images_generations(
|
||||
}
|
||||
}
|
||||
|
||||
/// `POST /v1/systemone` (#338): typed decisions over the TypeSafe Jev
|
||||
/// wire protocol, from a loaded decision model (Laya). Parsing and the
|
||||
/// transport-level limits happen here; routing, validation, admission,
|
||||
/// scoring and shaping in `CandleHarness::systemone`.
|
||||
///
|
||||
/// Errors carry the #63 envelope *and* the reference server's top-level
|
||||
/// `detail` string, so OpenAI-style and Jev/laya clients both read them.
|
||||
async fn systemone(
|
||||
State(state): State<Arc<NeuronState>>,
|
||||
headers: axum::http::HeaderMap,
|
||||
body: axum::body::Bytes,
|
||||
) -> axum::response::Response {
|
||||
use crate::harness::candle::SystemOneError;
|
||||
use cortex_core::decisions::{SystemOneRequest, limits};
|
||||
|
||||
let Some(candle) = state.candle.as_ref().map(Arc::clone) else {
|
||||
return decision_error_response(503, "candle harness not enabled on this neuron");
|
||||
};
|
||||
if body.len() > limits::MAX_BODY_BYTES {
|
||||
return decision_error_response(413, "request body too large");
|
||||
}
|
||||
let request = match SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET) {
|
||||
Ok(r) => r,
|
||||
Err(e) => return decision_error_response(e.status, &e.detail),
|
||||
};
|
||||
let served = headers
|
||||
.get(cortex_core::catalogue::DECISION_MODEL_HEADER)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(str::to_string);
|
||||
let principal = principal_key(&headers);
|
||||
|
||||
let started = std::time::Instant::now();
|
||||
match candle
|
||||
.systemone(served.as_deref(), request, principal.as_deref())
|
||||
.await
|
||||
{
|
||||
Ok(response) => {
|
||||
let ms = started.elapsed().as_secs_f64() * 1e3;
|
||||
let mut out = Json(response).into_response();
|
||||
for (name, value) in [
|
||||
("server-timing", format!("inference;dur={ms:.2}")),
|
||||
("x-inference-time-ms", format!("{ms:.2}")),
|
||||
] {
|
||||
if let Ok(v) = axum::http::HeaderValue::from_str(&value) {
|
||||
out.headers_mut().insert(name, v);
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
Err(SystemOneError::Request(e)) => decision_error_response(e.status, &e.detail),
|
||||
Err(SystemOneError::Inference(e)) => {
|
||||
tracing::warn!(error = %e, "systemone request failed");
|
||||
inference_error_response(e)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A rejected decision request: the #63 envelope, plus the `detail` field
|
||||
/// the reference (FastAPI) server answers with.
|
||||
fn decision_error_response(status: u16, detail: &str) -> axum::response::Response {
|
||||
use cortex_core::error_envelope::OpenAiError;
|
||||
let (kind, code) = match status {
|
||||
404 => ("invalid_request_error", "model_not_found"),
|
||||
413 => ("invalid_request_error", "request_too_large"),
|
||||
422 => ("invalid_request_error", "invalid_decision_request"),
|
||||
400..=499 => ("invalid_request_error", "invalid_request"),
|
||||
_ => ("api_error", "service_unavailable"),
|
||||
};
|
||||
let env = OpenAiError::new(status, kind, code, detail.to_string());
|
||||
let mut body = env.body();
|
||||
if let Some(obj) = body.as_object_mut() {
|
||||
obj.insert("detail".into(), json!(detail));
|
||||
}
|
||||
let status = StatusCode::from_u16(status).unwrap_or(StatusCode::BAD_REQUEST);
|
||||
(status, Json(body)).into_response()
|
||||
}
|
||||
|
||||
fn finish_reason_from_str(s: &str) -> crate::wire::FinishReason {
|
||||
use crate::wire::FinishReason;
|
||||
match s {
|
||||
|
||||
@@ -94,6 +94,27 @@ struct Args {
|
||||
/// How many recent tokens the repeat penalty considers.
|
||||
#[arg(long, default_value_t = 64)]
|
||||
repeat_last_n: usize,
|
||||
/// Tensor-parallel degree. 1 loads the model on one card; 2+ uses
|
||||
/// the worker pool, which is the only way to reach a model too big
|
||||
/// for a single GPU — the 27B is 54 GB in bf16 and the single-GPU
|
||||
/// path is bf16-only.
|
||||
#[arg(long, default_value_t = 1)]
|
||||
tensor_parallel: u32,
|
||||
/// Quantisation for the TP load (`q6k`, `q8_0`, …). Only the TP
|
||||
/// loader applies it.
|
||||
#[arg(long)]
|
||||
quant: Option<String>,
|
||||
/// The binary the worker ranks are spawned from. Defaults to
|
||||
/// `neuron` beside this executable — the probe itself does not
|
||||
/// implement `--worker`.
|
||||
#[arg(long)]
|
||||
worker_binary: Option<PathBuf>,
|
||||
/// How many verify-shaped forwards to time after the measurement
|
||||
/// loop. A single sample moved 57.3 -> 50.1 ms between otherwise
|
||||
/// identical runs — ±7 ms on a 45 ms baseline, enough to move the
|
||||
/// predicted speedup across the gate on its own.
|
||||
#[arg(long, default_value_t = 9)]
|
||||
verify_samples: usize,
|
||||
/// Write the summary as JSON here as well as to stdout.
|
||||
#[arg(long)]
|
||||
json: Option<PathBuf>,
|
||||
@@ -217,14 +238,12 @@ fn prefill(
|
||||
last.context("empty prompt")
|
||||
}
|
||||
|
||||
struct Timings {
|
||||
target_decode_ms: f64,
|
||||
draft_step_ms: f64,
|
||||
verify_ms: f64,
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
if args.tensor_parallel > 1 {
|
||||
return run_tp(args).await;
|
||||
}
|
||||
let dev = pick_device(&args.device)?;
|
||||
let dtype = if matches!(dev, Device::Cpu) {
|
||||
DType::F32
|
||||
@@ -310,7 +329,6 @@ fn main() -> Result<()> {
|
||||
let mut accepted_hist = vec![0usize; args.draft_len + 1];
|
||||
let mut accepted_hist_penalised = vec![0usize; args.draft_len + 1];
|
||||
let mut penalty_disagreements = 0usize;
|
||||
let mut timings: Option<Timings> = None;
|
||||
let mut draft_ms_total = 0f64;
|
||||
let mut draft_rounds = 0usize;
|
||||
let mut target_ms_total = 0f64;
|
||||
@@ -394,50 +412,6 @@ fn main() -> Result<()> {
|
||||
target_ms_total += target_elapsed;
|
||||
target_steps += 1;
|
||||
}
|
||||
|
||||
// One verify-shaped forward, once, for the speedup arithmetic:
|
||||
// K+1 positions in a single pass is what a real round costs the
|
||||
// target, and it is a different kernel from a 1-token decode.
|
||||
if round == args.warmup {
|
||||
let block: Vec<u32> = std::iter::once(next)
|
||||
.chain(std::iter::repeat_n(next, args.draft_len))
|
||||
.collect();
|
||||
let blk = Tensor::new(block.as_slice(), &dev)?.unsqueeze(0)?;
|
||||
let snap_t = Instant::now();
|
||||
let _ = model.forward_multi(&blk, pos + 1)?;
|
||||
sync(&dev)?;
|
||||
let verify_ms = snap_t.elapsed().as_secs_f64() * 1000.0;
|
||||
// That forward advanced the target's cache by K+1; put it
|
||||
// back by re-prefilling is not possible here, so the probe
|
||||
// measures verify cost on a throwaway clone of the state.
|
||||
timings = Some(Timings {
|
||||
target_decode_ms: target_elapsed,
|
||||
draft_step_ms: draft_elapsed / args.draft_len as f64,
|
||||
verify_ms,
|
||||
});
|
||||
// Re-align: the verify pass wrote K+1 positions into the
|
||||
// target's KV that the observer never committed.
|
||||
model.clear_kv_cache();
|
||||
head.clear_kv_cache();
|
||||
let replay: Vec<u32> = ids
|
||||
.iter()
|
||||
.copied()
|
||||
.chain(produced.iter().copied())
|
||||
.collect();
|
||||
let (lg, hd) = prefill(
|
||||
&mut model,
|
||||
&mut head,
|
||||
&rotary,
|
||||
&replay,
|
||||
args.prefill_chunk,
|
||||
dtype,
|
||||
&dev,
|
||||
&sync,
|
||||
)?;
|
||||
last_logits = lg.i((0, 0, ..))?.to_dtype(DType::F32)?.to_vec1()?;
|
||||
hidden_cur = hd;
|
||||
pending.clear();
|
||||
}
|
||||
}
|
||||
|
||||
// Anything still pending never got its K true tokens; score what is
|
||||
@@ -456,6 +430,24 @@ fn main() -> Result<()> {
|
||||
}
|
||||
let _ = &mut accepted_hist_penalised;
|
||||
|
||||
// ── verify cost, sampled ────────────────────────────────────────
|
||||
//
|
||||
// Timed after the loop because each sample advances the target's
|
||||
// cache by K+1 positions — harmless once measuring is done, and it
|
||||
// spares the loop the re-prefill that undoing one sample required.
|
||||
let mut verify_samples = Vec::with_capacity(args.verify_samples);
|
||||
for i in 0..args.verify_samples {
|
||||
let block: Vec<u32> = std::iter::repeat_n(produced[0], args.draft_len + 1).collect();
|
||||
let blk = Tensor::new(block.as_slice(), &dev)?.unsqueeze(0)?;
|
||||
let at = n + produced.len() + i * (args.draft_len + 1);
|
||||
let t0 = Instant::now();
|
||||
let _ = model.forward_multi(&blk, at)?;
|
||||
sync(&dev)?;
|
||||
verify_samples.push(t0.elapsed().as_secs_f64() * 1000.0);
|
||||
}
|
||||
verify_samples.sort_by(f64::total_cmp);
|
||||
let verify_ms = verify_samples[verify_samples.len() / 2];
|
||||
|
||||
// ── report ──────────────────────────────────────────────────────
|
||||
let rounds: usize = accepted_hist.iter().sum();
|
||||
let mean_accepted: f64 = accepted_hist
|
||||
@@ -464,13 +456,12 @@ fn main() -> Result<()> {
|
||||
.map(|(k, n)| (k * n) as f64)
|
||||
.sum::<f64>()
|
||||
/ rounds.max(1) as f64;
|
||||
let t = timings.context("no timing round ran; --steps must exceed 0")?;
|
||||
let target_ms = target_ms_total / target_steps.max(1) as f64;
|
||||
let draft_ms = draft_ms_total / (draft_rounds.max(1) * args.draft_len) as f64;
|
||||
let c = draft_ms / target_ms;
|
||||
// One round commits 1 + accepted tokens and costs one verify plus K
|
||||
// drafts, against (1 + accepted) plain decode steps.
|
||||
let round_cost = t.verify_ms + args.draft_len as f64 * draft_ms;
|
||||
let round_cost = verify_ms + args.draft_len as f64 * draft_ms;
|
||||
let speedup = (1.0 + mean_accepted) * target_ms / round_cost;
|
||||
|
||||
let summary = serde_json::json!({
|
||||
@@ -482,7 +473,10 @@ fn main() -> Result<()> {
|
||||
"rounds_scored": rounds,
|
||||
"target_decode_ms": target_ms,
|
||||
"draft_step_ms": draft_ms,
|
||||
"verify_ms_k_plus_1": t.verify_ms,
|
||||
"verify_ms_k_plus_1": verify_ms,
|
||||
"verify_ms_min": verify_samples.first(),
|
||||
"verify_ms_max": verify_samples.last(),
|
||||
"verify_sample_count": verify_samples.len(),
|
||||
"c_draft_over_target": c,
|
||||
"mean_accepted_tokens": mean_accepted,
|
||||
"accepted_histogram": accepted_hist,
|
||||
@@ -490,8 +484,6 @@ fn main() -> Result<()> {
|
||||
"penalty_argmax_disagreements": penalty_disagreements,
|
||||
"penalty_argmax_disagreement_rate":
|
||||
penalty_disagreements as f64 / produced.len().max(1) as f64,
|
||||
"first_target_decode_ms_sample": t.target_decode_ms,
|
||||
"draft_step_ms_sample": t.draft_step_ms,
|
||||
});
|
||||
let pretty = serde_json::to_string_pretty(&summary)?;
|
||||
println!("{pretty}");
|
||||
@@ -500,3 +492,239 @@ fn main() -> Result<()> {
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Tensor-parallel measurement (#96).
|
||||
///
|
||||
/// The only route to the 27B: 54 GB in bf16 does not fit one 5090, and
|
||||
/// quantisation is reachable only through the TP loader. The shape
|
||||
/// mirrors the single-GPU path — prefill, then draft-and-rewind at each
|
||||
/// decode step — with two differences:
|
||||
///
|
||||
/// - the target's step fans out to every rank, because the row-parallel
|
||||
/// `AllReduce`s only complete when all of them arrive;
|
||||
/// - the draft head is unsharded on the leader, so drafting issues no
|
||||
/// collectives at all and cannot leave a rank waiting.
|
||||
#[cfg(feature = "cuda")]
|
||||
async fn run_tp(args: Args) -> Result<()> {
|
||||
use neuron::harness::device_worker::DeviceWorkerHandle;
|
||||
use neuron::harness::tp::WorkerPool;
|
||||
|
||||
let cfg_raw =
|
||||
std::fs::read_to_string(args.model.join("config.json")).context("read config.json")?;
|
||||
let config: Config = Config::from_config_json(&cfg_raw).context("parse config.json")?;
|
||||
let text_cfg = config.text_config.clone();
|
||||
anyhow::ensure!(
|
||||
MtpHead::present_in(&text_cfg),
|
||||
"this checkpoint declares no MTP head (mtp_num_hidden_layers = 0)"
|
||||
);
|
||||
|
||||
let mut shards: Vec<PathBuf> = std::fs::read_dir(&args.model)
|
||||
.context("read model dir")?
|
||||
.filter_map(|e| e.ok().map(|e| e.path()))
|
||||
.filter(|p| p.extension().is_some_and(|x| x == "safetensors"))
|
||||
.collect();
|
||||
shards.sort();
|
||||
anyhow::ensure!(!shards.is_empty(), "no .safetensors in {:?}", args.model);
|
||||
|
||||
let worker_bin = match &args.worker_binary {
|
||||
Some(p) => p.clone(),
|
||||
None => std::env::current_exe()
|
||||
.context("resolve current_exe")?
|
||||
.parent()
|
||||
.context("executable has no parent directory")?
|
||||
.join("neuron"),
|
||||
};
|
||||
anyhow::ensure!(
|
||||
worker_bin.exists(),
|
||||
"worker binary {worker_bin:?} not found — build `neuron` beside this probe, \
|
||||
or pass --worker-binary"
|
||||
);
|
||||
|
||||
let devices: Vec<u32> = (0..args.tensor_parallel).collect();
|
||||
eprintln!(
|
||||
"tp-{}: spawning workers from {worker_bin:?} on devices {devices:?}…",
|
||||
args.tensor_parallel
|
||||
);
|
||||
let leader = DeviceWorkerHandle::spawn(devices[0])?;
|
||||
let mut pool =
|
||||
WorkerPool::spawn(&worker_bin, args.tensor_parallel, &devices, leader.clone()).await?;
|
||||
pool.init_nccl(devices[0]).await?;
|
||||
|
||||
let leader_device = Device::new_cuda(devices[0] as usize)?;
|
||||
let model_id = "mtp-probe";
|
||||
let t0 = Instant::now();
|
||||
let handle = pool
|
||||
.load_dense_shard(
|
||||
model_id,
|
||||
&cfg_raw,
|
||||
&shards,
|
||||
&leader_device,
|
||||
DType::BF16,
|
||||
args.quant.clone(),
|
||||
)
|
||||
.await?;
|
||||
eprintln!("shards loaded: {:.1} s", t0.elapsed().as_secs_f64());
|
||||
|
||||
let shard_strs: Vec<String> = shards
|
||||
.iter()
|
||||
.map(|p| p.to_string_lossy().into_owned())
|
||||
.collect();
|
||||
leader
|
||||
.tp_load_mtp_head(handle, cfg_raw.clone(), shard_strs)
|
||||
.await?;
|
||||
eprintln!("draft head loaded on the leader");
|
||||
|
||||
let tok = tokenizers::Tokenizer::from_file(args.model.join("tokenizer.json"))
|
||||
.map_err(|e| anyhow::anyhow!("load tokenizer: {e}"))?;
|
||||
let prompt = match &args.prompt_file {
|
||||
Some(p) => std::fs::read_to_string(p).context("read prompt file")?,
|
||||
None => "Explain, step by step, how a bicycle derailleur shifts gears.".into(),
|
||||
};
|
||||
let ids: Vec<u32> = tok
|
||||
.encode(prompt.as_str(), false)
|
||||
.map_err(|e| anyhow::anyhow!("tokenize: {e}"))?
|
||||
.get_ids()
|
||||
.to_vec();
|
||||
let n = ids.len();
|
||||
anyhow::ensure!(n >= 2, "prompt is too short to prefill");
|
||||
eprintln!("prompt: {n} tokens");
|
||||
|
||||
// ── prefill, chunked, target then head ──────────────────────────
|
||||
let t0 = Instant::now();
|
||||
let mut last_logits: Vec<f32> = Vec::new();
|
||||
let mut at = 0usize;
|
||||
while at < n {
|
||||
let end = (at + args.prefill_chunk).min(n);
|
||||
last_logits = pool
|
||||
.generate_step(model_id, handle, ids[at..end].to_vec(), at)
|
||||
.await?;
|
||||
let head_end = end.min(n - 1);
|
||||
if head_end > at {
|
||||
leader
|
||||
.tp_mtp_prefill_chunk(handle, ids[at + 1..head_end + 1].to_vec(), at)
|
||||
.await?;
|
||||
}
|
||||
at = end;
|
||||
}
|
||||
eprintln!("prefill ({n} tokens): {:.1} s", t0.elapsed().as_secs_f64());
|
||||
|
||||
// ── draft / verify observation ──────────────────────────────────
|
||||
let mut produced: Vec<u32> = Vec::new();
|
||||
let mut pending: Vec<(usize, Vec<u32>)> = Vec::new();
|
||||
let mut accepted_hist = vec![0usize; args.draft_len + 1];
|
||||
let mut penalty_disagreements = 0usize;
|
||||
let mut draft_ms_total = 0f64;
|
||||
let mut target_ms_total = 0f64;
|
||||
let mut measured = 0usize;
|
||||
|
||||
for round in 0..args.warmup + args.steps {
|
||||
let measuring = round >= args.warmup;
|
||||
let pos = n + produced.len();
|
||||
let next = argmax(&last_logits);
|
||||
let next_penalised = argmax_with_penalty(
|
||||
&last_logits,
|
||||
&produced[produced.len().saturating_sub(args.repeat_last_n)..],
|
||||
args.repeat_penalty,
|
||||
);
|
||||
if next != next_penalised {
|
||||
penalty_disagreements += 1;
|
||||
}
|
||||
produced.push(next);
|
||||
|
||||
pending.retain(|(start, drafted)| {
|
||||
if produced.len() - start < drafted.len() {
|
||||
return true;
|
||||
}
|
||||
let truth = &produced[*start..*start + drafted.len()];
|
||||
let acc = drafted
|
||||
.iter()
|
||||
.zip(truth.iter())
|
||||
.take_while(|(d, t)| d == t)
|
||||
.count();
|
||||
accepted_hist[acc] += 1;
|
||||
false
|
||||
});
|
||||
|
||||
let t0 = Instant::now();
|
||||
let drafted = leader
|
||||
.tp_mtp_draft(handle, next, pos, args.draft_len)
|
||||
.await?;
|
||||
let draft_elapsed = t0.elapsed().as_secs_f64() * 1000.0;
|
||||
leader.tp_mtp_advance(handle, next, pos).await?;
|
||||
|
||||
let t0 = Instant::now();
|
||||
last_logits = pool
|
||||
.generate_step(model_id, handle, vec![next], pos)
|
||||
.await?;
|
||||
let target_elapsed = t0.elapsed().as_secs_f64() * 1000.0;
|
||||
|
||||
if measuring {
|
||||
pending.push((produced.len(), drafted));
|
||||
draft_ms_total += draft_elapsed;
|
||||
target_ms_total += target_elapsed;
|
||||
measured += 1;
|
||||
}
|
||||
}
|
||||
|
||||
// ── verify cost, sampled ────────────────────────────────────────
|
||||
//
|
||||
// After the loop, because each sample advances the target's cache by
|
||||
// K+1. One sample was not enough: it moved 57.3 -> 50.1 ms between
|
||||
// otherwise identical runs, wide enough to decide the gate by itself.
|
||||
let mut verify_samples = Vec::with_capacity(args.verify_samples);
|
||||
for i in 0..args.verify_samples {
|
||||
let block: Vec<u32> = std::iter::repeat_n(produced[0], args.draft_len + 1).collect();
|
||||
let at = n + produced.len() + i * (args.draft_len + 1);
|
||||
let t0 = Instant::now();
|
||||
let _ = pool
|
||||
.generate_step_multi(model_id, handle, block, at)
|
||||
.await?;
|
||||
verify_samples.push(t0.elapsed().as_secs_f64() * 1000.0);
|
||||
}
|
||||
verify_samples.sort_by(f64::total_cmp);
|
||||
let verify_ms = verify_samples[verify_samples.len() / 2];
|
||||
|
||||
let rounds: usize = accepted_hist.iter().sum();
|
||||
let mean_accepted: f64 = accepted_hist
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(k, c)| (k * c) as f64)
|
||||
.sum::<f64>()
|
||||
/ rounds.max(1) as f64;
|
||||
let target_ms = target_ms_total / measured.max(1) as f64;
|
||||
let draft_ms = draft_ms_total / (measured.max(1) * args.draft_len) as f64;
|
||||
let round_cost = verify_ms + args.draft_len as f64 * draft_ms;
|
||||
let summary = serde_json::json!({
|
||||
"model": args.model.to_string_lossy(),
|
||||
"tensor_parallel": args.tensor_parallel,
|
||||
"quant": args.quant,
|
||||
"prompt_tokens": n,
|
||||
"draft_len": args.draft_len,
|
||||
"rounds_scored": rounds,
|
||||
"target_decode_ms": target_ms,
|
||||
"draft_step_ms": draft_ms,
|
||||
"verify_ms_k_plus_1": verify_ms,
|
||||
"verify_ms_min": verify_samples.first(),
|
||||
"verify_ms_max": verify_samples.last(),
|
||||
"verify_sample_count": verify_samples.len(),
|
||||
"c_draft_over_target": draft_ms / target_ms,
|
||||
"mean_accepted_tokens": mean_accepted,
|
||||
"accepted_histogram": accepted_hist,
|
||||
"predicted_speedup": (1.0 + mean_accepted) * target_ms / round_cost,
|
||||
"penalty_argmax_disagreements": penalty_disagreements,
|
||||
"penalty_argmax_disagreement_rate":
|
||||
penalty_disagreements as f64 / produced.len().max(1) as f64,
|
||||
});
|
||||
let pretty = serde_json::to_string_pretty(&summary)?;
|
||||
println!("{pretty}");
|
||||
if let Some(p) = &args.json {
|
||||
std::fs::write(p, &pretty).context("write --json")?;
|
||||
}
|
||||
pool.unload_model(model_id).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
async fn run_tp(_args: Args) -> Result<()> {
|
||||
anyhow::bail!("--tensor-parallel > 1 requires a build with the `cuda` feature")
|
||||
}
|
||||
|
||||
@@ -0,0 +1,478 @@
|
||||
//! ModernBERT, the bidirectional encoder under every Laya checkpoint.
|
||||
//!
|
||||
//! Follows the structure of `candle_transformers::models::modernbert`
|
||||
//! (itself a port of transformers' `ModernBertModel`), vendored rather
|
||||
//! than used as-is because that module only runs correctly in f32:
|
||||
//!
|
||||
//! - its attention masks are built in F32 whatever the model dtype, so
|
||||
//! an f16/bf16 load fails at the first `broadcast_add`;
|
||||
//! - it computes the rotary angles `position × inv_freq` in the model
|
||||
//! dtype, which in f16 already loses a quarter-radian of phase by
|
||||
//! position 500;
|
||||
//! - it loads each layer's `attn_norm` with `.ok()`, so a weight-prefix
|
||||
//! mistake silently yields a model with no pre-attention norms.
|
||||
//!
|
||||
//! Here the angles are computed in f32 and only the resulting cos/sin
|
||||
//! tables are cast; attention scores are softmaxed in f32 against f32
|
||||
//! masks; and `attn_norm` is required on every layer except layer 0,
|
||||
//! which ModernBERT builds as an identity.
|
||||
//!
|
||||
//! Precision follows PyTorch autocast, which is how the reference runs on
|
||||
//! a GPU: the linear layers (and so the attention inputs) run in the
|
||||
//! load dtype, while the residual stream, the embeddings and every
|
||||
//! LayerNorm stay in f32. Casting the residual stream to bf16 as well
|
||||
//! drifts the released English checkpoint's logits by ~0.25 over its 28
|
||||
//! layers; keeping it in f32 is what autocast does and costs little.
|
||||
//!
|
||||
//! Masked positions use a large *finite* negative rather than `-inf`.
|
||||
//! A padded query near the end of a short row can find every key in
|
||||
//! its sliding window padded; with `-inf` that row's softmax is NaN,
|
||||
//! and NaN in a padded row's value vector poisons the valid rows that
|
||||
//! give it zero weight (`0 × NaN = NaN`). A finite value makes such a
|
||||
//! row uniform instead, and valid rows never look at it.
|
||||
|
||||
use anyhow::{Context, Result, ensure};
|
||||
use candle_core::{D, DType, Device, Module, Tensor};
|
||||
use candle_nn::{Embedding, LayerNorm, Linear, VarBuilder};
|
||||
|
||||
use super::config::{EncoderConfig, LayerKind};
|
||||
|
||||
/// Additive mask value for a position that must not be attended.
|
||||
pub(crate) const MASKED: f32 = -1e9;
|
||||
|
||||
/// Apply `lin` in its weight's dtype, casting `x` in if needed. The
|
||||
/// result stays in that dtype; callers cast back to f32 where it joins
|
||||
/// the residual stream.
|
||||
pub(crate) fn project(lin: &Linear, x: &Tensor) -> candle_core::Result<Tensor> {
|
||||
x.to_dtype(lin.weight().dtype())?.apply(lin)
|
||||
}
|
||||
|
||||
/// A bias-free LayerNorm that takes candle's fused kernel.
|
||||
///
|
||||
/// `candle_nn::LayerNorm` dispatches to the fused `layer_norm` op only
|
||||
/// when it has a bias; without one it falls back to six elementwise
|
||||
/// kernels (mean, subtract, square, mean, divide, scale), which was a
|
||||
/// quarter of the encoder's GPU time. ModernBERT's norms have no bias,
|
||||
/// so they get an explicit zero one: same result, one kernel.
|
||||
fn fused_layer_norm_no_bias(size: usize, eps: f64, vb: VarBuilder) -> Result<LayerNorm> {
|
||||
let weight = vb.get(size, "weight")?;
|
||||
let bias = weight.zeros_like()?;
|
||||
Ok(LayerNorm::new(weight, bias, eps))
|
||||
}
|
||||
|
||||
/// Precomputed rotary tables for one base, `(max_pos, head_dim / 2)`.
|
||||
struct Rotary {
|
||||
cos: Tensor,
|
||||
sin: Tensor,
|
||||
}
|
||||
|
||||
impl Rotary {
|
||||
fn new(cfg: &EncoderConfig, theta: f64, dtype: DType, device: &Device) -> Result<Self> {
|
||||
let dim = cfg.head_dim();
|
||||
let inv_freq: Vec<f32> = (0..dim)
|
||||
.step_by(2)
|
||||
.map(|i| (1.0 / theta.powf(i as f64 / dim as f64)) as f32)
|
||||
.collect();
|
||||
let n = inv_freq.len();
|
||||
let inv_freq = Tensor::from_vec(inv_freq, (1, n), device)?;
|
||||
let max_pos = cfg.max_position_embeddings;
|
||||
let pos = Tensor::arange(0u32, max_pos as u32, device)?
|
||||
.to_dtype(DType::F32)?
|
||||
.reshape((max_pos, 1))?;
|
||||
let freqs = pos.matmul(&inv_freq)?;
|
||||
Ok(Self {
|
||||
cos: freqs.cos()?.to_dtype(dtype)?,
|
||||
sin: freqs.sin()?.to_dtype(dtype)?,
|
||||
})
|
||||
}
|
||||
|
||||
fn apply(&self, q: &Tensor, k: &Tensor, seq_len: usize) -> Result<(Tensor, Tensor)> {
|
||||
let cos = self.cos.narrow(0, 0, seq_len)?;
|
||||
let sin = self.sin.narrow(0, 0, seq_len)?;
|
||||
Ok((
|
||||
candle_nn::rotary_emb::rope(&q.contiguous()?, &cos, &sin)?,
|
||||
candle_nn::rotary_emb::rope(&k.contiguous()?, &cos, &sin)?,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
struct Attention {
|
||||
wqkv: Linear,
|
||||
wo: Linear,
|
||||
n_heads: usize,
|
||||
head_dim: usize,
|
||||
kind: LayerKind,
|
||||
}
|
||||
|
||||
impl Attention {
|
||||
fn forward(&self, xs: &Tensor, rotary: &Rotary, mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (b, l, d) = xs.dims3()?;
|
||||
let qkv = project(&self.wqkv, xs)?
|
||||
.reshape((b, l, 3, self.n_heads, self.head_dim))?
|
||||
.permute((2, 0, 3, 1, 4))?;
|
||||
// Rotary in the projection dtype, as under autocast; the scores
|
||||
// are then taken in f32 (see `scaled_attention`).
|
||||
let q = qkv.get(0)?;
|
||||
let k = qkv.get(1)?;
|
||||
let v = qkv.get(2)?.contiguous()?;
|
||||
let (q, k) = rotary.apply(&q, &k, l)?;
|
||||
let ctx = scaled_attention(&q, &k, &v, mask, self.head_dim)?;
|
||||
let ctx = ctx.transpose(1, 2)?.reshape((b, l, d))?;
|
||||
Ok(ctx.apply(&self.wo)?.to_dtype(DType::F32)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// A batch packed without padding for the flash-attention path: every
|
||||
/// row's tokens back to back, `(total, …)`. Linear layers then do no work
|
||||
/// for padding, and flash varlen keeps rows from attending to each other.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub struct Packed {
|
||||
/// `(b + 1,)` u32 cumulative row lengths, as flash varlen takes them.
|
||||
pub cu_seqlens: Tensor,
|
||||
/// Longest row.
|
||||
pub max_len: usize,
|
||||
/// `(total,)` u32: each token's position within its own row.
|
||||
pub positions: Tensor,
|
||||
/// Where each row starts in the packed order.
|
||||
pub offsets: Vec<usize>,
|
||||
}
|
||||
|
||||
#[cfg(feature = "flash-attn")]
|
||||
impl Packed {
|
||||
pub fn new(lens: &[usize], device: &Device) -> Result<Self> {
|
||||
let mut cu = Vec::with_capacity(lens.len() + 1);
|
||||
let mut offsets = Vec::with_capacity(lens.len());
|
||||
let mut positions = Vec::with_capacity(lens.iter().sum());
|
||||
let mut total = 0u32;
|
||||
cu.push(0u32);
|
||||
for &n in lens {
|
||||
offsets.push(total as usize);
|
||||
positions.extend(0..n as u32);
|
||||
total += n as u32;
|
||||
cu.push(total);
|
||||
}
|
||||
let n = cu.len();
|
||||
Ok(Self {
|
||||
cu_seqlens: Tensor::from_vec(cu, n, device)?,
|
||||
max_len: lens.iter().copied().max().unwrap_or(0),
|
||||
positions: Tensor::from_vec(positions, total as usize, device)?,
|
||||
offsets,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// Bidirectional flash attention over a packed batch. `q`, `k`, `v` are
|
||||
/// `(total, h, dh)`; `half_window` limits each query to keys within that
|
||||
/// distance on either side (ModernBERT's sliding layers), `None` attends
|
||||
/// the whole row. Scores are accumulated in f32 inside the kernel.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub(crate) fn flash_attention(
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
v: &Tensor,
|
||||
packed: &Packed,
|
||||
half_window: Option<usize>,
|
||||
) -> Result<Tensor> {
|
||||
let head_dim = q.dim(2)?;
|
||||
let scale = (head_dim as f64).powf(-0.5) as f32;
|
||||
Ok(candle_flash_attn::flash_attn_varlen_windowed(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
&packed.cu_seqlens,
|
||||
&packed.cu_seqlens,
|
||||
packed.max_len,
|
||||
packed.max_len,
|
||||
scale,
|
||||
half_window,
|
||||
half_window,
|
||||
)?)
|
||||
}
|
||||
|
||||
/// Split a packed `(total, 3·d)` qkv projection into contiguous
|
||||
/// `(total, h, dh)` q, k and v.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub(crate) fn split_qkv(
|
||||
qkv: &Tensor,
|
||||
n_heads: usize,
|
||||
head_dim: usize,
|
||||
) -> Result<(Tensor, Tensor, Tensor)> {
|
||||
let t = qkv.dim(0)?;
|
||||
let qkv = qkv.reshape((t, 3, n_heads, head_dim))?;
|
||||
let part = |i: usize| -> Result<Tensor> { Ok(qkv.narrow(1, i, 1)?.squeeze(1)?.contiguous()?) };
|
||||
Ok((part(0)?, part(1)?, part(2)?))
|
||||
}
|
||||
|
||||
#[cfg(feature = "flash-attn")]
|
||||
impl Rotary {
|
||||
/// Rotate packed `(total, h, dh)` q and k by each token's position.
|
||||
fn apply_packed(&self, q: &Tensor, k: &Tensor, positions: &Tensor) -> Result<(Tensor, Tensor)> {
|
||||
let cos = self.cos.index_select(positions, 0)?;
|
||||
let sin = self.sin.index_select(positions, 0)?;
|
||||
let rot = |x: &Tensor| -> Result<Tensor> {
|
||||
Ok(candle_nn::rotary_emb::rope_thd(&x.unsqueeze(0)?, &cos, &sin)?.squeeze(0)?)
|
||||
};
|
||||
Ok((rot(q)?, rot(k)?))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "flash-attn")]
|
||||
impl Attention {
|
||||
/// `xs` is the packed f32 residual stream `(total, d)`.
|
||||
fn forward_packed(
|
||||
&self,
|
||||
xs: &Tensor,
|
||||
rotary: &Rotary,
|
||||
packed: &Packed,
|
||||
half_window: usize,
|
||||
) -> Result<Tensor> {
|
||||
let (t, d) = xs.dims2()?;
|
||||
let (q, k, v) = split_qkv(&project(&self.wqkv, xs)?, self.n_heads, self.head_dim)?;
|
||||
let (q, k) = rotary.apply_packed(&q, &k, &packed.positions)?;
|
||||
let window = match self.kind {
|
||||
LayerKind::Sliding => Some(half_window),
|
||||
LayerKind::Global => None,
|
||||
};
|
||||
let ctx = flash_attention(&q, &k, &v, packed, window)?.reshape((t, d))?;
|
||||
Ok(ctx.apply(&self.wo)?.to_dtype(DType::F32)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// `softmax(q kᵀ / √d + mask) v`, scores and softmax in f32.
|
||||
///
|
||||
/// `q`, `k`, `v` are `(b, h, l, d)`; `mask` is an f32 additive mask
|
||||
/// broadcastable to `(b, h, l, l)`, or `None` when nothing is masked —
|
||||
/// which saves a bandwidth-bound pass over the `(b, h, l, l)` scores.
|
||||
pub(crate) fn scaled_attention(
|
||||
q: &Tensor,
|
||||
k: &Tensor,
|
||||
v: &Tensor,
|
||||
mask: Option<&Tensor>,
|
||||
head_dim: usize,
|
||||
) -> Result<Tensor> {
|
||||
let dtype = v.dtype();
|
||||
let scale = (head_dim as f64).powf(-0.5);
|
||||
// Scores in f32 end to end: a bf16 `q kᵀ` is rounded to 8 mantissa
|
||||
// bits before the softmax ever sees it, which on this model's
|
||||
// logit-sized scores is most of the drift from the reference's GPU
|
||||
// path (whose SDPA kernels keep scores in f32).
|
||||
// The scale goes on q, not on the `(b, h, l, l)` scores: the same
|
||||
// product, on a tensor `l / d` times smaller.
|
||||
let scores = (q.to_dtype(DType::F32)? * scale)?.matmul(&k.to_dtype(DType::F32)?.t()?)?;
|
||||
let scores = match mask {
|
||||
Some(mask) => scores.broadcast_add(mask)?,
|
||||
None => scores,
|
||||
};
|
||||
let probs = candle_nn::ops::softmax_last_dim(&scores)?;
|
||||
Ok(probs.to_dtype(dtype)?.matmul(v)?)
|
||||
}
|
||||
|
||||
struct Mlp {
|
||||
wi: Linear,
|
||||
wo: Linear,
|
||||
}
|
||||
|
||||
impl Module for Mlp {
|
||||
/// f32 in, f32 out; the projections run in the weights' dtype.
|
||||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||||
let xs = project(&self.wi, xs)?;
|
||||
let parts = xs.chunk(2, D::Minus1)?;
|
||||
(parts[0].gelu_erf()? * &parts[1])?
|
||||
.apply(&self.wo)?
|
||||
.to_dtype(DType::F32)
|
||||
}
|
||||
}
|
||||
|
||||
struct Layer {
|
||||
attn_norm: Option<LayerNorm>,
|
||||
attn: Attention,
|
||||
mlp_norm: LayerNorm,
|
||||
mlp: Mlp,
|
||||
}
|
||||
|
||||
/// The two additive masks one forward needs, both f32. Either is `None`
|
||||
/// when it would mask nothing: `global` for a batch with no padding (every
|
||||
/// single-question request), `sliding` additionally when no row is longer
|
||||
/// than the window reaches.
|
||||
pub struct Masks {
|
||||
/// `(b, 1, 1, l)`: padded keys masked. Global layers and the
|
||||
/// decision head use this one.
|
||||
pub global: Option<Tensor>,
|
||||
/// `(b, 1, l, l)`: padded keys, plus keys outside the window.
|
||||
pub sliding: Option<Tensor>,
|
||||
}
|
||||
|
||||
impl Masks {
|
||||
/// Build both masks from per-row valid lengths. Rows are
|
||||
/// right-padded: positions `>= lens[i]` of row `i` are padding.
|
||||
pub fn new(
|
||||
lens: &[usize],
|
||||
seq_len: usize,
|
||||
half_window: usize,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let b = lens.len();
|
||||
let padded = lens.iter().any(|&n| n < seq_len);
|
||||
let windowed = seq_len > half_window + 1;
|
||||
let global = if padded {
|
||||
let mut pad = Vec::with_capacity(b * seq_len);
|
||||
for &n in lens {
|
||||
pad.extend((0..seq_len).map(|j| if j < n { 0f32 } else { MASKED }));
|
||||
}
|
||||
Some(Tensor::from_vec(pad, (b, 1, 1, seq_len), device)?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
if !windowed {
|
||||
let sliding = global.clone();
|
||||
return Ok(Self { global, sliding });
|
||||
}
|
||||
// Built on the device from two aranges, not on the host: at
|
||||
// 1024 tokens a host-built window is a 4 MB upload per forward.
|
||||
let pos = Tensor::arange(0u32, seq_len as u32, device)?.to_dtype(DType::F32)?;
|
||||
let dist = pos
|
||||
.reshape((seq_len, 1))?
|
||||
.broadcast_sub(&pos.reshape((1, seq_len))?)?
|
||||
.abs()?;
|
||||
let outside = dist.gt(half_window as f64)?;
|
||||
let zeros = Tensor::zeros((seq_len, seq_len), DType::F32, device)?;
|
||||
let masked = Tensor::full(MASKED, (seq_len, seq_len), device)?;
|
||||
let window = outside
|
||||
.where_cond(&masked, &zeros)?
|
||||
.reshape((1, 1, seq_len, seq_len))?;
|
||||
let sliding = match &global {
|
||||
Some(g) => g.broadcast_add(&window)?,
|
||||
None => window,
|
||||
};
|
||||
Ok(Self {
|
||||
global,
|
||||
sliding: Some(sliding),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub struct ModernBert {
|
||||
embeddings: Embedding,
|
||||
embed_norm: LayerNorm,
|
||||
layers: Vec<Layer>,
|
||||
final_norm: LayerNorm,
|
||||
global_rotary: Rotary,
|
||||
sliding_rotary: Rotary,
|
||||
half_window: usize,
|
||||
}
|
||||
|
||||
impl ModernBert {
|
||||
/// Load from a VarBuilder rooted at the encoder's prefix (`encoder`
|
||||
/// in a Laya checkpoint, `model` in a stock ModernBERT one). Linear
|
||||
/// weights load in the VarBuilder's dtype; embeddings and norms in
|
||||
/// f32.
|
||||
pub fn load(vb: VarBuilder, cfg: &EncoderConfig) -> Result<Self> {
|
||||
let d = cfg.hidden_size;
|
||||
let f32_vb = vb.clone().set_dtype(DType::F32);
|
||||
let embeddings = candle_nn::embedding(
|
||||
cfg.vocab_size,
|
||||
d,
|
||||
f32_vb.pp("embeddings").pp("tok_embeddings"),
|
||||
)
|
||||
.context("encoder.embeddings.tok_embeddings")?;
|
||||
let embed_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, f32_vb.pp("embeddings.norm"))?;
|
||||
let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
|
||||
for (i, kind) in cfg.layer_kinds.iter().copied().enumerate() {
|
||||
let lvb = vb.pp(format!("layers.{i}"));
|
||||
let norm_vb = f32_vb.pp(format!("layers.{i}"));
|
||||
// ModernBERT builds layer 0's attn_norm as an identity and
|
||||
// saves nothing for it; every other layer must have one.
|
||||
let attn_norm = if lvb.contains_tensor("attn_norm.weight") {
|
||||
Some(fused_layer_norm_no_bias(
|
||||
d,
|
||||
cfg.norm_eps,
|
||||
norm_vb.pp("attn_norm"),
|
||||
)?)
|
||||
} else {
|
||||
ensure!(
|
||||
i == 0,
|
||||
"encoder layer {i} has no attn_norm.weight — wrong weight prefix, \
|
||||
or not a ModernBERT checkpoint"
|
||||
);
|
||||
None
|
||||
};
|
||||
let attn = Attention {
|
||||
wqkv: candle_nn::linear_no_bias(d, 3 * d, lvb.pp("attn.Wqkv"))?,
|
||||
wo: candle_nn::linear_no_bias(d, d, lvb.pp("attn.Wo"))?,
|
||||
n_heads: cfg.num_attention_heads,
|
||||
head_dim: cfg.head_dim(),
|
||||
kind,
|
||||
};
|
||||
let mlp_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, norm_vb.pp("mlp_norm"))?;
|
||||
let mlp = Mlp {
|
||||
wi: candle_nn::linear_no_bias(d, 2 * cfg.intermediate_size, lvb.pp("mlp.Wi"))?,
|
||||
wo: candle_nn::linear_no_bias(cfg.intermediate_size, d, lvb.pp("mlp.Wo"))?,
|
||||
};
|
||||
layers.push(Layer {
|
||||
attn_norm,
|
||||
attn,
|
||||
mlp_norm,
|
||||
mlp,
|
||||
});
|
||||
}
|
||||
let final_norm = fused_layer_norm_no_bias(d, cfg.norm_eps, f32_vb.pp("final_norm"))?;
|
||||
let device = vb.device();
|
||||
Ok(Self {
|
||||
embeddings,
|
||||
embed_norm,
|
||||
layers,
|
||||
final_norm,
|
||||
global_rotary: Rotary::new(cfg, cfg.global_rope_theta, vb.dtype(), device)?,
|
||||
sliding_rotary: Rotary::new(cfg, cfg.local_rope_theta, vb.dtype(), device)?,
|
||||
half_window: cfg.local_attention / 2,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn half_window(&self) -> usize {
|
||||
self.half_window
|
||||
}
|
||||
|
||||
/// `ids` is `(b, l)` u32; returns the final hidden states `(b, l, d)`
|
||||
/// in f32.
|
||||
pub fn forward(&self, ids: &Tensor, masks: &Masks) -> Result<Tensor> {
|
||||
let mut xs = ids.apply(&self.embeddings)?.apply(&self.embed_norm)?;
|
||||
for layer in &self.layers {
|
||||
let normed = match &layer.attn_norm {
|
||||
Some(norm) => xs.apply(norm)?,
|
||||
None => xs.clone(),
|
||||
};
|
||||
let (rotary, mask) = match layer.attn.kind {
|
||||
LayerKind::Global => (&self.global_rotary, masks.global.as_ref()),
|
||||
LayerKind::Sliding => (&self.sliding_rotary, masks.sliding.as_ref()),
|
||||
};
|
||||
xs = (&xs + layer.attn.forward(&normed, rotary, mask)?)?;
|
||||
let mlp_out = xs.apply(&layer.mlp_norm)?.apply(&layer.mlp)?;
|
||||
xs = (xs + mlp_out)?;
|
||||
}
|
||||
Ok(xs.apply(&self.final_norm)?)
|
||||
}
|
||||
|
||||
/// The flash path: `ids` is the packed `(total,)` u32 token stream;
|
||||
/// returns the packed final hidden states `(total, d)` in f32.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub fn forward_packed(&self, ids: &Tensor, packed: &Packed) -> Result<Tensor> {
|
||||
let mut xs = ids.apply(&self.embeddings)?.apply(&self.embed_norm)?;
|
||||
for layer in &self.layers {
|
||||
let normed = match &layer.attn_norm {
|
||||
Some(norm) => xs.apply(norm)?,
|
||||
None => xs.clone(),
|
||||
};
|
||||
let rotary = match layer.attn.kind {
|
||||
LayerKind::Global => &self.global_rotary,
|
||||
LayerKind::Sliding => &self.sliding_rotary,
|
||||
};
|
||||
xs = (&xs
|
||||
+ layer
|
||||
.attn
|
||||
.forward_packed(&normed, rotary, packed, self.half_window)?)?;
|
||||
let mlp_out = xs.apply(&layer.mlp_norm)?.apply(&layer.mlp)?;
|
||||
xs = (xs + mlp_out)?;
|
||||
}
|
||||
Ok(xs.apply(&self.final_norm)?)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,358 @@
|
||||
//! The two config files a Laya checkpoint ships, and what they mean.
|
||||
//!
|
||||
//! A checkpoint directory holds `rl_agent_config.json` (the decision
|
||||
//! model: head depth, token budgets, calibration temperatures) and
|
||||
//! `encoder/config.json` (the ModernBERT backbone, written by
|
||||
//! transformers 5). There is no top-level `config.json`.
|
||||
//!
|
||||
//! The encoder config is parsed here rather than through
|
||||
//! `candle_transformers::models::modernbert::Config` because
|
||||
//! transformers 5 moved the rotary bases into a nested
|
||||
//! `rope_parameters.{full_attention,sliding_attention}.rope_theta`
|
||||
//! block, which that struct cannot read; the older flat
|
||||
//! `global_rope_theta` / `local_rope_theta` fields are still honoured
|
||||
//! when a checkpoint carries them. mmBERT (the multilingual
|
||||
//! checkpoint's backbone) is the case that matters: both of its bases
|
||||
//! are 160,000, where the flat-field defaults would silently give the
|
||||
//! sliding layers 10,000.
|
||||
|
||||
use anyhow::{Context, Result, bail, ensure};
|
||||
use serde::Deserialize;
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::Path;
|
||||
|
||||
/// `rl_agent_config.json`.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct LayaConfig {
|
||||
/// The backbone the head was trained on, e.g.
|
||||
/// `answerdotai/ModernBERT-large`. Informational: the weights come
|
||||
/// from the checkpoint, the shape from `encoder/config.json`.
|
||||
#[serde(default)]
|
||||
pub encoder: Option<String>,
|
||||
#[serde(default = "default_head_layers")]
|
||||
pub head_layers: usize,
|
||||
#[serde(default = "default_max_len")]
|
||||
pub max_len: usize,
|
||||
#[serde(default = "default_head_max_len")]
|
||||
pub head_max_len: usize,
|
||||
/// One act head output per entry, plus one for "answer". The act
|
||||
/// head's width is `act_costs.len() + 1`.
|
||||
#[serde(default)]
|
||||
pub act_costs: BTreeMap<String, f64>,
|
||||
/// Per-question-type temperatures, indexed choice / score / noul.
|
||||
#[serde(default = "default_temperature")]
|
||||
pub temperature: Vec<f64>,
|
||||
/// Temperatures fitted per `"<type>:<cardinality bucket>"`.
|
||||
#[serde(default)]
|
||||
pub temperature_by_options: BTreeMap<String, f64>,
|
||||
/// Autocast dtype the reference runs on GPU (`"bf16"` or `"fp16"`).
|
||||
#[serde(default)]
|
||||
pub amp_dtype: Option<String>,
|
||||
}
|
||||
|
||||
fn default_head_layers() -> usize {
|
||||
2
|
||||
}
|
||||
fn default_max_len() -> usize {
|
||||
512
|
||||
}
|
||||
fn default_head_max_len() -> usize {
|
||||
192
|
||||
}
|
||||
fn default_temperature() -> Vec<f64> {
|
||||
vec![1.0, 1.0, 1.0]
|
||||
}
|
||||
|
||||
impl LayaConfig {
|
||||
pub fn n_act(&self) -> usize {
|
||||
self.act_costs.len() + 1
|
||||
}
|
||||
}
|
||||
|
||||
/// Whether an attention layer sees the whole sequence or a window.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum LayerKind {
|
||||
Global,
|
||||
Sliding,
|
||||
}
|
||||
|
||||
/// The ModernBERT backbone shape, resolved from `encoder/config.json`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct EncoderConfig {
|
||||
pub vocab_size: usize,
|
||||
pub hidden_size: usize,
|
||||
pub num_hidden_layers: usize,
|
||||
pub num_attention_heads: usize,
|
||||
pub intermediate_size: usize,
|
||||
pub max_position_embeddings: usize,
|
||||
pub norm_eps: f64,
|
||||
pub pad_token_id: u32,
|
||||
/// Full sliding-window width. A token attends to keys at distance
|
||||
/// `<= local_attention / 2` on either side.
|
||||
pub local_attention: usize,
|
||||
pub global_rope_theta: f64,
|
||||
pub local_rope_theta: f64,
|
||||
pub layer_kinds: Vec<LayerKind>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RawEncoderConfig {
|
||||
vocab_size: usize,
|
||||
hidden_size: usize,
|
||||
num_hidden_layers: usize,
|
||||
num_attention_heads: usize,
|
||||
intermediate_size: usize,
|
||||
max_position_embeddings: usize,
|
||||
#[serde(default)]
|
||||
norm_eps: Option<f64>,
|
||||
#[serde(default)]
|
||||
layer_norm_eps: Option<f64>,
|
||||
pad_token_id: u32,
|
||||
#[serde(default = "default_global_every")]
|
||||
global_attn_every_n_layers: usize,
|
||||
local_attention: usize,
|
||||
#[serde(default)]
|
||||
layer_types: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
rope_parameters: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
global_rope_theta: Option<f64>,
|
||||
#[serde(default)]
|
||||
local_rope_theta: Option<f64>,
|
||||
#[serde(default)]
|
||||
norm_bias: bool,
|
||||
#[serde(default)]
|
||||
attention_bias: bool,
|
||||
#[serde(default)]
|
||||
mlp_bias: bool,
|
||||
#[serde(default)]
|
||||
hidden_activation: Option<String>,
|
||||
}
|
||||
|
||||
fn default_global_every() -> usize {
|
||||
3
|
||||
}
|
||||
|
||||
/// transformers' ModernBERT defaults for checkpoints that predate
|
||||
/// `rope_parameters` and also omit the flat fields.
|
||||
const DEFAULT_GLOBAL_ROPE_THETA: f64 = 160_000.0;
|
||||
const DEFAULT_LOCAL_ROPE_THETA: f64 = 10_000.0;
|
||||
|
||||
impl EncoderConfig {
|
||||
pub fn from_json(text: &str) -> Result<Self> {
|
||||
let raw: RawEncoderConfig =
|
||||
serde_json::from_str(text).context("parse ModernBERT encoder config")?;
|
||||
|
||||
// Variants this port does not implement are refused, not
|
||||
// approximated: each would load without error and answer wrong.
|
||||
ensure!(
|
||||
!raw.norm_bias && !raw.attention_bias && !raw.mlp_bias,
|
||||
"ModernBERT with norm/attention/mlp biases is not supported"
|
||||
);
|
||||
if let Some(act) = raw.hidden_activation.as_deref() {
|
||||
ensure!(
|
||||
act == "gelu",
|
||||
"ModernBERT hidden_activation {act:?} is not supported (only exact \"gelu\")"
|
||||
);
|
||||
}
|
||||
ensure!(
|
||||
raw.hidden_size.is_multiple_of(raw.num_attention_heads),
|
||||
"hidden_size {} is not divisible by num_attention_heads {}",
|
||||
raw.hidden_size,
|
||||
raw.num_attention_heads
|
||||
);
|
||||
|
||||
let (global_rope_theta, local_rope_theta) = rope_thetas(&raw)?;
|
||||
|
||||
let layer_kinds = match &raw.layer_types {
|
||||
Some(types) => {
|
||||
ensure!(
|
||||
types.len() == raw.num_hidden_layers,
|
||||
"layer_types has {} entries for {} layers",
|
||||
types.len(),
|
||||
raw.num_hidden_layers
|
||||
);
|
||||
types
|
||||
.iter()
|
||||
.map(|t| match t.as_str() {
|
||||
"full_attention" => Ok(LayerKind::Global),
|
||||
"sliding_attention" => Ok(LayerKind::Sliding),
|
||||
other => bail!("unknown ModernBERT layer type {other:?}"),
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?
|
||||
}
|
||||
None => (0..raw.num_hidden_layers)
|
||||
.map(|i| {
|
||||
if i % raw.global_attn_every_n_layers == 0 {
|
||||
LayerKind::Global
|
||||
} else {
|
||||
LayerKind::Sliding
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
vocab_size: raw.vocab_size,
|
||||
hidden_size: raw.hidden_size,
|
||||
num_hidden_layers: raw.num_hidden_layers,
|
||||
num_attention_heads: raw.num_attention_heads,
|
||||
intermediate_size: raw.intermediate_size,
|
||||
max_position_embeddings: raw.max_position_embeddings,
|
||||
norm_eps: raw.norm_eps.or(raw.layer_norm_eps).unwrap_or(1e-5),
|
||||
pad_token_id: raw.pad_token_id,
|
||||
local_attention: raw.local_attention,
|
||||
global_rope_theta,
|
||||
local_rope_theta,
|
||||
layer_kinds,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn head_dim(&self) -> usize {
|
||||
self.hidden_size / self.num_attention_heads
|
||||
}
|
||||
}
|
||||
|
||||
/// Resolve the two rotary bases. `rope_parameters` (transformers 5)
|
||||
/// wins over the flat fields; within it, a per-layer-type block wins
|
||||
/// over a shared top-level `rope_theta`.
|
||||
fn rope_thetas(raw: &RawEncoderConfig) -> Result<(f64, f64)> {
|
||||
let mut global = raw.global_rope_theta.unwrap_or(DEFAULT_GLOBAL_ROPE_THETA);
|
||||
let mut local = raw.local_rope_theta.unwrap_or(DEFAULT_LOCAL_ROPE_THETA);
|
||||
if let Some(rope) = &raw.rope_parameters {
|
||||
let shared = rope.get("rope_theta").and_then(|v| v.as_f64());
|
||||
for (key, slot) in [
|
||||
("full_attention", &mut global),
|
||||
("sliding_attention", &mut local),
|
||||
] {
|
||||
let block = rope.get(key);
|
||||
if let Some(rope_type) = block
|
||||
.and_then(|b| b.get("rope_type"))
|
||||
.and_then(|v| v.as_str())
|
||||
{
|
||||
ensure!(
|
||||
rope_type == "default",
|
||||
"rope_type {rope_type:?} for {key} is not supported"
|
||||
);
|
||||
}
|
||||
let theta = block
|
||||
.and_then(|b| b.get("rope_theta"))
|
||||
.and_then(|v| v.as_f64())
|
||||
.or(shared);
|
||||
if let Some(theta) = theta {
|
||||
*slot = theta;
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok((global, local))
|
||||
}
|
||||
|
||||
/// Both configs of one checkpoint directory.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CheckpointConfig {
|
||||
pub laya: LayaConfig,
|
||||
pub encoder: EncoderConfig,
|
||||
}
|
||||
|
||||
impl CheckpointConfig {
|
||||
/// Read `rl_agent_config.json` and `encoder/config.json` from a
|
||||
/// checkpoint directory (the repo root, or one of its subfolders).
|
||||
pub fn from_dir(dir: &Path) -> Result<Self> {
|
||||
let laya_path = dir.join("rl_agent_config.json");
|
||||
let laya: LayaConfig = serde_json::from_str(
|
||||
&std::fs::read_to_string(&laya_path)
|
||||
.with_context(|| format!("read {}", laya_path.display()))?,
|
||||
)
|
||||
.with_context(|| format!("parse {}", laya_path.display()))?;
|
||||
let enc_path = dir.join("encoder").join("config.json");
|
||||
let encoder = EncoderConfig::from_json(
|
||||
&std::fs::read_to_string(&enc_path)
|
||||
.with_context(|| format!("read {}", enc_path.display()))?,
|
||||
)
|
||||
.with_context(|| format!("parse {}", enc_path.display()))?;
|
||||
ensure!(
|
||||
laya.temperature.len() == 3,
|
||||
"rl_agent_config temperature must have 3 entries (choice, score, noul), got {}",
|
||||
laya.temperature.len()
|
||||
);
|
||||
Ok(Self { laya, encoder })
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn base() -> serde_json::Value {
|
||||
serde_json::json!({
|
||||
"vocab_size": 100, "hidden_size": 64, "num_hidden_layers": 4,
|
||||
"num_attention_heads": 4, "intermediate_size": 96,
|
||||
"max_position_embeddings": 512, "norm_eps": 1e-5, "pad_token_id": 0,
|
||||
"global_attn_every_n_layers": 3, "local_attention": 16,
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nested_rope_parameters_override_flat_defaults() {
|
||||
let mut v = base();
|
||||
v["rope_parameters"] = serde_json::json!({
|
||||
"full_attention": {"rope_theta": 160000.0, "rope_type": "default"},
|
||||
"sliding_attention": {"rope_theta": 160000.0, "rope_type": "default"},
|
||||
});
|
||||
let c = EncoderConfig::from_json(&v.to_string()).unwrap();
|
||||
assert_eq!(c.global_rope_theta, 160_000.0);
|
||||
// The mmBERT case: the flat default would have been 10,000.
|
||||
assert_eq!(c.local_rope_theta, 160_000.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn flat_fields_are_read_without_rope_parameters() {
|
||||
let mut v = base();
|
||||
v["global_rope_theta"] = serde_json::json!(80000.0);
|
||||
v["local_rope_theta"] = serde_json::json!(5000.0);
|
||||
let c = EncoderConfig::from_json(&v.to_string()).unwrap();
|
||||
assert_eq!((c.global_rope_theta, c.local_rope_theta), (80000.0, 5000.0));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn layer_kinds_follow_every_n_when_layer_types_absent() {
|
||||
let c = EncoderConfig::from_json(&base().to_string()).unwrap();
|
||||
assert_eq!(
|
||||
c.layer_kinds,
|
||||
vec![
|
||||
LayerKind::Global,
|
||||
LayerKind::Sliding,
|
||||
LayerKind::Sliding,
|
||||
LayerKind::Global
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_layer_types_win() {
|
||||
let mut v = base();
|
||||
v["layer_types"] = serde_json::json!([
|
||||
"sliding_attention",
|
||||
"full_attention",
|
||||
"full_attention",
|
||||
"sliding_attention"
|
||||
]);
|
||||
let c = EncoderConfig::from_json(&v.to_string()).unwrap();
|
||||
assert_eq!(c.layer_kinds[0], LayerKind::Sliding);
|
||||
assert_eq!(c.layer_kinds[1], LayerKind::Global);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn unsupported_variants_are_refused() {
|
||||
let mut v = base();
|
||||
v["attention_bias"] = serde_json::json!(true);
|
||||
assert!(EncoderConfig::from_json(&v.to_string()).is_err());
|
||||
let mut v = base();
|
||||
v["hidden_activation"] = serde_json::json!("gelu_pytorch_tanh");
|
||||
assert!(EncoderConfig::from_json(&v.to_string()).is_err());
|
||||
let mut v = base();
|
||||
v["rope_parameters"] =
|
||||
serde_json::json!({"full_attention": {"rope_type": "yarn", "rope_theta": 1.0}});
|
||||
assert!(EncoderConfig::from_json(&v.to_string()).is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
//! The decision head that sits on top of the encoder.
|
||||
//!
|
||||
//! Mirrors the reference `DecisionModel` (the `laya` SDK's
|
||||
//! `common.py`):
|
||||
//!
|
||||
//! 1. a per-question-type embedding added to every position;
|
||||
//! 2. `head_layers` × PyTorch `nn.TransformerEncoderLayer` with
|
||||
//! `norm_first=True`: pre-norm self-attention through a fused
|
||||
//! `in_proj` with bias, then a pre-norm feed-forward whose
|
||||
//! activation is **ReLU** — PyTorch's default, which the reference
|
||||
//! never overrides;
|
||||
//! 3. `scorer`, applied at each option's `[MASK]` marker: LayerNorm →
|
||||
//! Linear → exact GELU → Linear to one logit;
|
||||
//! 4. `act_head`, over the first position's hidden state plus four
|
||||
//! summary features of the (uncalibrated) answer distribution.
|
||||
|
||||
use anyhow::Result;
|
||||
use candle_core::{DType, Tensor};
|
||||
use candle_nn::{Embedding, LayerNorm, Linear, VarBuilder};
|
||||
|
||||
use super::backbone::{project, scaled_attention};
|
||||
|
||||
/// PyTorch's `LayerNorm` default.
|
||||
const TORCH_LN_EPS: f64 = 1e-5;
|
||||
|
||||
struct HeadLayer {
|
||||
norm1: LayerNorm,
|
||||
in_proj: Linear,
|
||||
out_proj: Linear,
|
||||
norm2: LayerNorm,
|
||||
linear1: Linear,
|
||||
linear2: Linear,
|
||||
n_heads: usize,
|
||||
head_dim: usize,
|
||||
}
|
||||
|
||||
impl HeadLayer {
|
||||
fn load(vb: VarBuilder, d: usize) -> Result<Self> {
|
||||
// `nn.TransformerEncoder` is built with `nhead = max(1, d // 64)`
|
||||
// and `dim_feedforward = 4 * d`.
|
||||
let n_heads = (d / 64).max(1);
|
||||
let f32_vb = vb.clone().set_dtype(DType::F32);
|
||||
Ok(Self {
|
||||
norm1: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("norm1"))?,
|
||||
in_proj: linear_torch_mha_in_proj(d, vb.pp("self_attn"))?,
|
||||
out_proj: candle_nn::linear(d, d, vb.pp("self_attn.out_proj"))?,
|
||||
norm2: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("norm2"))?,
|
||||
linear1: candle_nn::linear(d, 4 * d, vb.pp("linear1"))?,
|
||||
linear2: candle_nn::linear(4 * d, d, vb.pp("linear2"))?,
|
||||
n_heads,
|
||||
head_dim: d / n_heads,
|
||||
})
|
||||
}
|
||||
|
||||
/// `xs` is the f32 residual stream; projections run in the weights'
|
||||
/// dtype and join it back in f32, as under autocast.
|
||||
fn forward(&self, xs: &Tensor, key_mask: Option<&Tensor>) -> Result<Tensor> {
|
||||
let (b, l, d) = xs.dims3()?;
|
||||
let qkv = project(&self.in_proj, &xs.apply(&self.norm1)?)?
|
||||
.reshape((b, l, 3, self.n_heads, self.head_dim))?
|
||||
.permute((2, 0, 3, 1, 4))?;
|
||||
let q = qkv.get(0)?.contiguous()?;
|
||||
let k = qkv.get(1)?.contiguous()?;
|
||||
let v = qkv.get(2)?.contiguous()?;
|
||||
let ctx = scaled_attention(&q, &k, &v, key_mask, self.head_dim)?
|
||||
.transpose(1, 2)?
|
||||
.reshape((b, l, d))?;
|
||||
let xs = (xs + ctx.apply(&self.out_proj)?.to_dtype(DType::F32)?)?;
|
||||
let ff = project(&self.linear1, &xs.apply(&self.norm2)?)?
|
||||
.relu()?
|
||||
.apply(&self.linear2)?
|
||||
.to_dtype(DType::F32)?;
|
||||
Ok((xs + ff)?)
|
||||
}
|
||||
|
||||
/// The flash path over a packed `(total, d)` f32 stream.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
fn forward_packed(&self, xs: &Tensor, packed: &super::backbone::Packed) -> Result<Tensor> {
|
||||
use super::backbone::{flash_attention, split_qkv};
|
||||
let (t, d) = xs.dims2()?;
|
||||
let qkv = project(&self.in_proj, &xs.apply(&self.norm1)?)?;
|
||||
let (q, k, v) = split_qkv(&qkv, self.n_heads, self.head_dim)?;
|
||||
let ctx = flash_attention(&q, &k, &v, packed, None)?.reshape((t, d))?;
|
||||
let xs = (xs + ctx.apply(&self.out_proj)?.to_dtype(DType::F32)?)?;
|
||||
let ff = project(&self.linear1, &xs.apply(&self.norm2)?)?
|
||||
.relu()?
|
||||
.apply(&self.linear2)?
|
||||
.to_dtype(DType::F32)?;
|
||||
Ok((xs + ff)?)
|
||||
}
|
||||
}
|
||||
|
||||
/// `nn.MultiheadAttention` stores q/k/v as one `in_proj_weight`
|
||||
/// `(3d, d)` and `in_proj_bias` `(3d)`, rows ordered q then k then v —
|
||||
/// exactly the layout one fused `Linear` wants.
|
||||
fn linear_torch_mha_in_proj(d: usize, vb: VarBuilder) -> Result<Linear> {
|
||||
let w = vb.get((3 * d, d), "in_proj_weight")?;
|
||||
let b = vb.get(3 * d, "in_proj_bias")?;
|
||||
Ok(Linear::new(w, Some(b)))
|
||||
}
|
||||
|
||||
pub struct DecisionHead {
|
||||
type_emb: Embedding,
|
||||
layers: Vec<HeadLayer>,
|
||||
scorer_norm: LayerNorm,
|
||||
scorer_fc: Linear,
|
||||
scorer_out: Linear,
|
||||
act_fc: Linear,
|
||||
act_out: Linear,
|
||||
}
|
||||
|
||||
impl DecisionHead {
|
||||
/// `vb` is the checkpoint root; the head's tensors are top-level.
|
||||
/// Linear weights load in `vb`'s dtype; the type embedding and norms
|
||||
/// in f32.
|
||||
pub fn load(vb: VarBuilder, d: usize, head_layers: usize, n_act: usize) -> Result<Self> {
|
||||
let layers = (0..head_layers)
|
||||
.map(|i| HeadLayer::load(vb.pp(format!("head.layers.{i}")), d))
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
let f32_vb = vb.clone().set_dtype(DType::F32);
|
||||
Ok(Self {
|
||||
type_emb: candle_nn::embedding(3, d, f32_vb.pp("type_emb"))?,
|
||||
layers,
|
||||
scorer_norm: candle_nn::layer_norm(d, TORCH_LN_EPS, f32_vb.pp("scorer.0"))?,
|
||||
scorer_fc: candle_nn::linear(d, d, vb.pp("scorer.1"))?,
|
||||
scorer_out: candle_nn::linear(d, 1, vb.pp("scorer.3"))?,
|
||||
act_fc: candle_nn::linear(d + 4, 256, vb.pp("act_head.0"))?,
|
||||
act_out: candle_nn::linear(256, n_act, vb.pp("act_head.2"))?,
|
||||
})
|
||||
}
|
||||
|
||||
/// Type embedding plus the transformer layers. `h` is the encoder
|
||||
/// output `(b, l, d)`, `qtype` is `(b,)` u32.
|
||||
pub fn contextualise(
|
||||
&self,
|
||||
h: &Tensor,
|
||||
qtype: &Tensor,
|
||||
key_mask: Option<&Tensor>,
|
||||
) -> Result<Tensor> {
|
||||
let mut h = h.broadcast_add(&qtype.apply(&self.type_emb)?.unsqueeze(1)?)?;
|
||||
for layer in &self.layers {
|
||||
h = layer.forward(&h, key_mask)?;
|
||||
}
|
||||
Ok(h)
|
||||
}
|
||||
|
||||
/// The flash path: `h` is the packed encoder output `(total, d)` and
|
||||
/// `qtype` the per-token question type `(total,)` u32.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
pub fn contextualise_packed(
|
||||
&self,
|
||||
h: &Tensor,
|
||||
qtype: &Tensor,
|
||||
packed: &super::backbone::Packed,
|
||||
) -> Result<Tensor> {
|
||||
let mut h = (h + qtype.apply(&self.type_emb)?)?;
|
||||
for layer in &self.layers {
|
||||
h = layer.forward_packed(&h, packed)?;
|
||||
}
|
||||
Ok(h)
|
||||
}
|
||||
|
||||
/// Score gathered marker states `(b, k, d)` → `(b, k)` f32 logits.
|
||||
pub fn score(&self, markers: &Tensor) -> Result<Tensor> {
|
||||
let x = project(&self.scorer_fc, &markers.apply(&self.scorer_norm)?)?
|
||||
.gelu_erf()?
|
||||
.apply(&self.scorer_out)?;
|
||||
Ok(x.squeeze(2)?.to_dtype(DType::F32)?)
|
||||
}
|
||||
|
||||
/// Act logits `(b, n_act)` f32 from the pooled first position
|
||||
/// `(b, d)` and the four answer-distribution features `(b, 4)`.
|
||||
pub fn act(&self, pooled: &Tensor, feats: &Tensor) -> Result<Tensor> {
|
||||
let x = Tensor::cat(&[pooled, &feats.to_dtype(pooled.dtype())?], 1)?;
|
||||
let x = project(&self.act_fc, &x)?
|
||||
.gelu_erf()?
|
||||
.apply(&self.act_out)?;
|
||||
Ok(x.to_dtype(DType::F32)?)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
//! Laya — a non-autoregressive decision model (`convaiinnovations/laya`).
|
||||
//!
|
||||
//! Not a language model: a ModernBERT encoder plus a small decision
|
||||
//! head that scores the options of a typed question in one forward
|
||||
//! pass. The caller assembles each question into one token sequence
|
||||
//! (`[CLS] <type> question: <instructions> [SEP] [MASK] opt0 [MASK]
|
||||
//! opt1 … [SEP] state [SEP]`) and says where each option's `[MASK]`
|
||||
//! marker sits; this module returns one raw score per marker plus the
|
||||
//! act head's distribution. Calibration (temperature scaling) and
|
||||
//! answer shaping happen after, CPU-side, from the checkpoint's
|
||||
//! `rl_agent_config.json`.
|
||||
//!
|
||||
//! The reference implementation is the `laya` SDK's `DecisionModel`
|
||||
//! (`laya/common.py`), not the `rl_common.py` shipped in the model
|
||||
//! repository, which has diverged from it.
|
||||
//!
|
||||
//! - [`config`] — the two config files a checkpoint ships.
|
||||
//! - [`backbone`] — ModernBERT, vendored so it runs outside f32.
|
||||
//! - [`head`] — type embedding, the head's transformer layers, the
|
||||
//! marker scorer and the act head.
|
||||
//! - [`model`] — the composition, batch in, [`DecisionRow`]s out.
|
||||
|
||||
pub mod backbone;
|
||||
pub mod config;
|
||||
pub mod head;
|
||||
pub mod model;
|
||||
|
||||
pub use config::CheckpointConfig;
|
||||
pub use model::LayaModel;
|
||||
|
||||
/// Question type, as the checkpoint's type embedding indexes it.
|
||||
pub const QTYPE_CHOICE: u8 = 0;
|
||||
pub const QTYPE_SCORE: u8 = 1;
|
||||
pub const QTYPE_NOUL: u8 = 2;
|
||||
|
||||
/// One assembled question: its token ids, where each option's
|
||||
/// `[MASK]` marker sits, and its type.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DecisionSeq {
|
||||
pub input_ids: Vec<u32>,
|
||||
pub markers: Vec<usize>,
|
||||
/// [`QTYPE_CHOICE`], [`QTYPE_SCORE`] or [`QTYPE_NOUL`].
|
||||
pub qtype: u8,
|
||||
}
|
||||
|
||||
/// Sequences scored in one forward. One request's questions normally
|
||||
/// travel together, so the state they share is encoded in one pass.
|
||||
#[derive(Debug, Clone, Default, PartialEq)]
|
||||
pub struct DecisionBatch {
|
||||
pub seqs: Vec<DecisionSeq>,
|
||||
}
|
||||
|
||||
/// One sequence's result, CPU-side.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DecisionRow {
|
||||
/// One uncalibrated score per marker, in marker order.
|
||||
pub logits: Vec<f32>,
|
||||
/// Softmax over the act head's outputs. Index 0 is what the
|
||||
/// reference reports as `action.act_probability`.
|
||||
pub act: Vec<f32>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
@@ -0,0 +1,260 @@
|
||||
//! A whole Laya checkpoint: encoder + decision head, batch in, raw
|
||||
//! scores out.
|
||||
|
||||
use anyhow::{Context, Result, ensure};
|
||||
use candle_core::{DType, Device, Tensor};
|
||||
use candle_nn::VarBuilder;
|
||||
use std::path::Path;
|
||||
|
||||
use super::backbone::{Masks, ModernBert};
|
||||
use super::config::CheckpointConfig;
|
||||
use super::head::DecisionHead;
|
||||
use super::{DecisionBatch, DecisionRow};
|
||||
|
||||
pub struct LayaModel {
|
||||
encoder: ModernBert,
|
||||
head: DecisionHead,
|
||||
config: CheckpointConfig,
|
||||
device: Device,
|
||||
/// Precision of the linear layers (the load dtype).
|
||||
dtype: DType,
|
||||
}
|
||||
|
||||
/// A validated batch, flattened: what both forward paths start from.
|
||||
struct Flat {
|
||||
lens: Vec<usize>,
|
||||
qtypes: Vec<u32>,
|
||||
/// Each row's marker positions, padded to `k_max` with 0.
|
||||
marker_idx: Vec<u32>,
|
||||
k_max: usize,
|
||||
seq_len: usize,
|
||||
}
|
||||
|
||||
impl LayaModel {
|
||||
/// Load one checkpoint directory (the repo root, or one of its
|
||||
/// subfolders): `rl_agent_config.json`, `encoder/config.json` and
|
||||
/// `model.safetensors`.
|
||||
pub fn load(dir: &Path, device: &Device, dtype: DType) -> Result<Self> {
|
||||
let config = CheckpointConfig::from_dir(dir)?;
|
||||
let weights = dir.join("model.safetensors");
|
||||
// SAFETY: the file is memory-mapped read-only for the lifetime
|
||||
// of the VarBuilder; nothing writes to a checkpoint in place.
|
||||
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[&weights], dtype, device) }
|
||||
.with_context(|| format!("mmap {}", weights.display()))?;
|
||||
Self::from_var_builder(vb, config, device)
|
||||
}
|
||||
|
||||
pub fn from_var_builder(
|
||||
vb: VarBuilder,
|
||||
config: CheckpointConfig,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let dtype = vb.dtype();
|
||||
let encoder = ModernBert::load(vb.pp("encoder"), &config.encoder)?;
|
||||
let head = DecisionHead::load(
|
||||
vb.clone(),
|
||||
config.encoder.hidden_size,
|
||||
config.laya.head_layers,
|
||||
config.laya.n_act(),
|
||||
)?;
|
||||
Ok(Self {
|
||||
encoder,
|
||||
head,
|
||||
config,
|
||||
device: device.clone(),
|
||||
dtype,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &CheckpointConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Precision the linear layers were loaded in.
|
||||
pub fn dtype(&self) -> DType {
|
||||
self.dtype
|
||||
}
|
||||
|
||||
/// One forward over every sequence in `batch`.
|
||||
///
|
||||
/// On a CUDA device in reduced precision, in a build with the
|
||||
/// `flash-attn` feature, the batch is packed without padding and
|
||||
/// attention runs in flash varlen kernels. Otherwise rows are
|
||||
/// right-padded to the longest, as the reference collates them, and
|
||||
/// attention is eager; padding changes nothing about an unpadded
|
||||
/// row's result.
|
||||
pub fn forward(&self, batch: &DecisionBatch) -> Result<Vec<DecisionRow>> {
|
||||
let flat = self.validate(batch)?;
|
||||
#[cfg(feature = "flash-attn")]
|
||||
if self.device.is_cuda() && self.dtype != DType::F32 && flash_enabled() {
|
||||
return self.forward_packed(batch, &flat);
|
||||
}
|
||||
self.forward_padded(batch, &flat)
|
||||
}
|
||||
|
||||
fn validate(&self, batch: &DecisionBatch) -> Result<Flat> {
|
||||
let seqs = &batch.seqs;
|
||||
ensure!(!seqs.is_empty(), "empty decision batch");
|
||||
let max_pos = self.config.encoder.max_position_embeddings;
|
||||
let seq_len = seqs.iter().map(|s| s.input_ids.len()).max().unwrap_or(0);
|
||||
ensure!(seq_len > 0, "decision batch has an empty sequence");
|
||||
ensure!(
|
||||
seq_len <= max_pos,
|
||||
"sequence of {seq_len} tokens exceeds max_position_embeddings {max_pos}"
|
||||
);
|
||||
let vocab = self.config.encoder.vocab_size as u32;
|
||||
let k_max = seqs.iter().map(|s| s.markers.len()).max().unwrap_or(0);
|
||||
ensure!(
|
||||
k_max > 0,
|
||||
"decision batch has a sequence with no option markers"
|
||||
);
|
||||
let mut lens = Vec::with_capacity(seqs.len());
|
||||
let mut qtypes = Vec::with_capacity(seqs.len());
|
||||
let mut marker_idx = Vec::with_capacity(seqs.len() * k_max);
|
||||
for s in seqs {
|
||||
let n = s.input_ids.len();
|
||||
ensure!(n > 0, "decision batch has an empty sequence");
|
||||
ensure!(!s.markers.is_empty(), "sequence has no option markers");
|
||||
ensure!(s.qtype < 3, "question type {} out of range", s.qtype);
|
||||
if let Some(&bad) = s.input_ids.iter().find(|&&t| t >= vocab) {
|
||||
anyhow::bail!("token id {bad} outside vocabulary of {vocab}");
|
||||
}
|
||||
if let Some(&bad) = s.markers.iter().find(|&&m| m >= n) {
|
||||
anyhow::bail!("marker position {bad} outside a {n}-token sequence");
|
||||
}
|
||||
lens.push(n);
|
||||
qtypes.push(s.qtype as u32);
|
||||
marker_idx.extend(s.markers.iter().map(|&m| m as u32));
|
||||
// Absent slots gather position 0; their scores are discarded.
|
||||
marker_idx.extend(std::iter::repeat_n(0u32, k_max - s.markers.len()));
|
||||
}
|
||||
Ok(Flat {
|
||||
lens,
|
||||
qtypes,
|
||||
marker_idx,
|
||||
k_max,
|
||||
seq_len,
|
||||
})
|
||||
}
|
||||
|
||||
/// The eager path: rows right-padded to `(b, seq_len)`.
|
||||
fn forward_padded(&self, batch: &DecisionBatch, flat: &Flat) -> Result<Vec<DecisionRow>> {
|
||||
let b = batch.seqs.len();
|
||||
let seq_len = flat.seq_len;
|
||||
let pad = self.config.encoder.pad_token_id;
|
||||
let mut ids = Vec::with_capacity(b * seq_len);
|
||||
for s in &batch.seqs {
|
||||
ids.extend_from_slice(&s.input_ids);
|
||||
ids.extend(std::iter::repeat_n(pad, seq_len - s.input_ids.len()));
|
||||
}
|
||||
let ids = Tensor::from_vec(ids, (b, seq_len), &self.device)?;
|
||||
let masks = Masks::new(
|
||||
&flat.lens,
|
||||
seq_len,
|
||||
self.encoder.half_window(),
|
||||
&self.device,
|
||||
)?;
|
||||
let qtype = Tensor::from_vec(flat.qtypes.clone(), b, &self.device)?;
|
||||
|
||||
let h = self.encoder.forward(&ids, &masks)?;
|
||||
let h = self.head.contextualise(&h, &qtype, masks.global.as_ref())?;
|
||||
let d = self.config.encoder.hidden_size;
|
||||
let idx = Tensor::from_vec(flat.marker_idx.clone(), (b, flat.k_max, 1), &self.device)?
|
||||
.broadcast_as((b, flat.k_max, d))?
|
||||
.contiguous()?;
|
||||
let marker_h = h.gather(&idx, 1)?;
|
||||
let pooled = h.narrow(1, 0, 1)?.squeeze(1)?;
|
||||
self.finish(batch, &marker_h, &pooled)
|
||||
}
|
||||
|
||||
/// The flash path: tokens packed without padding, `(total, d)`.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
fn forward_packed(&self, batch: &DecisionBatch, flat: &Flat) -> Result<Vec<DecisionRow>> {
|
||||
use super::backbone::Packed;
|
||||
let b = batch.seqs.len();
|
||||
let packed = Packed::new(&flat.lens, &self.device)?;
|
||||
let total: usize = flat.lens.iter().sum();
|
||||
let mut ids = Vec::with_capacity(total);
|
||||
let mut qtype_tok = Vec::with_capacity(total);
|
||||
for (s, &q) in batch.seqs.iter().zip(&flat.qtypes) {
|
||||
ids.extend_from_slice(&s.input_ids);
|
||||
qtype_tok.extend(std::iter::repeat_n(q, s.input_ids.len()));
|
||||
}
|
||||
let ids = Tensor::from_vec(ids, total, &self.device)?;
|
||||
let qtype_tok = Tensor::from_vec(qtype_tok, total, &self.device)?;
|
||||
|
||||
let h = self.encoder.forward_packed(&ids, &packed)?;
|
||||
let h = self.head.contextualise_packed(&h, &qtype_tok, &packed)?;
|
||||
let d = self.config.encoder.hidden_size;
|
||||
// Marker positions, and each row's first token, in packed order.
|
||||
let global: Vec<u32> = flat
|
||||
.marker_idx
|
||||
.chunks(flat.k_max)
|
||||
.zip(&packed.offsets)
|
||||
.flat_map(|(ms, &off)| ms.iter().map(move |&m| off as u32 + m))
|
||||
.collect();
|
||||
let global = Tensor::from_vec(global, b * flat.k_max, &self.device)?;
|
||||
let marker_h = h.index_select(&global, 0)?.reshape((b, flat.k_max, d))?;
|
||||
let firsts: Vec<u32> = packed.offsets.iter().map(|&o| o as u32).collect();
|
||||
let pooled = h.index_select(&Tensor::from_vec(firsts, b, &self.device)?, 0)?;
|
||||
self.finish(batch, &marker_h, &pooled)
|
||||
}
|
||||
|
||||
/// Score the gathered marker states `(b, k_max, d)` and run the act
|
||||
/// head over the pooled first positions `(b, d)`.
|
||||
fn finish(
|
||||
&self,
|
||||
batch: &DecisionBatch,
|
||||
marker_h: &Tensor,
|
||||
pooled: &Tensor,
|
||||
) -> Result<Vec<DecisionRow>> {
|
||||
let b = batch.seqs.len();
|
||||
let logits: Vec<Vec<f32>> = self.head.score(marker_h)?.to_vec2()?;
|
||||
let mut feats = Vec::with_capacity(b * 4);
|
||||
let mut row_logits = Vec::with_capacity(b);
|
||||
for (s, full) in batch.seqs.iter().zip(logits) {
|
||||
let row: Vec<f32> = full[..s.markers.len()].to_vec();
|
||||
feats.extend_from_slice(&act_features(&row));
|
||||
row_logits.push(row);
|
||||
}
|
||||
let feats = Tensor::from_vec(feats, (b, 4), &self.device)?;
|
||||
let act_logits = self.head.act(pooled, &feats)?;
|
||||
let act = candle_nn::ops::softmax_last_dim(&act_logits)?.to_vec2::<f32>()?;
|
||||
Ok(row_logits
|
||||
.into_iter()
|
||||
.zip(act)
|
||||
.map(|(logits, act)| DecisionRow { logits, act })
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
/// Runtime kill switch for the flash path: `NEURON_FLASH_ATTN=0` (or
|
||||
/// `false`) forces eager attention without a rebuild — the same switch
|
||||
/// the text architectures read. Read once.
|
||||
#[cfg(feature = "flash-attn")]
|
||||
fn flash_enabled() -> bool {
|
||||
static ENABLED: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
|
||||
*ENABLED
|
||||
.get_or_init(|| !std::env::var("NEURON_FLASH_ATTN").is_ok_and(|v| v == "0" || v == "false"))
|
||||
}
|
||||
|
||||
/// The act head's four inputs, from one row's uncalibrated logits:
|
||||
/// top-1 probability, top-1 minus top-2, entropy normalised by
|
||||
/// `ln(max(k, 2))`, and `k / 255`.
|
||||
///
|
||||
/// The reference computes these over the batch's padded width, with
|
||||
/// absent option slots filled with a logit of -1e4. Those slots carry
|
||||
/// exactly zero probability in f32, so computing per row over only the
|
||||
/// real options gives the same result.
|
||||
pub fn act_features(logits: &[f32]) -> [f32; 4] {
|
||||
let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
let exps: Vec<f32> = logits.iter().map(|&z| (z - max).exp()).collect();
|
||||
let sum: f32 = exps.iter().sum();
|
||||
let mut p: Vec<f32> = exps.iter().map(|e| e / sum).collect();
|
||||
let k = (logits.len().max(2)) as f32;
|
||||
let ent = -p.iter().map(|&x| x * x.max(1e-9).ln()).sum::<f32>() / k.ln();
|
||||
p.sort_by(|a, b| b.total_cmp(a));
|
||||
let top1 = p[0];
|
||||
let top2 = p.get(1).copied().unwrap_or(0.0);
|
||||
[top1, top1 - top2, ent, k / 255.0]
|
||||
}
|
||||
@@ -0,0 +1,450 @@
|
||||
//! Parity against the upstream reference.
|
||||
//!
|
||||
//! `testdata/laya/tiny/` is a toy-sized checkpoint built and scored by
|
||||
//! the SDK's own `DecisionModel` (`script/generate_laya_fixture.py`);
|
||||
//! it runs in CI. `testdata/laya/reference.json` holds the released
|
||||
//! checkpoints' outputs over a real corpus (`script/laya-reference.py`);
|
||||
//! scoring it needs the weights, so that test is ignored unless
|
||||
//! `LAYA_REFERENCE_DIR` names a snapshot of `convaiinnovations/laya`.
|
||||
|
||||
use super::*;
|
||||
use candle_core::{DType, Device};
|
||||
use serde::Deserialize;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
fn testdata() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TinyExpected {
|
||||
seqs: Vec<TinySeq>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TinySeq {
|
||||
input_ids: Vec<u32>,
|
||||
markers: Vec<usize>,
|
||||
qtype: u8,
|
||||
logits: Vec<f32>,
|
||||
act: Vec<f32>,
|
||||
}
|
||||
|
||||
fn load_tiny() -> (LayaModel, TinyExpected) {
|
||||
let dir = testdata().join("tiny");
|
||||
let model = LayaModel::load(&dir, &Device::Cpu, DType::F32).expect("load tiny checkpoint");
|
||||
let expected: TinyExpected = serde_json::from_str(
|
||||
&std::fs::read_to_string(dir.join("expected.json")).expect("read expected.json"),
|
||||
)
|
||||
.expect("parse expected.json");
|
||||
(model, expected)
|
||||
}
|
||||
|
||||
fn batch_of(seqs: &[TinySeq]) -> DecisionBatch {
|
||||
DecisionBatch {
|
||||
seqs: seqs
|
||||
.iter()
|
||||
.map(|s| DecisionSeq {
|
||||
input_ids: s.input_ids.clone(),
|
||||
markers: s.markers.clone(),
|
||||
qtype: s.qtype,
|
||||
})
|
||||
.collect(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Device and precision for the weight-backed tests: `LAYA_DEVICE`
|
||||
/// (`cpu`, or `cuda:N` in a build with the `cuda` feature) and
|
||||
/// `LAYA_DTYPE` (`f32`, `bf16` or `f16`). Defaults: CPU, f32.
|
||||
fn test_device() -> (Device, DType) {
|
||||
let device = match std::env::var("LAYA_DEVICE").as_deref() {
|
||||
Ok("cpu") | Err(_) => Device::Cpu,
|
||||
Ok(s) => {
|
||||
let ordinal = s
|
||||
.strip_prefix("cuda:")
|
||||
.and_then(|n| n.parse().ok())
|
||||
.unwrap_or_else(|| panic!("LAYA_DEVICE={s}: expected cpu or cuda:N"));
|
||||
Device::new_cuda(ordinal).unwrap_or_else(|e| panic!("LAYA_DEVICE={s}: {e}"))
|
||||
}
|
||||
};
|
||||
let dtype = match std::env::var("LAYA_DTYPE").as_deref() {
|
||||
Ok("f32") | Err(_) => DType::F32,
|
||||
Ok("bf16") => DType::BF16,
|
||||
Ok("f16") => DType::F16,
|
||||
Ok(s) => panic!("LAYA_DTYPE={s}: expected f32, bf16 or f16"),
|
||||
};
|
||||
(device, dtype)
|
||||
}
|
||||
|
||||
/// Logit tolerance against the f32 reference for a given precision.
|
||||
/// Reduced precision is judged mainly on argmax agreement; the bound
|
||||
/// only catches a broken kernel, not rounding.
|
||||
fn logit_tolerance(dtype: DType) -> f32 {
|
||||
if dtype == DType::F32 { 1e-3 } else { 1.0 }
|
||||
}
|
||||
|
||||
fn max_abs_diff(a: &[f32], b: &[f32]) -> f32 {
|
||||
assert_eq!(a.len(), b.len(), "length mismatch");
|
||||
a.iter()
|
||||
.zip(b)
|
||||
.map(|(x, y)| (x - y).abs())
|
||||
.fold(0.0, f32::max)
|
||||
}
|
||||
|
||||
/// The whole forward — encoder with both attention kinds and rotary
|
||||
/// bases, padding, type embedding, head layers, marker gather, scorer
|
||||
/// and act head — against the reference, in f32.
|
||||
#[test]
|
||||
fn tiny_checkpoint_matches_reference() {
|
||||
let (model, expected) = load_tiny();
|
||||
let rows = model.forward(&batch_of(&expected.seqs)).expect("forward");
|
||||
assert_eq!(rows.len(), expected.seqs.len());
|
||||
for (i, (got, want)) in rows.iter().zip(&expected.seqs).enumerate() {
|
||||
let dl = max_abs_diff(&got.logits, &want.logits);
|
||||
let da = max_abs_diff(&got.act, &want.act);
|
||||
assert!(
|
||||
dl < 1e-5 && da < 1e-5,
|
||||
"row {i}: logit diff {dl:e}, act diff {da:e}\n got {:?} / {:?}\nwant {:?} / {:?}",
|
||||
got.logits,
|
||||
got.act,
|
||||
want.logits,
|
||||
want.act
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// A row's result must not depend on what it was batched with: padded
|
||||
/// keys are masked out of every attention, including the head's.
|
||||
#[test]
|
||||
fn padding_does_not_change_a_row() {
|
||||
let (model, expected) = load_tiny();
|
||||
let batched = model.forward(&batch_of(&expected.seqs)).expect("batched");
|
||||
for (i, seq) in expected.seqs.iter().enumerate() {
|
||||
let alone = model
|
||||
.forward(&batch_of(std::slice::from_ref(seq)))
|
||||
.expect("single");
|
||||
let dl = max_abs_diff(&alone[0].logits, &batched[i].logits);
|
||||
let da = max_abs_diff(&alone[0].act, &batched[i].act);
|
||||
assert!(dl < 1e-5 && da < 1e-5, "row {i}: {dl:e} / {da:e}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn malformed_batches_are_refused() {
|
||||
let (model, _) = load_tiny();
|
||||
let seq = |ids: Vec<u32>, markers: Vec<usize>, qtype| DecisionBatch {
|
||||
seqs: vec![DecisionSeq {
|
||||
input_ids: ids,
|
||||
markers,
|
||||
qtype,
|
||||
}],
|
||||
};
|
||||
assert!(model.forward(&DecisionBatch::default()).is_err());
|
||||
assert!(model.forward(&seq(vec![1, 2, 3], vec![], 0)).is_err());
|
||||
assert!(model.forward(&seq(vec![1, 2, 3], vec![3], 0)).is_err());
|
||||
assert!(model.forward(&seq(vec![1, 2, 3], vec![1], 3)).is_err());
|
||||
assert!(model.forward(&seq(vec![1, 500, 3], vec![1], 0)).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn act_features_match_the_reference_definition() {
|
||||
// Two equal logits: top1 = top2 = 0.5, entropy exactly 1.
|
||||
let f = model::act_features(&[0.3, 0.3]);
|
||||
assert!((f[0] - 0.5).abs() < 1e-6 && f[1].abs() < 1e-6 && (f[2] - 1.0).abs() < 1e-6);
|
||||
assert!((f[3] - 2.0 / 255.0).abs() < 1e-7);
|
||||
// One option: probability 1, no second, k clamps to 2.
|
||||
let f = model::act_features(&[4.0]);
|
||||
assert_eq!(f[0], 1.0);
|
||||
assert_eq!(f[1], 1.0);
|
||||
assert!(f[2].abs() < 1e-6);
|
||||
assert!((f[3] - 2.0 / 255.0).abs() < 1e-7);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------- released checkpoints
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Reference {
|
||||
cases: Vec<RefCase>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefCase {
|
||||
name: String,
|
||||
items: Vec<RefItem>,
|
||||
response: serde_json::Value,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct RefItem {
|
||||
input_ids: Vec<u32>,
|
||||
markers: Vec<usize>,
|
||||
qtype: u8,
|
||||
logits: Vec<f32>,
|
||||
act: Vec<f32>,
|
||||
}
|
||||
|
||||
/// Every case in `reference.json`, scored by the checkpoint the
|
||||
/// reference routed it to, on the device and precision
|
||||
/// [`test_device`] selects (f32 on the CPU by default).
|
||||
///
|
||||
/// `LAYA_REFERENCE_DIR` is a snapshot of `convaiinnovations/laya` at the
|
||||
/// revision recorded in the fixture header: the root holds the English
|
||||
/// checkpoint, `multilingual/` and `typed-decisions/` the other two.
|
||||
#[test]
|
||||
#[ignore = "needs the released weights: set LAYA_REFERENCE_DIR"]
|
||||
fn released_checkpoints_match_reference() {
|
||||
let Ok(root) = std::env::var("LAYA_REFERENCE_DIR") else {
|
||||
panic!("LAYA_REFERENCE_DIR is not set");
|
||||
};
|
||||
let root = PathBuf::from(root);
|
||||
let (device, dtype) = test_device();
|
||||
let tolerance = logit_tolerance(dtype);
|
||||
let reference: Reference = serde_json::from_str(
|
||||
&std::fs::read_to_string(testdata().join("reference.json")).expect("read reference.json"),
|
||||
)
|
||||
.expect("parse reference.json");
|
||||
|
||||
// The reference's own GPU run (bf16 autocast), for judging reduced
|
||||
// precision against the drift upstream itself serves with.
|
||||
let torch_bf16: Option<Reference> =
|
||||
std::fs::read_to_string(testdata().join("reference_cuda_bf16.json"))
|
||||
.ok()
|
||||
.map(|t| serde_json::from_str(&t).expect("parse reference_cuda_bf16.json"));
|
||||
// Per checkpoint: summed centred drift of ours and of the reference's
|
||||
// bf16 run, their maxima, and the row count.
|
||||
let mut drift: std::collections::BTreeMap<String, (f32, f32, f32, f32, usize)> =
|
||||
Default::default();
|
||||
|
||||
let mut models = std::collections::BTreeMap::new();
|
||||
let mut stats: std::collections::BTreeMap<String, (usize, f32, f32, usize)> =
|
||||
Default::default();
|
||||
// Forward wall time and tokens encoded, per checkpoint: a first
|
||||
// CPU number, not a benchmark (debug builds are far slower).
|
||||
let mut timing: std::collections::BTreeMap<String, (f64, usize)> = Default::default();
|
||||
for case in &reference.cases {
|
||||
let ckpt = case.response["routing"]["model"]
|
||||
.as_str()
|
||||
.expect("routing.model")
|
||||
.to_string();
|
||||
let model = models.entry(ckpt.clone()).or_insert_with(|| {
|
||||
let dir = match ckpt.as_str() {
|
||||
"english" => root.clone(),
|
||||
other => root.join(other),
|
||||
};
|
||||
LayaModel::load(&dir, &device, dtype).unwrap_or_else(|e| panic!("load {ckpt}: {e:#}"))
|
||||
});
|
||||
let batch = DecisionBatch {
|
||||
seqs: case
|
||||
.items
|
||||
.iter()
|
||||
.map(|it| DecisionSeq {
|
||||
input_ids: it.input_ids.clone(),
|
||||
markers: it.markers.clone(),
|
||||
qtype: it.qtype,
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
let started = std::time::Instant::now();
|
||||
let rows = model
|
||||
.forward(&batch)
|
||||
.unwrap_or_else(|e| panic!("{}: {e:#}", case.name));
|
||||
let t = timing.entry(ckpt.clone()).or_default();
|
||||
t.0 += started.elapsed().as_secs_f64() * 1e3;
|
||||
t.1 += batch.seqs.iter().map(|s| s.input_ids.len()).sum::<usize>();
|
||||
let torch_case = torch_bf16
|
||||
.as_ref()
|
||||
.and_then(|r| r.cases.iter().find(|c| c.name == case.name));
|
||||
for (i, (got, want)) in rows.iter().zip(&case.items).enumerate() {
|
||||
let d = drift.entry(ckpt.clone()).or_default();
|
||||
let ours = centred_diff(&got.logits, &want.logits);
|
||||
d.0 += ours;
|
||||
d.2 = d.2.max(ours);
|
||||
if let Some(t) = torch_case.and_then(|c| c.items.get(i)) {
|
||||
let theirs = centred_diff(&t.logits, &want.logits);
|
||||
d.1 += theirs;
|
||||
d.3 = d.3.max(theirs);
|
||||
}
|
||||
d.4 += 1;
|
||||
}
|
||||
let entry = stats.entry(ckpt.clone()).or_default();
|
||||
for (got, want) in rows.iter().zip(&case.items) {
|
||||
let dl = max_abs_diff(&got.logits, &want.logits);
|
||||
let da = max_abs_diff(&got.act, &want.act);
|
||||
let argmax = |v: &[f32]| {
|
||||
v.iter()
|
||||
.enumerate()
|
||||
.max_by(|a, b| a.1.total_cmp(b.1))
|
||||
.map(|(i, _)| i)
|
||||
};
|
||||
entry.0 += 1;
|
||||
entry.1 = entry.1.max(dl);
|
||||
entry.2 = entry.2.max(da);
|
||||
if argmax(&got.logits) == argmax(&want.logits) {
|
||||
entry.3 += 1;
|
||||
}
|
||||
if dtype != DType::F32 && std::env::var_os("LAYA_VERBOSE").is_some() {
|
||||
eprintln!(" {:<40} {ckpt:<16} |Δlogit| {dl:.3}", case.name);
|
||||
}
|
||||
assert!(
|
||||
dl < tolerance,
|
||||
"{} ({ckpt}): logit diff {dl:e}\n got {:?}\nwant {:?}",
|
||||
case.name,
|
||||
got.logits,
|
||||
want.logits
|
||||
);
|
||||
assert!(da < tolerance, "{} ({ckpt}): act diff {da:e}", case.name);
|
||||
}
|
||||
}
|
||||
for (ckpt, (n, dl, da, agree)) in &stats {
|
||||
let (ms, tokens) = timing[ckpt];
|
||||
eprintln!(
|
||||
"[{device:?} {dtype:?}] {ckpt}: {n} questions, max |Δlogit| {dl:.2e}, max |Δact| {da:.2e}, argmax {agree}/{n}; \
|
||||
{tokens} tokens in {ms:.0} ms ({:.2} ms/token)",
|
||||
ms / tokens as f64
|
||||
);
|
||||
assert_eq!(agree, n, "{ckpt}: argmax disagreement");
|
||||
}
|
||||
// Reduced precision: centred drift (softmax ignores a constant shift)
|
||||
// against f32, next to the reference's own bf16 drift. Per-case drift
|
||||
// is chaotic, so the gate is the mean: no worse than twice upstream's,
|
||||
// plus 0.01 — about one bf16 ulp at these logits' magnitude — so a
|
||||
// checkpoint with a single fixture case whose reference drift happens
|
||||
// to round to nothing is not held to zero.
|
||||
for (ckpt, (ours, theirs, ours_max, theirs_max, n)) in &drift {
|
||||
let (ours, theirs) = (ours / *n as f32, theirs / *n as f32);
|
||||
eprintln!(
|
||||
"[{device:?} {dtype:?}] {ckpt}: centred drift mean {ours:.4} (reference bf16 {theirs:.4}), \
|
||||
max {ours_max:.3} (reference bf16 {theirs_max:.3})"
|
||||
);
|
||||
if dtype != DType::F32 && torch_bf16.is_some() {
|
||||
assert!(
|
||||
ours <= 2.0 * theirs + 0.01,
|
||||
"{ckpt}: mean drift {ours:.4} exceeds twice the reference's own bf16 drift {theirs:.4}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Largest difference between two logit rows after removing each row's
|
||||
/// mean — the part of a difference a softmax can see.
|
||||
fn centred_diff(a: &[f32], b: &[f32]) -> f32 {
|
||||
let mean = |v: &[f32]| v.iter().sum::<f32>() / v.len() as f32;
|
||||
let (ma, mb) = (mean(a), mean(b));
|
||||
a.iter()
|
||||
.zip(b)
|
||||
.map(|(x, y)| ((x - ma) - (y - mb)).abs())
|
||||
.fold(0.0, f32::max)
|
||||
}
|
||||
|
||||
/// Forward latency for the four workloads #337 benchmarks against the
|
||||
/// PyTorch reference, plus batch throughput, on the English checkpoint.
|
||||
///
|
||||
/// Sequences are real ones from `reference.json`: `choice_basic` is the
|
||||
/// short state (41 tokens), `multi_question` five questions over it, and
|
||||
/// `state_right_truncated` a full 512-token sequence (repeated for the
|
||||
/// five-question and throughput batches). The timed span is the whole
|
||||
/// `forward` call — host-to-device copy, forward, and the CPU readback
|
||||
/// that ends it — so it is directly comparable to the reference's
|
||||
/// end-to-end `predict()` minus tokenisation.
|
||||
///
|
||||
/// Run with `cargo test --release -p neuron [--features cuda] --
|
||||
/// --ignored --nocapture laya_forward_timing`, with `LAYA_REFERENCE_DIR`
|
||||
/// and optionally `LAYA_DEVICE` / `LAYA_DTYPE` set.
|
||||
#[test]
|
||||
#[ignore = "benchmark: needs the released weights (LAYA_REFERENCE_DIR)"]
|
||||
fn laya_forward_timing() {
|
||||
let root = PathBuf::from(std::env::var("LAYA_REFERENCE_DIR").expect("LAYA_REFERENCE_DIR"));
|
||||
let (device, dtype) = test_device();
|
||||
let reference: Reference = serde_json::from_str(
|
||||
&std::fs::read_to_string(testdata().join("reference.json")).expect("read reference.json"),
|
||||
)
|
||||
.expect("parse reference.json");
|
||||
let seqs_of = |name: &str| -> Vec<DecisionSeq> {
|
||||
reference
|
||||
.cases
|
||||
.iter()
|
||||
.find(|c| c.name == name)
|
||||
.unwrap_or_else(|| panic!("no case {name}"))
|
||||
.items
|
||||
.iter()
|
||||
.map(|it| DecisionSeq {
|
||||
input_ids: it.input_ids.clone(),
|
||||
markers: it.markers.clone(),
|
||||
qtype: it.qtype,
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
let long = seqs_of("state_right_truncated").remove(0);
|
||||
let workloads: Vec<(String, DecisionBatch)> = vec![
|
||||
(
|
||||
"short x1q".into(),
|
||||
DecisionBatch {
|
||||
seqs: seqs_of("choice_basic"),
|
||||
},
|
||||
),
|
||||
(
|
||||
"short x5q".into(),
|
||||
DecisionBatch {
|
||||
seqs: seqs_of("multi_question"),
|
||||
},
|
||||
),
|
||||
(
|
||||
"512tok x1q".into(),
|
||||
DecisionBatch {
|
||||
seqs: vec![long.clone()],
|
||||
},
|
||||
),
|
||||
(
|
||||
"512tok x5q".into(),
|
||||
DecisionBatch {
|
||||
seqs: vec![long.clone(); 5],
|
||||
},
|
||||
),
|
||||
(
|
||||
"512tok x16".into(),
|
||||
DecisionBatch {
|
||||
seqs: vec![long.clone(); 16],
|
||||
},
|
||||
),
|
||||
(
|
||||
"512tok x64".into(),
|
||||
DecisionBatch {
|
||||
seqs: vec![long.clone(); 64],
|
||||
},
|
||||
),
|
||||
];
|
||||
|
||||
let model = LayaModel::load(&root, &device, dtype).expect("load English checkpoint");
|
||||
// `LAYA_TIMING_ONLY` keeps the workloads whose name contains it, and
|
||||
// `LAYA_TIMING_ITERS` sets the timed iterations — for profiling one
|
||||
// workload under a tracer.
|
||||
let only = std::env::var("LAYA_TIMING_ONLY").ok();
|
||||
let iters: usize = std::env::var("LAYA_TIMING_ITERS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(20);
|
||||
for (name, batch) in &workloads {
|
||||
if only.as_deref().is_some_and(|o| !name.contains(o)) {
|
||||
continue;
|
||||
}
|
||||
for _ in 0..3 {
|
||||
model.forward(batch).expect("warm-up forward");
|
||||
}
|
||||
let mut ms: Vec<f64> = (0..iters)
|
||||
.map(|_| {
|
||||
let started = std::time::Instant::now();
|
||||
model.forward(batch).expect("forward");
|
||||
started.elapsed().as_secs_f64() * 1e3
|
||||
})
|
||||
.collect();
|
||||
ms.sort_by(f64::total_cmp);
|
||||
let median = ms[ms.len() / 2];
|
||||
let tokens: usize = batch.seqs.iter().map(|s| s.input_ids.len()).sum();
|
||||
eprintln!(
|
||||
"[{device:?} {dtype:?}] {name:>11}: median {median:7.1} ms, min {:7.1} ms, \
|
||||
{tokens} tokens, {:.0} seq/s",
|
||||
ms[0],
|
||||
batch.seqs.len() as f64 / (median / 1e3)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -20,6 +20,7 @@
|
||||
//! TP-aware analogues live in `harness/tp/tp_<family>.rs` and follow
|
||||
//! the pattern set by `tp_qwen3.rs`.
|
||||
|
||||
pub mod laya;
|
||||
pub mod phase_probe;
|
||||
pub mod qwen3_5;
|
||||
pub mod qwen4_exp;
|
||||
|
||||
@@ -143,6 +143,9 @@ pub enum LoadedHandle {
|
||||
Tp(Arc<TpLoadedModel>),
|
||||
/// Z-Image text-to-image pipeline (#198).
|
||||
Image(Arc<LoadedImageModel>),
|
||||
/// Laya decision model (#336): typed questions in, calibrated
|
||||
/// option scores out, no text generation.
|
||||
Decision(Arc<super::decision_model::LoadedDecisionModel>),
|
||||
}
|
||||
|
||||
impl LoadedHandle {
|
||||
@@ -152,6 +155,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.model_id,
|
||||
LoadedHandle::Image(m) => &m.model_id,
|
||||
LoadedHandle::Decision(m) => &m.model_id,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -162,6 +166,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.spec,
|
||||
LoadedHandle::Image(m) => &m.spec,
|
||||
LoadedHandle::Decision(m) => &m.spec,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -171,6 +176,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.devices.clone(),
|
||||
LoadedHandle::Image(m) => m.devices.clone(),
|
||||
LoadedHandle::Decision(m) => m.devices.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -184,6 +190,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.poisoned.load(Ordering::Acquire),
|
||||
LoadedHandle::Image(m) => m.poisoned.load(Ordering::Acquire),
|
||||
LoadedHandle::Decision(m) => m.poisoned.load(Ordering::Acquire),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -195,6 +202,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => (m.admission.in_flight(), m.admission.queue_depth()),
|
||||
LoadedHandle::Image(m) => (m.admission.in_flight(), m.admission.queue_depth()),
|
||||
LoadedHandle::Decision(m) => (m.admission.in_flight(), m.admission.queue_depth()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -207,6 +215,9 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
|
||||
LoadedHandle::Image(m) => (m.admission.max_in_flight(), m.admission.max_queue_depth()),
|
||||
LoadedHandle::Decision(m) => {
|
||||
(m.admission.max_in_flight(), m.admission.max_queue_depth())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -230,6 +241,10 @@ impl LoadedHandle {
|
||||
m.admission.anon_in_flight(),
|
||||
m.admission.anon_max_in_flight(),
|
||||
),
|
||||
LoadedHandle::Decision(m) => (
|
||||
m.admission.anon_in_flight(),
|
||||
m.admission.anon_max_in_flight(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -240,6 +255,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.admission.close(),
|
||||
LoadedHandle::Image(m) => m.admission.close(),
|
||||
LoadedHandle::Decision(m) => m.admission.close(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -251,6 +267,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.admission.rejections(),
|
||||
LoadedHandle::Image(m) => m.admission.rejections(),
|
||||
LoadedHandle::Decision(m) => m.admission.rejections(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -262,6 +279,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.admission,
|
||||
LoadedHandle::Image(m) => &m.admission,
|
||||
LoadedHandle::Decision(m) => &m.admission,
|
||||
};
|
||||
(admission.kv_budget_mb(), admission.kv_available_mb())
|
||||
}
|
||||
@@ -275,7 +293,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Tp(m) => &m.prefill_rate,
|
||||
// Image models have no token throughput; steps/sec lives in
|
||||
// the per-request timing instead.
|
||||
LoadedHandle::Image(_) => return (0.0, 0.0),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return (0.0, 0.0),
|
||||
};
|
||||
(ema.get().unwrap_or(0.0), ema.decode().unwrap_or(0.0))
|
||||
}
|
||||
@@ -287,12 +305,15 @@ impl LoadedHandle {
|
||||
if matches!(self, LoadedHandle::Image(_)) {
|
||||
return vec!["image".to_string()];
|
||||
}
|
||||
if matches!(self, LoadedHandle::Decision(_)) {
|
||||
return vec!["decision".to_string()];
|
||||
}
|
||||
let mut caps = vec!["text".to_string()];
|
||||
let has_vision = match self {
|
||||
LoadedHandle::Single(m) => m.has_vision,
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.has_vision,
|
||||
LoadedHandle::Image(_) => unreachable!("handled above"),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => unreachable!("handled above"),
|
||||
};
|
||||
if has_vision {
|
||||
caps.push("vision".to_string());
|
||||
@@ -338,7 +359,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Tp(m) => (m.last_free_mb.load(Ordering::Acquire), m.context_profile),
|
||||
// Image generation has its own VRAM path and no prefill
|
||||
// floor, so this check would say nothing useful about it.
|
||||
LoadedHandle::Image(_) => return None,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
|
||||
};
|
||||
// `0` is the CPU-load and missing-worker sentinel — the same one
|
||||
// `check_vram` skips on. Nothing to assert either way.
|
||||
@@ -419,7 +440,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.tool_call_tokens.is_some(),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.tool_call_tokens.is_some(),
|
||||
LoadedHandle::Image(_) => false,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -430,7 +451,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.reasoning_tokens.is_some(),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.reasoning_tokens.is_some(),
|
||||
LoadedHandle::Image(_) => false,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => false,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -441,7 +462,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.preserve_thinking,
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.preserve_thinking,
|
||||
LoadedHandle::Image(_) => None,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -451,7 +472,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.reasoning_efforts.clone(),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.reasoning_efforts.clone(),
|
||||
LoadedHandle::Image(_) => Default::default(),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => Default::default(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -477,7 +498,7 @@ impl LoadedHandle {
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.query_vram_tightest_free_mb().await,
|
||||
// No context-limit derivation for image models; nothing to cache.
|
||||
LoadedHandle::Image(_) => return,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return,
|
||||
};
|
||||
// Don't clobber a good cached value with a transient `0`
|
||||
// (worker gone/poisoned sentinel).
|
||||
@@ -486,7 +507,7 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.last_free_mb.store(free, Ordering::Release),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.last_free_mb.store(free, Ordering::Release),
|
||||
LoadedHandle::Image(_) => {}
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -515,7 +536,7 @@ impl LoadedHandle {
|
||||
m.prefill_rate.get(),
|
||||
),
|
||||
// Token limits are meaningless for image generation.
|
||||
LoadedHandle::Image(_) => return None,
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => return None,
|
||||
};
|
||||
let rate = rate.unwrap_or(cfg.bootstrap_prefill_tok_per_sec);
|
||||
// The KV budget admission will actually grant. Advertising past
|
||||
@@ -525,7 +546,9 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => m.admission.kv_budget_mb(),
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => m.admission.kv_budget_mb(),
|
||||
LoadedHandle::Image(_) => unreachable!("derived_limit returned above"),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
|
||||
unreachable!("derived_limit returned above")
|
||||
}
|
||||
};
|
||||
let limit = super::context_limit::derive_limit(
|
||||
&profile,
|
||||
@@ -544,7 +567,9 @@ impl LoadedHandle {
|
||||
LoadedHandle::Single(m) => &m.derived_input_cap,
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(m) => &m.derived_input_cap,
|
||||
LoadedHandle::Image(_) => unreachable!("derived_limit returned above"),
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
|
||||
unreachable!("derived_limit returned above")
|
||||
}
|
||||
};
|
||||
cap.store(input, Ordering::Release);
|
||||
}
|
||||
@@ -3282,7 +3307,7 @@ impl CandleHarness {
|
||||
LoadedHandle::Tp(m) => {
|
||||
return self.inference_tp_stream(m, request, principal, wire).await;
|
||||
}
|
||||
LoadedHandle::Image(_) => {
|
||||
LoadedHandle::Image(_) | LoadedHandle::Decision(_) => {
|
||||
return Err(InferenceError::WrongModality {
|
||||
model_id: request.model.clone(),
|
||||
});
|
||||
@@ -4001,6 +4026,11 @@ impl Harness for CandleHarness {
|
||||
if matches!(plan.format, super::preflight::SourceFormat::Diffusers) {
|
||||
return self.load_image_model(spec, &source_id).await;
|
||||
}
|
||||
// Laya decision models (#336) take the decision path: one
|
||||
// device, unquantized (preflight rejected tp>1 and quant).
|
||||
if matches!(plan.format, super::preflight::SourceFormat::DecisionModel) {
|
||||
return self.load_decision_model(spec, &source_id).await;
|
||||
}
|
||||
|
||||
let tp_size = spec.tensor_parallel.unwrap_or(1);
|
||||
if tp_size > 1 {
|
||||
@@ -4367,6 +4397,11 @@ impl Harness for CandleHarness {
|
||||
);
|
||||
}
|
||||
}
|
||||
LoadedHandle::Decision(dm) => {
|
||||
// Each checkpoint's weights drop on the thread that owns
|
||||
// them (the device worker, for a CUDA load).
|
||||
dm.release().await;
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
LoadedHandle::Tp(tp) => {
|
||||
// Try to recover the inner TpLoadedModel so we can move
|
||||
@@ -4541,6 +4576,132 @@ impl CandleHarness {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Load a Laya decision model (#336): resolve the checkpoints' files,
|
||||
/// build them on the device worker (or the CPU), and register the
|
||||
/// model with `capabilities: ["decision"]`.
|
||||
///
|
||||
/// The repo root is required. A released family also carries
|
||||
/// `multilingual/` and `typed-decisions/` beside it (#339); each one
|
||||
/// the repo has is loaded too, so the `/v1/systemone` handler can
|
||||
/// route non-English states to the multilingual checkpoint and honour
|
||||
/// a request that pins one. A repo without them (a single fine-tune)
|
||||
/// serves its root alone.
|
||||
async fn load_decision_model(
|
||||
&self,
|
||||
spec: &ModelSpec,
|
||||
source_id: &cortex_core::source::ModelSourceId,
|
||||
) -> Result<()> {
|
||||
let devices = spec.devices.clone().unwrap_or_else(|| vec![0]);
|
||||
let device = Self::pick_device(&devices)?;
|
||||
#[cfg(feature = "cuda")]
|
||||
let worker = if matches!(device, Device::Cuda(_)) {
|
||||
Some(self.ensure_device_worker(devices[0]).await?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
// CPU build: always the CPU backend. Correct, but at hundreds of
|
||||
// milliseconds per question it is a fallback, not a serving tier.
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
let worker = {
|
||||
let _ = &device;
|
||||
None
|
||||
};
|
||||
|
||||
let files = self.resolve_decision_files(source_id, None).await?;
|
||||
let root = super::decision_model::DecisionCheckpoint::load(
|
||||
super::decision_model::ROOT_CHECKPOINT.to_string(),
|
||||
files,
|
||||
&spec.model_id,
|
||||
worker.clone(),
|
||||
)
|
||||
.await?;
|
||||
let dtype = root.dtype;
|
||||
let mut checkpoints = vec![root];
|
||||
for sub in super::decision::route::FAMILY
|
||||
.iter()
|
||||
.filter(|c| **c != super::decision_model::ROOT_CHECKPOINT)
|
||||
{
|
||||
let files = match self.resolve_decision_files(source_id, Some(sub)).await {
|
||||
Ok(files) => files,
|
||||
Err(e) => {
|
||||
tracing::info!(
|
||||
model = %spec.model_id,
|
||||
checkpoint = sub,
|
||||
error = %format!("{e:#}"),
|
||||
"decision checkpoint not available; serving without it"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
checkpoints.push(
|
||||
super::decision_model::DecisionCheckpoint::load(
|
||||
sub.to_string(),
|
||||
files,
|
||||
&spec.model_id,
|
||||
worker.clone(),
|
||||
)
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
let names: Vec<&str> = checkpoints.iter().map(|c| c.name.as_str()).collect();
|
||||
tracing::info!(model = %spec.model_id, checkpoints = ?names, "decision checkpoints loaded");
|
||||
|
||||
let loaded = super::decision_model::LoadedDecisionModel {
|
||||
model_id: spec.model_id.clone(),
|
||||
spec: spec.clone(),
|
||||
devices,
|
||||
checkpoints,
|
||||
poisoned: Arc::new(AtomicBool::new(false)),
|
||||
admission: crate::harness::admission::AdmissionController::new(&self.admission_cfg),
|
||||
};
|
||||
let mut models = self.models.write().await;
|
||||
models.insert(
|
||||
spec.model_id.clone(),
|
||||
LoadedHandle::Decision(Arc::new(loaded)),
|
||||
);
|
||||
tracing::info!(model = %spec.model_id, dtype = ?dtype, "decision model loaded");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fetch one Laya checkpoint's files through hf-hub: the repo root
|
||||
/// for `subfolder = None`, else that subfolder. Only the files the
|
||||
/// checkpoint needs are fetched, not the whole repo.
|
||||
async fn resolve_decision_files(
|
||||
&self,
|
||||
source_id: &cortex_core::source::ModelSourceId,
|
||||
subfolder: Option<&str>,
|
||||
) -> Result<super::decision_model::CheckpointFiles> {
|
||||
let api = self.hf_api_for(&source_id.scheme)?;
|
||||
let repo = api.model(source_id.repo_path());
|
||||
let display_id = source_id.to_string();
|
||||
let prefix = subfolder.map(|s| format!("{s}/")).unwrap_or_default();
|
||||
let mut fetched = Vec::new();
|
||||
for file in [
|
||||
"rl_agent_config.json",
|
||||
"encoder/config.json",
|
||||
"model.safetensors",
|
||||
"tokenizer/tokenizer.json",
|
||||
"tokenizer/tokenizer_config.json",
|
||||
] {
|
||||
let path = format!("{prefix}{file}");
|
||||
fetched.push(
|
||||
repo.get(&path)
|
||||
.await
|
||||
.with_context(|| format!("fetch {path} from {display_id}"))?,
|
||||
);
|
||||
}
|
||||
// hf-hub materialises a snapshot as a directory tree mirroring
|
||||
// the repo, so the checkpoint directory is where its config sits.
|
||||
let dir = fetched[0]
|
||||
.parent()
|
||||
.context("checkpoint config has no parent directory")?
|
||||
.to_path_buf();
|
||||
Ok(super::decision_model::CheckpointFiles {
|
||||
dir,
|
||||
tokenizer: fetched[3].clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Resolve a diffusers-layout Z-Image repo into local component
|
||||
/// paths, downloading through hf-hub as needed. Verifies the
|
||||
/// pipeline class before fetching multi-GB weight shards.
|
||||
@@ -4654,6 +4815,189 @@ impl CandleHarness {
|
||||
}
|
||||
}
|
||||
|
||||
/// The loaded decision model `model_id` (#336), for the
|
||||
/// `/v1/systemone` handler: it takes the admission permit, scores
|
||||
/// through [`super::decision_model::DecisionCheckpoint::decide`],
|
||||
/// and shapes the answers. A model of another kind is a
|
||||
/// `WrongModality`, as for images.
|
||||
pub async fn decision_model(
|
||||
&self,
|
||||
model_id: &str,
|
||||
) -> Result<Arc<super::decision_model::LoadedDecisionModel>, InferenceError> {
|
||||
let handle = {
|
||||
let models = self.models.read().await;
|
||||
models.get(model_id).cloned()
|
||||
};
|
||||
match handle {
|
||||
Some(LoadedHandle::Decision(m)) => Ok(m),
|
||||
Some(_) => Err(InferenceError::WrongModality {
|
||||
model_id: model_id.to_string(),
|
||||
}),
|
||||
None => Err(InferenceError::ModelNotLoaded(model_id.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Answer one `/v1/systemone` request (#338): pick the served decision
|
||||
/// model and the checkpoint of it that answers, assemble the question
|
||||
/// sequences, score them in one forward pass, and shape the answers.
|
||||
///
|
||||
/// The served model is `served` when the gateway routed the request
|
||||
/// (the `x-helexa-decision-model` header), else the model the body's
|
||||
/// `model` names if one is loaded, else the one decision model this
|
||||
/// node serves. The body's `model` is then read a second time, as the
|
||||
/// reference server reads it: as an optional pin to a checkpoint of
|
||||
/// the family (`multilingual`, `typed-decisions`, …).
|
||||
///
|
||||
/// Questions are validated and tokenised before admission, so a
|
||||
/// malformed request never takes a slot.
|
||||
pub async fn systemone(
|
||||
&self,
|
||||
served: Option<&str>,
|
||||
request: cortex_core::decisions::SystemOneRequest,
|
||||
principal: Option<&str>,
|
||||
) -> Result<cortex_core::decisions::SystemOneResponse, SystemOneError> {
|
||||
use super::decision::{self, route};
|
||||
|
||||
let model = match served {
|
||||
Some(id) => self.decision_model(id).await?,
|
||||
None => {
|
||||
self.default_decision_model(request.model.as_deref())
|
||||
.await?
|
||||
}
|
||||
};
|
||||
if model.poisoned.load(Ordering::Acquire) {
|
||||
return Err(InferenceError::Other(anyhow::anyhow!(
|
||||
"decision model '{}' is poisoned; unload and reload it",
|
||||
model.model_id
|
||||
))
|
||||
.into());
|
||||
}
|
||||
|
||||
let pin = route::resolve_pin(request.model.as_deref());
|
||||
let mut routing = route::route(&model.model_id, &request.state, &request.questions, pin);
|
||||
let checkpoint = match model.checkpoint(Some(&routing.model)) {
|
||||
Some(c) => c,
|
||||
None if pin.is_some() => {
|
||||
return Err(decision::DecisionRequestError::new(
|
||||
404,
|
||||
format!(
|
||||
"checkpoint '{}' is not served by '{}' on this node",
|
||||
routing.model, model.model_id
|
||||
),
|
||||
)
|
||||
.into());
|
||||
}
|
||||
// Routed (not pinned) to a checkpoint this repo does not
|
||||
// carry: answer with the default one, and say so.
|
||||
None => {
|
||||
let fallback = model.checkpoint(None).ok_or_else(|| {
|
||||
InferenceError::Other(anyhow::anyhow!(
|
||||
"decision model '{}' has no checkpoints",
|
||||
model.model_id
|
||||
))
|
||||
})?;
|
||||
routing.reason = format!(
|
||||
"{}; the {} checkpoint is not loaded, so {} answered",
|
||||
routing.reason, routing.model, fallback.name
|
||||
);
|
||||
routing.model = fallback.name.clone();
|
||||
routing.repo = route::repo_of(&model.model_id, &fallback.name);
|
||||
fallback
|
||||
}
|
||||
};
|
||||
let lang = route::detected_language(&routing).map(str::to_string);
|
||||
|
||||
let prepared = decision::prepare(
|
||||
&request.state,
|
||||
&request.questions,
|
||||
&checkpoint.calibration,
|
||||
&checkpoint.tokenizer,
|
||||
decision::Budget {
|
||||
max_len: request.max_len,
|
||||
head_max_len: request.head_max_len,
|
||||
},
|
||||
)?;
|
||||
|
||||
let _permit = model.admission.enter(principal).await?;
|
||||
let outputs: Vec<decision::RowOutput> = if prepared.rows.is_empty() {
|
||||
Vec::new()
|
||||
} else {
|
||||
let batch = super::arch::laya::DecisionBatch {
|
||||
seqs: prepared
|
||||
.rows
|
||||
.iter()
|
||||
.map(|r| super::arch::laya::DecisionSeq {
|
||||
input_ids: r.input_ids.clone(),
|
||||
markers: r.markers.clone(),
|
||||
qtype: r.qtype as u8,
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
checkpoint
|
||||
.decide(batch)
|
||||
.await
|
||||
.map_err(|e| InferenceError::Other(anyhow::anyhow!("decision forward: {e:#}")))?
|
||||
.into_iter()
|
||||
.map(|row| decision::RowOutput {
|
||||
act_probability: row.act.first().copied().unwrap_or(0.0),
|
||||
logits: row.logits,
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
let (answers, usage) = decision::shape(
|
||||
&prepared,
|
||||
&outputs,
|
||||
&checkpoint.calibration,
|
||||
lang.as_deref(),
|
||||
);
|
||||
Ok(decision::shape::response(answers, usage, Some(routing)))
|
||||
}
|
||||
|
||||
/// Register an already-built decision model, bypassing the hub fetch
|
||||
/// `load_model` does — for tests that serve local checkpoints.
|
||||
#[cfg(test)]
|
||||
pub(crate) async fn register_decision_model(
|
||||
&self,
|
||||
model: super::decision_model::LoadedDecisionModel,
|
||||
) {
|
||||
self.models.write().await.insert(
|
||||
model.model_id.clone(),
|
||||
LoadedHandle::Decision(Arc::new(model)),
|
||||
);
|
||||
}
|
||||
|
||||
/// The decision model to serve a request no gateway routed: the one
|
||||
/// `requested` names if it is loaded, else the only (first) loaded
|
||||
/// decision model.
|
||||
async fn default_decision_model(
|
||||
&self,
|
||||
requested: Option<&str>,
|
||||
) -> Result<Arc<super::decision_model::LoadedDecisionModel>, InferenceError> {
|
||||
let models = self.models.read().await;
|
||||
if let Some(id) = requested
|
||||
&& let Some(handle) = models.get(id)
|
||||
{
|
||||
return match handle {
|
||||
LoadedHandle::Decision(m) => Ok(Arc::clone(m)),
|
||||
_ => Err(InferenceError::WrongModality {
|
||||
model_id: id.to_string(),
|
||||
}),
|
||||
};
|
||||
}
|
||||
let mut decision: Vec<&Arc<super::decision_model::LoadedDecisionModel>> = models
|
||||
.values()
|
||||
.filter_map(|h| match h {
|
||||
LoadedHandle::Decision(m) => Some(m),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
decision.sort_by(|a, b| a.model_id.cmp(&b.model_id));
|
||||
decision
|
||||
.first()
|
||||
.map(|m| Arc::clone(m))
|
||||
.ok_or_else(|| InferenceError::ModelNotLoaded("no decision model is loaded".into()))
|
||||
}
|
||||
|
||||
/// Run one text-to-image generation (#198). This is the internal
|
||||
/// path the `/v1/images/generations` handler (#200) fronts;
|
||||
/// admission, validation, tokenization, and PNG encoding happen
|
||||
@@ -6278,6 +6622,33 @@ pub enum InferenceError {
|
||||
Other(#[from] anyhow::Error),
|
||||
}
|
||||
|
||||
/// Why a `/v1/systemone` request was not answered (#338): a request the
|
||||
/// reference server would also reject (its status and `detail` text), or
|
||||
/// a serving failure shared with the other endpoints.
|
||||
#[derive(Debug)]
|
||||
pub enum SystemOneError {
|
||||
Request(super::decision::DecisionRequestError),
|
||||
Inference(InferenceError),
|
||||
}
|
||||
|
||||
impl From<super::decision::DecisionRequestError> for SystemOneError {
|
||||
fn from(e: super::decision::DecisionRequestError) -> Self {
|
||||
SystemOneError::Request(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<InferenceError> for SystemOneError {
|
||||
fn from(e: InferenceError) -> Self {
|
||||
SystemOneError::Inference(e)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<super::admission::AdmissionRejection> for SystemOneError {
|
||||
fn from(e: super::admission::AdmissionRejection) -> Self {
|
||||
SystemOneError::Inference(e.into())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<super::admission::AdmissionRejection> for InferenceError {
|
||||
fn from(rejection: super::admission::AdmissionRejection) -> Self {
|
||||
use super::admission::AdmissionRejection;
|
||||
|
||||
@@ -57,6 +57,9 @@ pub fn classify_files(
|
||||
) -> Vec<String> {
|
||||
match preflight::classify(filenames) {
|
||||
SourceFormat::Diffusers => vec!["image".to_string()],
|
||||
// A decision model answers typed questions (`/v1/systemone`);
|
||||
// it generates no text, so it must not claim "text".
|
||||
SourceFormat::DecisionModel => vec!["decision".to_string()],
|
||||
// A repo with no recognised weights (tokenizer-only, or an empty
|
||||
// entry) is not servable, so claiming "text" would be a lie of
|
||||
// the same kind this module exists to remove.
|
||||
@@ -101,6 +104,22 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn laya_repo_is_a_decision_model() {
|
||||
let files = [
|
||||
"rl_agent_config.json",
|
||||
"encoder/config.json",
|
||||
"model.safetensors",
|
||||
"tokenizer/tokenizer.json",
|
||||
];
|
||||
assert_eq!(
|
||||
classify_files(&files, || panic!(
|
||||
"config must not be read for a decision model"
|
||||
)),
|
||||
vec!["decision".to_string()]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dense_repo_without_vision_config_is_text_only() {
|
||||
let files = ["config.json", "tokenizer.json", "model.safetensors"];
|
||||
|
||||
@@ -0,0 +1,233 @@
|
||||
//! Per-checkpoint decision settings, read from `rl_agent_config.json`.
|
||||
//!
|
||||
//! Parsed with [`Json`], not `serde_json`: the fitted temperatures are
|
||||
//! written with 17 significant digits, and `serde_json`'s default float
|
||||
//! parser is not correctly rounded (it reads `1.7601518630981445` as
|
||||
//! `…443`).
|
||||
|
||||
use cortex_core::decisions::Json;
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// Default token budgets when a checkpoint's config omits them.
|
||||
pub const DEFAULT_MAX_LEN: usize = 512;
|
||||
pub const DEFAULT_HEAD_MAX_LEN: usize = 192;
|
||||
|
||||
/// Bounds on a usable calibration temperature. A fitted value below 1
|
||||
/// sharpens instead of softening; the shipped `choice:11+` bucket is 0.1006,
|
||||
/// which would publish a coin flip as a certainty, so the reference refuses
|
||||
/// to apply anything outside this range.
|
||||
pub const TEMP_MIN: f64 = 0.5;
|
||||
pub const TEMP_MAX: f64 = 5.0;
|
||||
|
||||
/// A temperature confined to [`TEMP_MIN`, `TEMP_MAX`]; a non-finite or
|
||||
/// non-numeric value falls back to 1.0.
|
||||
pub fn clamp_temperature(t: Option<f64>) -> f64 {
|
||||
match t {
|
||||
Some(t) if t.is_finite() => t.clamp(TEMP_MIN, TEMP_MAX),
|
||||
_ => 1.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Python's `float(v)` for a JSON value: numbers, bools, and numeric strings.
|
||||
fn py_float(v: &Json) -> Option<f64> {
|
||||
match v {
|
||||
Json::Int(d) => d.parse().ok(),
|
||||
Json::Float(x) => Some(*x),
|
||||
Json::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
|
||||
Json::String(s) => s.trim().parse().ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// One set of calibration temperatures: per question type, overridden per
|
||||
/// `(type, option-count bucket)` (see [`temp_bucket`]). Already clamped.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Temperatures {
|
||||
/// Indexed by [`super::QType`] (`choice`, `score`, `noul`).
|
||||
pub by_type: [f64; 3],
|
||||
/// `"choice:3-5"` → temperature, and so on.
|
||||
pub by_options: HashMap<String, f64>,
|
||||
}
|
||||
|
||||
impl Temperatures {
|
||||
/// The temperature for a question of type index `qtype` with `k` options.
|
||||
pub fn for_question(&self, qtype: usize, k: usize) -> f64 {
|
||||
self.by_options
|
||||
.get(&temp_bucket(qtype, k))
|
||||
.copied()
|
||||
.unwrap_or(self.by_type[qtype])
|
||||
}
|
||||
|
||||
fn from_json(
|
||||
temperature: Option<&Json>,
|
||||
by_options: Option<&Json>,
|
||||
fallback: Option<[f64; 3]>,
|
||||
) -> Result<Self, String> {
|
||||
let raw = match temperature {
|
||||
None => fallback.unwrap_or([1.0; 3]).map(Some),
|
||||
Some(Json::Array(a)) if a.len() == 3 => {
|
||||
[py_float(&a[0]), py_float(&a[1]), py_float(&a[2])]
|
||||
}
|
||||
Some(other) => {
|
||||
return Err(format!(
|
||||
"temperature must be a list of 3 floats, got {}",
|
||||
other.py_repr()
|
||||
));
|
||||
}
|
||||
};
|
||||
let by_options = match by_options {
|
||||
Some(Json::Object(m)) => m
|
||||
.iter()
|
||||
.map(|(k, v)| (k.clone(), clamp_temperature(py_float(v))))
|
||||
.collect(),
|
||||
_ => HashMap::new(),
|
||||
};
|
||||
Ok(Self {
|
||||
by_type: raw.map(clamp_temperature),
|
||||
by_options,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
/// The calibration bucket of a question: its type and option count.
|
||||
pub fn temp_bucket(qtype: usize, k: usize) -> String {
|
||||
let size = match k {
|
||||
0..=2 => "2",
|
||||
3..=5 => "3-5",
|
||||
6..=10 => "6-10",
|
||||
_ => "11+",
|
||||
};
|
||||
format!("{}:{size}", super::QType::NAMES[qtype])
|
||||
}
|
||||
|
||||
/// What the sequence builder and the calibration need from a checkpoint.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct DecisionConfig {
|
||||
/// Total sequence budget.
|
||||
pub max_len: usize,
|
||||
/// Budget for the question text and every option together.
|
||||
pub head_max_len: usize,
|
||||
pub temperatures: Temperatures,
|
||||
/// Per-language overrides, keyed by the lower-cased primary subtag
|
||||
/// (`"de"` for `de-AT`). Empty for the shipped checkpoints; an operator
|
||||
/// may fit them per language.
|
||||
pub lang_temperatures: HashMap<String, Temperatures>,
|
||||
}
|
||||
|
||||
impl DecisionConfig {
|
||||
/// Parse a checkpoint's `rl_agent_config.json`.
|
||||
pub fn from_json(text: &str) -> Result<Self, String> {
|
||||
let v = Json::parse(text.as_bytes()).map_err(|e| format!("rl_agent_config.json: {e}"))?;
|
||||
let usize_or = |key: &str, default: usize| match v.get(key) {
|
||||
Some(Json::Int(d)) => d.parse().unwrap_or(default),
|
||||
_ => default,
|
||||
};
|
||||
Ok(Self {
|
||||
max_len: usize_or("max_len", DEFAULT_MAX_LEN),
|
||||
head_max_len: usize_or("head_max_len", DEFAULT_HEAD_MAX_LEN),
|
||||
temperatures: Temperatures::from_json(
|
||||
v.get("temperature"),
|
||||
v.get("temperature_by_options"),
|
||||
None,
|
||||
)?,
|
||||
lang_temperatures: HashMap::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Add per-language temperatures, the shape the reference takes them in:
|
||||
/// `{"de": {"temperature": [..3], "temperature_by_options": {..}}}`. A
|
||||
/// language that omits `temperature` inherits the checkpoint's raw
|
||||
/// per-type values.
|
||||
pub fn with_lang_temperatures(mut self, spec: &Json) -> Result<Self, String> {
|
||||
let Json::Object(m) = spec else {
|
||||
return Err("lang_temperatures must be an object".into());
|
||||
};
|
||||
for (lang, cfg) in m {
|
||||
let t = Temperatures::from_json(
|
||||
cfg.get("temperature"),
|
||||
cfg.get("temperature_by_options"),
|
||||
Some(self.temperatures.by_type),
|
||||
)
|
||||
.map_err(|e| format!("language override {lang:?}: {e}"))?;
|
||||
let primary = lang.split('-').next().unwrap_or(lang).to_lowercase();
|
||||
self.lang_temperatures.insert(primary, t);
|
||||
}
|
||||
Ok(self)
|
||||
}
|
||||
|
||||
/// The temperatures to apply for a request detected (or declared) as
|
||||
/// `lang`, which may carry a region (`de-AT`).
|
||||
pub fn temperatures_for(&self, lang: Option<&str>) -> &Temperatures {
|
||||
lang.and_then(|l| {
|
||||
let primary = l.split('-').next().unwrap_or(l).to_lowercase();
|
||||
self.lang_temperatures.get(&primary)
|
||||
})
|
||||
.unwrap_or(&self.temperatures)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const ENGLISH: &str = r#"{"encoder": "answerdotai/ModernBERT-large", "head_layers": 2,
|
||||
"max_len": 512, "head_max_len": 192,
|
||||
"temperature": [1.6369030475616455, 1.2514300346374512, 1.983399510383606],
|
||||
"temperature_by_options": {"choice:3-5": 1.7601518630981445, "choice:11+": 0.10058280825614929}}"#;
|
||||
|
||||
#[test]
|
||||
fn shipped_choice_11_bucket_is_clamped() {
|
||||
let c = DecisionConfig::from_json(ENGLISH).unwrap();
|
||||
assert_eq!(c.temperatures.for_question(0, 12), TEMP_MIN);
|
||||
assert_eq!(c.temperatures.for_question(0, 4), 1.7601518630981445);
|
||||
// No bucket: the per-type value.
|
||||
assert_eq!(c.temperatures.for_question(1, 7), 1.2514300346374512);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn buckets_split_where_the_reference_splits() {
|
||||
let b: Vec<_> = [1, 2, 3, 5, 6, 10, 11]
|
||||
.iter()
|
||||
.map(|&k| temp_bucket(2, k))
|
||||
.collect();
|
||||
assert_eq!(
|
||||
b,
|
||||
[
|
||||
"noul:2",
|
||||
"noul:2",
|
||||
"noul:3-5",
|
||||
"noul:3-5",
|
||||
"noul:6-10",
|
||||
"noul:6-10",
|
||||
"noul:11+"
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn language_overrides_match_on_primary_subtag() {
|
||||
let c = DecisionConfig::from_json(ENGLISH)
|
||||
.unwrap()
|
||||
.with_lang_temperatures(
|
||||
&Json::parse(br#"{"DE": {"temperature": [1.0, 1.4, 9.0]}}"#).unwrap(),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
c.temperatures_for(Some("de-AT")).by_type,
|
||||
[1.0, 1.4, TEMP_MAX]
|
||||
);
|
||||
assert_eq!(c.temperatures_for(Some("fr")), &c.temperatures);
|
||||
assert_eq!(c.temperatures_for(None), &c.temperatures);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_temperatures_are_refused_or_neutralised() {
|
||||
assert!(DecisionConfig::from_json(r#"{"temperature": [1.0]}"#).is_err());
|
||||
let c = DecisionConfig::from_json(r#"{"temperature": [" 2.5 ", "x", null]}"#).unwrap();
|
||||
assert_eq!(c.temperatures.by_type, [2.5, 1.0, 1.0]);
|
||||
assert_eq!(
|
||||
(c.max_len, c.head_max_len),
|
||||
(DEFAULT_MAX_LEN, DEFAULT_HEAD_MAX_LEN)
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
//! `CandleHarness::systemone` end to end: request in, routed, validated,
|
||||
//! scored and shaped response out (#338).
|
||||
//!
|
||||
//! The tiny checkpoint (`testdata/laya/tiny/`) exercises the plumbing in
|
||||
//! CI. The released-weights test replays every reference case through
|
||||
//! the same entry point the HTTP handler calls and compares the whole
|
||||
//! response — answers, usage and routing — with what the reference
|
||||
//! server returned; it needs `LAYA_REFERENCE_DIR`.
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
|
||||
use cortex_core::decisions::{Json, SystemOneRequest, limits};
|
||||
|
||||
use crate::harness::candle::{CandleHarness, SystemOneError};
|
||||
use crate::harness::decision_model::{CheckpointFiles, DecisionCheckpoint, LoadedDecisionModel};
|
||||
|
||||
fn testdata() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
|
||||
}
|
||||
|
||||
fn harness() -> Arc<CandleHarness> {
|
||||
CandleHarness::new(
|
||||
"http://localhost:13131".into(),
|
||||
&crate::config::CandleHarnessConfig::default(),
|
||||
)
|
||||
}
|
||||
|
||||
/// A decision model made of the given `(routing name, directory)`
|
||||
/// checkpoints, loaded on the CPU in f32.
|
||||
async fn decision_model(model_id: &str, checkpoints: &[(&str, PathBuf)]) -> LoadedDecisionModel {
|
||||
let mut loaded = Vec::new();
|
||||
for (name, dir) in checkpoints {
|
||||
let files = CheckpointFiles {
|
||||
tokenizer: if dir.join("tokenizer/tokenizer.json").exists() {
|
||||
dir.join("tokenizer/tokenizer.json")
|
||||
} else {
|
||||
dir.join("tokenizer.json")
|
||||
},
|
||||
dir: dir.clone(),
|
||||
};
|
||||
loaded.push(
|
||||
DecisionCheckpoint::load(name.to_string(), files, model_id, None)
|
||||
.await
|
||||
.unwrap_or_else(|e| panic!("load {name}: {e:#}")),
|
||||
);
|
||||
}
|
||||
LoadedDecisionModel {
|
||||
model_id: model_id.to_string(),
|
||||
spec: serde_json::from_value(serde_json::json!({
|
||||
"model_id": model_id,
|
||||
"harness": "candle",
|
||||
}))
|
||||
.expect("model spec"),
|
||||
devices: vec![0],
|
||||
checkpoints: loaded,
|
||||
poisoned: Arc::new(AtomicBool::new(false)),
|
||||
admission: crate::harness::admission::AdmissionController::new(
|
||||
&crate::config::AdmissionConfig::default(),
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn request(body: &str) -> SystemOneRequest {
|
||||
SystemOneRequest::parse(body.as_bytes(), limits::DEFAULT_MAX_TOKEN_BUDGET)
|
||||
.expect("valid request")
|
||||
}
|
||||
|
||||
async fn tiny_harness() -> Arc<CandleHarness> {
|
||||
let h = harness();
|
||||
h.register_decision_model(
|
||||
decision_model("test/tiny", &[("english", testdata().join("tiny"))]).await,
|
||||
)
|
||||
.await;
|
||||
h
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn answers_every_question_in_order_with_exact_usage() {
|
||||
let h = tiny_harness().await;
|
||||
let req = request(
|
||||
r#"{"state": "a b a", "questions": {
|
||||
"z_last_alphabetically": {"type": "noul", "instructions": "a?"},
|
||||
"a_first": {"type": "choice", "instructions": "b?", "criteria": ["a", "b", "a b"]}
|
||||
}}"#,
|
||||
);
|
||||
let resp = h.systemone(None, req, None).await.expect("answered");
|
||||
let json = serde_json::to_value(&resp).unwrap();
|
||||
// Answers come back in request order, not sorted.
|
||||
let keys: Vec<&str> = resp.answers.0.iter().map(|(k, _)| k.as_str()).collect();
|
||||
assert_eq!(keys, ["z_last_alphabetically", "a_first"]);
|
||||
assert_eq!(json["answers"]["z_last_alphabetically"]["type"], "noul");
|
||||
let probs = json["answers"]["a_first"]["probabilities"]
|
||||
.as_object()
|
||||
.unwrap();
|
||||
let total: f64 = probs.values().map(|p| p.as_f64().unwrap()).sum();
|
||||
assert!((total - 1.0).abs() < 1e-3, "{probs:?}");
|
||||
// Usage is the exact number of tokens encoded, output tokens zero.
|
||||
assert!(json["usage"]["input_tokens"].as_u64().unwrap() > 0);
|
||||
assert_eq!(json["usage"]["output_tokens"], 0);
|
||||
assert_eq!(json["routing"]["model"], "english");
|
||||
assert_eq!(json["routing"]["repo"], "test/tiny");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pinning_a_checkpoint_the_model_lacks_is_404() {
|
||||
let h = tiny_harness().await;
|
||||
let req = request(
|
||||
r#"{"state": "a", "model": "multilingual",
|
||||
"questions": {"q": {"type": "noul", "instructions": "a?"}}}"#,
|
||||
);
|
||||
match h.systemone(None, req, None).await {
|
||||
Err(SystemOneError::Request(e)) => assert_eq!(e.status, 404, "{}", e.detail),
|
||||
other => panic!("expected a 404, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn routing_to_a_missing_checkpoint_falls_back_and_says_so() {
|
||||
// Devanagari routes to `multilingual`, which this model lacks.
|
||||
let h = tiny_harness().await;
|
||||
let req = request(
|
||||
r#"{"state": "मुझे पैसे वापस चाहिए",
|
||||
"questions": {"q": {"type": "noul", "instructions": "a?"}}}"#,
|
||||
);
|
||||
let resp = h.systemone(None, req, None).await.expect("answered");
|
||||
let routing = resp.routing.expect("routing");
|
||||
assert_eq!(routing.model, "english");
|
||||
assert!(
|
||||
routing
|
||||
.reason
|
||||
.contains("multilingual checkpoint is not loaded"),
|
||||
"{}",
|
||||
routing.reason
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn invalid_questions_are_rejected_before_admission() {
|
||||
let h = tiny_harness().await;
|
||||
let req =
|
||||
request(r#"{"state": "a", "questions": {"q": {"type": "maybe", "instructions": "a?"}}}"#);
|
||||
match h.systemone(None, req, None).await {
|
||||
Err(SystemOneError::Request(e)) => assert_eq!(e.status, 422, "{}", e.detail),
|
||||
other => panic!("expected a 422, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn a_gateway_routed_model_that_is_not_a_decision_model_is_refused() {
|
||||
let h = tiny_harness().await;
|
||||
let req = request(r#"{"state": "a", "questions": {}}"#);
|
||||
match h.systemone(Some("not/loaded"), req, None).await {
|
||||
Err(SystemOneError::Inference(crate::harness::candle::InferenceError::ModelNotLoaded(
|
||||
id,
|
||||
))) => {
|
||||
assert_eq!(id, "not/loaded")
|
||||
}
|
||||
other => panic!("expected ModelNotLoaded, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Every reference case through `systemone`, compared with the reference
|
||||
/// server's whole response. Routing and usage must match exactly;
|
||||
/// probabilities to within the engine's measured f32 distance from the
|
||||
/// reference (a last-digit wobble after 4-decimal rounding); the chosen
|
||||
/// answer exactly.
|
||||
#[tokio::test]
|
||||
#[ignore = "needs the released weights: set LAYA_REFERENCE_DIR"]
|
||||
async fn released_family_answers_like_the_reference_server() {
|
||||
let root = PathBuf::from(std::env::var("LAYA_REFERENCE_DIR").expect("LAYA_REFERENCE_DIR"));
|
||||
let h = harness();
|
||||
h.register_decision_model(
|
||||
decision_model(
|
||||
"convaiinnovations/laya",
|
||||
&[
|
||||
("english", root.clone()),
|
||||
("multilingual", root.join("multilingual")),
|
||||
("typed-decisions", root.join("typed-decisions")),
|
||||
],
|
||||
)
|
||||
.await,
|
||||
)
|
||||
.await;
|
||||
|
||||
let text = std::fs::read_to_string(testdata().join("reference.json")).unwrap();
|
||||
let reference = Json::parse(text.as_bytes()).unwrap();
|
||||
let Some(Json::Array(cases)) = reference.get("cases") else {
|
||||
panic!("no cases")
|
||||
};
|
||||
let mut max_dp = 0.0f64;
|
||||
for case in cases {
|
||||
let name = case.get("name").and_then(Json::as_str).unwrap();
|
||||
let body = serde_json::to_vec(case.get("request").unwrap()).unwrap();
|
||||
let req = SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET).unwrap();
|
||||
let resp = h
|
||||
.systemone(Some("convaiinnovations/laya"), req, None)
|
||||
.await
|
||||
.unwrap_or_else(|e| panic!("{name}: {e:?}"));
|
||||
let want_json = case.get("response").unwrap();
|
||||
// Routing and usage compared as serialized text through `Json`:
|
||||
// exact, key order included. serde_json would re-round the
|
||||
// reference's floats on the way in.
|
||||
for key in ["routing", "usage", "model"] {
|
||||
let got = match key {
|
||||
"routing" => serde_json::to_string(&resp.routing).unwrap(),
|
||||
"usage" => serde_json::to_string(&resp.usage).unwrap(),
|
||||
_ => serde_json::to_string(&resp.model).unwrap(),
|
||||
};
|
||||
let want = serde_json::to_string(want_json.get(key).unwrap()).unwrap();
|
||||
assert_eq!(got, want, "{name}: {key}");
|
||||
}
|
||||
let got = serde_json::to_value(&resp).unwrap();
|
||||
let want: serde_json::Value =
|
||||
serde_json::from_str(&serde_json::to_string(want_json).unwrap()).unwrap();
|
||||
let (ga, wa) = (
|
||||
got["answers"].as_object().unwrap(),
|
||||
want["answers"].as_object().unwrap(),
|
||||
);
|
||||
assert_eq!(
|
||||
ga.keys().collect::<Vec<_>>(),
|
||||
wa.keys().collect::<Vec<_>>(),
|
||||
"{name}: answer ids"
|
||||
);
|
||||
for (qid, w) in wa {
|
||||
let g = &ga[qid];
|
||||
for key in ["choice", "type", "legend"] {
|
||||
assert_eq!(g.get(key), w.get(key), "{name}/{qid}: {key}");
|
||||
}
|
||||
if let Some(wp) = w["probabilities"].as_object() {
|
||||
for (opt, p) in wp {
|
||||
let d = (g["probabilities"][opt].as_f64().unwrap() - p.as_f64().unwrap()).abs();
|
||||
max_dp = max_dp.max(d);
|
||||
}
|
||||
}
|
||||
for key in ["noul", "score", "confidence", "answer_confidence"] {
|
||||
if let (Some(gv), Some(wv)) = (
|
||||
g.get(key).and_then(|v| v.as_f64()),
|
||||
w.get(key).and_then(|v| v.as_f64()),
|
||||
) {
|
||||
max_dp = max_dp.max((gv - wv).abs());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
eprintln!("released family: max |Δ| over probabilities and confidences {max_dp:.1e}");
|
||||
assert!(max_dp <= 5e-4, "max |Δ| {max_dp:e}");
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
//! Laya decision models: everything around the forward pass (#335).
|
||||
//!
|
||||
//! A Laya checkpoint answers typed questions over a *state* (text, a JSON
|
||||
//! document, a conversation) in one encoder forward pass: each question
|
||||
//! becomes one sequence,
|
||||
//!
|
||||
//! ```text
|
||||
//! [CLS] "<type> question: <instructions>" [SEP] [MASK] opt0 [MASK] opt1 … [SEP] state [SEP]
|
||||
//! ```
|
||||
//!
|
||||
//! and the decision head scores the hidden state at each option's `[MASK]`
|
||||
//! marker. This module is the part that touches no tensor:
|
||||
//!
|
||||
//! - [`question`] — validate and normalise the question definitions, and
|
||||
//! render each option's text;
|
||||
//! - [`sequence`] — assemble the token sequences under the model's token
|
||||
//! budgets ([`prepare`]);
|
||||
//! - [`shape`] — turn per-option logits into calibrated, typed answers
|
||||
//! ([`shape()`]);
|
||||
//! - [`config`] / [`tokenizer`] — the per-checkpoint pieces those need.
|
||||
//!
|
||||
//! The reference is the `laya` Python SDK (`laya/common.py`,
|
||||
//! `laya/agent.py`, `laya/serve.py`), whose behaviour this reproduces token
|
||||
//! for token: the parity tests replay fixtures recorded from it
|
||||
//! (`script/laya-reference.py`). The calibration arithmetic follows numpy's
|
||||
//! float32 evaluation order, so rounded probabilities match to the last
|
||||
//! published digit.
|
||||
//!
|
||||
//! The wire types (`/v1/systemone`) live in [`cortex_core::decisions`].
|
||||
|
||||
pub mod config;
|
||||
pub mod question;
|
||||
pub mod route;
|
||||
pub mod sequence;
|
||||
pub mod shape;
|
||||
pub mod tokenizer;
|
||||
|
||||
#[cfg(test)]
|
||||
mod handler_tests;
|
||||
|
||||
pub use config::DecisionConfig;
|
||||
pub use cortex_core::decisions::DecisionRequestError;
|
||||
pub use question::{QType, Question};
|
||||
pub use sequence::{Budget, OptionStats, Prepared, Row, prepare};
|
||||
pub use shape::{RowOutput, shape};
|
||||
pub use tokenizer::{DecisionTokenizer, SpecialTokens, TextEncoder};
|
||||
|
||||
/// The constant `model` field of every response: the decision head's name.
|
||||
/// Which checkpoint answered is reported in `routing`.
|
||||
pub const DECISION_MODEL_NAME: &str = "laya-rl-agent";
|
||||
|
||||
/// A request error in the reference's `ValueError` class: the server
|
||||
/// answers it with 422 and the message as `detail`.
|
||||
pub(crate) fn unprocessable(detail: impl Into<String>) -> DecisionRequestError {
|
||||
DecisionRequestError::new(422, detail)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod parity_tests;
|
||||
@@ -0,0 +1,209 @@
|
||||
//! Replays the reference fixtures (`testdata/laya/reference.json`, recorded
|
||||
//! by `script/laya-reference.py`) through [`prepare`] and [`shape`].
|
||||
//!
|
||||
//! Token ids and marker positions must match exactly. The answers and usage
|
||||
//! must serialize to exactly the same JSON — same keys, same order, same
|
||||
//! rounded values — when fed the reference's own raw logits.
|
||||
//!
|
||||
//! The English tokenizer is committed. The multilingual one is 34 MB, so the
|
||||
//! multilingual cases run only when `LAYA_ML_TOKENIZER_DIR` names a copy of
|
||||
//! the checkpoint's `multilingual/tokenizer/` directory.
|
||||
|
||||
use super::*;
|
||||
use cortex_core::decisions::{Json, SystemOneRequest, limits};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
fn testdata() -> PathBuf {
|
||||
Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya")
|
||||
}
|
||||
|
||||
fn reference() -> Json {
|
||||
Json::parse(&std::fs::read(testdata().join("reference.json")).expect("fixtures")).expect("json")
|
||||
}
|
||||
|
||||
fn to_value(j: &Json) -> serde_json::Value {
|
||||
serde_json::to_value(j).expect("serializable")
|
||||
}
|
||||
|
||||
fn tokenizer_for(checkpoint: &str) -> Option<DecisionTokenizer> {
|
||||
let dir = match checkpoint {
|
||||
// typed-decisions ships the English tokenizer, byte for byte.
|
||||
"english" | "typed-decisions" => testdata().join("tokenizer"),
|
||||
"multilingual" => PathBuf::from(std::env::var_os("LAYA_ML_TOKENIZER_DIR")?),
|
||||
other => panic!("unknown checkpoint {other}"),
|
||||
};
|
||||
Some(DecisionTokenizer::from_dir(&dir).expect("tokenizer"))
|
||||
}
|
||||
|
||||
fn config_for(reference: &Json, checkpoint: &str) -> DecisionConfig {
|
||||
let cfg = reference
|
||||
.get("header")
|
||||
.and_then(|h| h.get("checkpoints"))
|
||||
.and_then(|c| c.get(checkpoint))
|
||||
.and_then(|c| c.get("cfg"))
|
||||
.expect("checkpoint cfg");
|
||||
DecisionConfig::from_json(&serde_json::to_string(cfg).unwrap()).unwrap()
|
||||
}
|
||||
|
||||
fn array(j: &Json) -> &[Json] {
|
||||
match j {
|
||||
Json::Array(a) => a,
|
||||
other => panic!("expected a list, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn num(j: &Json) -> f64 {
|
||||
match j {
|
||||
Json::Int(d) => d.parse().unwrap(),
|
||||
Json::Float(x) => *x,
|
||||
other => panic!("expected a number, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Replay every case whose tokenizer is available; return how many ran per
|
||||
/// checkpoint.
|
||||
fn replay(filter: impl Fn(&str) -> bool) -> Vec<(String, usize)> {
|
||||
let reference = reference();
|
||||
let mut ran: Vec<(String, usize)> = Vec::new();
|
||||
for case in array(reference.get("cases").unwrap()) {
|
||||
let name = case.get("name").and_then(Json::as_str).unwrap();
|
||||
if !filter(name) {
|
||||
continue;
|
||||
}
|
||||
let response = case.get("response").unwrap();
|
||||
let routing = response.get("routing").unwrap();
|
||||
let checkpoint = routing.get("model").and_then(Json::as_str).unwrap();
|
||||
let Some(tok) = tokenizer_for(checkpoint) else {
|
||||
continue;
|
||||
};
|
||||
let cfg = config_for(&reference, checkpoint);
|
||||
|
||||
let body = serde_json::to_vec(case.get("request").unwrap()).unwrap();
|
||||
let req = SystemOneRequest::parse(&body, limits::DEFAULT_MAX_TOKEN_BUDGET)
|
||||
.unwrap_or_else(|e| panic!("{name}: {e}"));
|
||||
let budget = Budget {
|
||||
max_len: req.max_len,
|
||||
head_max_len: req.head_max_len,
|
||||
};
|
||||
let prepared = prepare(&req.state, &req.questions, &cfg, &tok, budget)
|
||||
.unwrap_or_else(|e| panic!("{name}: {e}"));
|
||||
|
||||
let items = array(case.get("items").unwrap());
|
||||
assert_eq!(prepared.rows.len(), items.len(), "{name}: row count");
|
||||
let mut outputs = Vec::new();
|
||||
for (i, (row, item)) in prepared.rows.iter().zip(items).enumerate() {
|
||||
let ids: Vec<u32> = array(item.get("input_ids").unwrap())
|
||||
.iter()
|
||||
.map(|v| num(v) as u32)
|
||||
.collect();
|
||||
let markers: Vec<usize> = array(item.get("markers").unwrap())
|
||||
.iter()
|
||||
.map(|v| num(v) as usize)
|
||||
.collect();
|
||||
assert_eq!(row.input_ids, ids, "{name}: row {i} input_ids");
|
||||
assert_eq!(row.markers, markers, "{name}: row {i} markers");
|
||||
assert_eq!(row.qtype as usize, num(item.get("qtype").unwrap()) as usize);
|
||||
let opts = item.get("options").unwrap();
|
||||
let s = prepared.stats[i];
|
||||
assert_eq!(
|
||||
(s.options, s.options_distinct, s.tokens_per_option),
|
||||
(
|
||||
num(opts.get("options").unwrap()) as usize,
|
||||
num(opts.get("options_distinct").unwrap()) as usize,
|
||||
match opts.get("tokens_per_option").unwrap() {
|
||||
Json::Null => None,
|
||||
v => Some(num(v) as usize),
|
||||
}
|
||||
),
|
||||
"{name}: row {i} option stats"
|
||||
);
|
||||
outputs.push(RowOutput {
|
||||
logits: array(item.get("logits").unwrap())
|
||||
.iter()
|
||||
.map(|v| num(v) as f32)
|
||||
.collect(),
|
||||
act_probability: num(&array(item.get("act").unwrap())[0]) as f32,
|
||||
});
|
||||
}
|
||||
|
||||
let lang = routing
|
||||
.get("detection")
|
||||
.and_then(|d| d.get("language"))
|
||||
.and_then(Json::as_str);
|
||||
let (answers, usage) = shape(&prepared, &outputs, &cfg, lang);
|
||||
// Serialized strings, not values: key order is part of the contract.
|
||||
assert_eq!(
|
||||
serde_json::to_string(&answers).unwrap(),
|
||||
serde_json::to_string(response.get("answers").unwrap()).unwrap(),
|
||||
"{name}: answers"
|
||||
);
|
||||
assert_eq!(
|
||||
serde_json::to_value(&usage).unwrap(),
|
||||
to_value(response.get("usage").unwrap()),
|
||||
"{name}: usage"
|
||||
);
|
||||
match ran.iter_mut().find(|(c, _)| c == checkpoint) {
|
||||
Some((_, n)) => *n += 1,
|
||||
None => ran.push((checkpoint.to_string(), 1)),
|
||||
}
|
||||
}
|
||||
ran
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn every_english_fixture_matches_the_reference() {
|
||||
let ran = replay(|_| true);
|
||||
let count = |c: &str| ran.iter().find(|(n, _)| n == c).map_or(0, |(_, n)| *n);
|
||||
// The corpus routes 19 cases to English and 1 to typed-decisions; a
|
||||
// drop means a case silently stopped running.
|
||||
assert_eq!(count("english"), 19, "{ran:?}");
|
||||
assert_eq!(count("typed-decisions"), 1, "{ran:?}");
|
||||
if std::env::var_os("LAYA_ML_TOKENIZER_DIR").is_some() {
|
||||
assert_eq!(count("multilingual"), 8, "{ran:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_questions_need_no_forward() {
|
||||
let reference = reference();
|
||||
let tok = tokenizer_for("english").unwrap();
|
||||
let cfg = config_for(&reference, "english");
|
||||
let prepared = prepare(
|
||||
&Json::String("x".into()),
|
||||
&[],
|
||||
&cfg,
|
||||
&tok,
|
||||
Budget::default(),
|
||||
)
|
||||
.unwrap();
|
||||
assert!(prepared.rows.is_empty());
|
||||
let (answers, usage) = shape(&prepared, &[], &cfg, None);
|
||||
assert_eq!(serde_json::to_string(&answers).unwrap(), "{}");
|
||||
assert_eq!(
|
||||
serde_json::to_string(&usage).unwrap(),
|
||||
r#"{"input_tokens":0,"output_tokens":0}"#
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn options_that_cannot_all_fit_are_refused() {
|
||||
let reference = reference();
|
||||
let tok = tokenizer_for("english").unwrap();
|
||||
let cfg = config_for(&reference, "english");
|
||||
let questions = Json::parse(
|
||||
br#"{"q": {"type": "choice", "instructions": "i", "criteria": ["a", "b", "c", "d"]}}"#,
|
||||
)
|
||||
.unwrap();
|
||||
let Json::Object(questions) = questions else {
|
||||
unreachable!()
|
||||
};
|
||||
let budget = Budget {
|
||||
max_len: Some(10),
|
||||
head_max_len: None,
|
||||
};
|
||||
let err = prepare(&Json::String("x".into()), &questions, &cfg, &tok, budget).unwrap_err();
|
||||
assert_eq!(
|
||||
(err.status, err.detail.as_str()),
|
||||
(422, "question 'q' options exceed head_max_len=192")
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,458 @@
|
||||
//! Question validation, normalisation and option rendering.
|
||||
//!
|
||||
//! Every rejection here is a 422 whose message names the question and what
|
||||
//! to fix — the same text the reference returns, so a client written
|
||||
//! against it sees the same errors.
|
||||
|
||||
use super::{DecisionRequestError, unprocessable};
|
||||
use cortex_core::decisions::{HashKey, Json};
|
||||
use std::collections::HashMap;
|
||||
|
||||
/// The three decision primitives. The discriminant is the index the model's
|
||||
/// type embedding and the calibration tables use.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum QType {
|
||||
/// Pick one of the caller's labels.
|
||||
Choice = 0,
|
||||
/// An ordinal level on the caller's rubric.
|
||||
Score = 1,
|
||||
/// Yes/no: the probability that a statement holds.
|
||||
Noul = 2,
|
||||
}
|
||||
|
||||
impl QType {
|
||||
pub const NAMES: [&'static str; 3] = ["choice", "score", "noul"];
|
||||
|
||||
pub fn name(self) -> &'static str {
|
||||
Self::NAMES[self as usize]
|
||||
}
|
||||
|
||||
fn parse(s: &str) -> Option<Self> {
|
||||
match s {
|
||||
"choice" => Some(Self::Choice),
|
||||
"score" => Some(Self::Score),
|
||||
"noul" => Some(Self::Noul),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A validated question, normalised the way the reference normalises it.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Question {
|
||||
pub id: String,
|
||||
pub kind: QType,
|
||||
/// The instruction text: the caller's string, or the JSON rendering of
|
||||
/// a structured value.
|
||||
pub instructions: String,
|
||||
/// Each option's rendered text, in answer-index order. For `noul` this
|
||||
/// is always `[false, true]`, so index 1 is P(true).
|
||||
pub options: Vec<String>,
|
||||
/// `choice`: each option's label as the caller wrote it (the answer key).
|
||||
/// `score`: each level's description (the legend). `noul`: empty.
|
||||
pub labels: Vec<Json>,
|
||||
}
|
||||
|
||||
impl Question {
|
||||
/// Validate one question definition and normalise it.
|
||||
pub fn from_definition(id: &str, def: &Json) -> Result<Self, DecisionRequestError> {
|
||||
check_question(id, def)?;
|
||||
Ok(to_internal(id, def))
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate and normalise every question, in request order. The first
|
||||
/// invalid question is reported; nothing is normalised until all pass.
|
||||
pub fn questions_from_request(
|
||||
questions: &[(String, Json)],
|
||||
) -> Result<Vec<Question>, DecisionRequestError> {
|
||||
for (id, def) in questions {
|
||||
check_question(id, def)?;
|
||||
}
|
||||
Ok(questions
|
||||
.iter()
|
||||
.map(|(id, def)| to_internal(id, def))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn q(id: &str) -> String {
|
||||
Json::String(id.to_string()).py_repr()
|
||||
}
|
||||
|
||||
fn check_question(id: &str, def: &Json) -> Result<(), DecisionRequestError> {
|
||||
let qid = q(id);
|
||||
if !matches!(def, Json::Object(_)) {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: definition must be a dict, got {}",
|
||||
def.py_type_name()
|
||||
)));
|
||||
}
|
||||
let t = def.get("type").unwrap_or(&Json::Null);
|
||||
let Some(kind) = t.as_str().and_then(QType::parse) else {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: unknown type {}; use one of ['choice', 'noul', 'score']",
|
||||
t.py_repr()
|
||||
)));
|
||||
};
|
||||
if def.get("instructions").is_none() {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: no 'instructions'; add the text the model should answer"
|
||||
)));
|
||||
}
|
||||
let crit = def.get("criteria");
|
||||
match kind {
|
||||
QType::Choice => check_choice(&qid, crit)?,
|
||||
QType::Score => {
|
||||
let Some(Json::Array(levels)) = crit else {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a score question takes 'criteria' as a list of level \
|
||||
descriptions, index 0 first"
|
||||
)));
|
||||
};
|
||||
if levels.is_empty() {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a score question needs at least one level"
|
||||
)));
|
||||
}
|
||||
if let Some(i) = levels.iter().position(|l| *l == Json::Null) {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: score level {i} is null; give every level a description, \
|
||||
index 0 first"
|
||||
)));
|
||||
}
|
||||
}
|
||||
QType::Noul => match crit {
|
||||
None | Some(Json::Null) => {}
|
||||
Some(Json::Object(m)) => {
|
||||
let mut keys: Vec<String> = m.iter().map(|(k, _)| k.to_lowercase()).collect();
|
||||
keys.sort();
|
||||
keys.dedup();
|
||||
if keys.iter().any(|k| k != "true" && k != "false") {
|
||||
let listed = Json::Array(keys.into_iter().map(Json::String).collect());
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a noul question takes 'criteria' keyed only \
|
||||
'true'/'false' (either or both, and omitted is fine), got {}. Those \
|
||||
keys are the option texts the model reads; any other key was silently \
|
||||
dropped and replaced with the defaults. If you want the answer worded \
|
||||
differently, keep 'criteria' keyed 'true'/'false' and set 'labels' \
|
||||
instead.",
|
||||
listed.py_repr()
|
||||
)));
|
||||
}
|
||||
}
|
||||
Some(_) => {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a noul question takes 'criteria' as a dict with optional \
|
||||
'true'/'false' descriptions, or omits it"
|
||||
)));
|
||||
}
|
||||
},
|
||||
}
|
||||
if let Some(labels) = def.get("labels") {
|
||||
if kind != QType::Noul {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: 'labels' is only supported for noul questions"
|
||||
)));
|
||||
}
|
||||
resolve_noul_labels(labels).map_err(|e| unprocessable(format!("question {qid}: {e}")))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn check_choice(qid: &str, crit: Option<&Json>) -> Result<(), DecisionRequestError> {
|
||||
let labels: Vec<Json> = match crit {
|
||||
Some(Json::Object(m)) => m.iter().map(|(k, _)| Json::String(k.clone())).collect(),
|
||||
Some(Json::Array(a)) => a.clone(),
|
||||
_ => {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a choice question takes 'criteria' as a dict of label -> \
|
||||
description, or a list of labels"
|
||||
)));
|
||||
}
|
||||
};
|
||||
if labels.is_empty() {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: a choice question needs at least one criterion"
|
||||
)));
|
||||
}
|
||||
for (i, label) in labels.iter().enumerate() {
|
||||
match label {
|
||||
Json::Array(_) | Json::Object(_) => {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: choice label {i} is a {}; a label is rendered as option \
|
||||
text and used as the answer key, so it must be a scalar (a string, number \
|
||||
or bool), got {}",
|
||||
label.py_type_name(),
|
||||
label.py_repr()
|
||||
)));
|
||||
}
|
||||
Json::Null => {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: choice label {i} is null; a label is rendered as option \
|
||||
text and used as the answer key, so it must be a string, number or bool \
|
||||
-- a null label renders as the text \"None\" while its answer key is \
|
||||
\"null\""
|
||||
)));
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
// Only the list form can repeat a key: an object's keys are unique.
|
||||
if matches!(crit, Some(Json::Array(_))) {
|
||||
let mut seen: HashMap<HashKey, usize> = HashMap::new();
|
||||
for (i, label) in labels.iter().enumerate() {
|
||||
let key = label.hash_key().expect("scalars are hashable");
|
||||
if key == HashKey::Unique {
|
||||
continue;
|
||||
}
|
||||
if let Some(first) = seen.get(&key) {
|
||||
return Err(unprocessable(format!(
|
||||
"question {qid}: choice label {i} ({}) repeats label {first}; the labels are \
|
||||
the answer keys, so every option needs its own (1, 1.0 and True are one key)",
|
||||
label.py_repr()
|
||||
)));
|
||||
}
|
||||
seen.insert(key, i);
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
const DEFAULT_NOUL_FALSE: &str = "no, the statement does not hold";
|
||||
const DEFAULT_NOUL_TRUE: &str = "yes, the statement holds";
|
||||
|
||||
/// The words a noul question shows the model for its two outcomes.
|
||||
fn resolve_noul_labels(labels: &Json) -> Result<(String, String), &'static str> {
|
||||
const BAD: &str =
|
||||
"noul labels must map exactly 'false' and 'true' to distinct non-empty strings";
|
||||
let Json::Object(m) = labels else {
|
||||
return if *labels == Json::Null {
|
||||
Ok(("false".into(), "true".into()))
|
||||
} else {
|
||||
Err(BAD)
|
||||
};
|
||||
};
|
||||
let mut keys: Vec<&str> = m.iter().map(|(k, _)| k.as_str()).collect();
|
||||
keys.sort_unstable();
|
||||
if keys != ["false", "true"] {
|
||||
return Err(BAD);
|
||||
}
|
||||
let (Some(Json::String(f)), Some(Json::String(t))) = (labels.get("false"), labels.get("true"))
|
||||
else {
|
||||
return Err(BAD);
|
||||
};
|
||||
let (f, t) = (py_strip(f), py_strip(t));
|
||||
if f.is_empty() || t.is_empty() || f == t {
|
||||
return Err(BAD);
|
||||
}
|
||||
Ok((f.to_string(), t.to_string()))
|
||||
}
|
||||
|
||||
/// Python's `str.strip()` with no argument: Unicode whitespace.
|
||||
fn py_strip(s: &str) -> &str {
|
||||
s.trim_matches(|c: char| c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c))
|
||||
}
|
||||
|
||||
/// `render_criterion`: a string as itself, anything else as JSON.
|
||||
fn render_criterion(v: &Json) -> String {
|
||||
match v {
|
||||
Json::String(s) => s.clone(),
|
||||
other => other.dumps(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Only `None` and `""` mean "no description"; `0` and `false` are values.
|
||||
fn is_blank(v: &Json) -> bool {
|
||||
matches!(v, Json::Null) || matches!(v, Json::String(s) if s.is_empty())
|
||||
}
|
||||
|
||||
fn to_internal(id: &str, def: &Json) -> Question {
|
||||
let kind = def
|
||||
.get("type")
|
||||
.and_then(Json::as_str)
|
||||
.and_then(QType::parse)
|
||||
.expect("validated");
|
||||
let instructions = match def.get("instructions").expect("validated") {
|
||||
Json::String(s) => s.clone(),
|
||||
other => other.dumps(),
|
||||
};
|
||||
let crit = def.get("criteria");
|
||||
let (options, labels) = match kind {
|
||||
QType::Choice => {
|
||||
let pairs: Vec<(Json, Json)> = match crit {
|
||||
Some(Json::Object(m)) => m
|
||||
.iter()
|
||||
.map(|(k, v)| (Json::String(k.clone()), v.clone()))
|
||||
.collect(),
|
||||
Some(Json::Array(a)) => a.iter().map(|l| (l.clone(), Json::Null)).collect(),
|
||||
_ => unreachable!("validated"),
|
||||
};
|
||||
let options = pairs
|
||||
.iter()
|
||||
.map(|(k, v)| {
|
||||
if is_blank(v) {
|
||||
k.py_str()
|
||||
} else {
|
||||
format!("{}: {}", k.py_str(), render_criterion(v))
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
(options, pairs.into_iter().map(|(k, _)| k).collect())
|
||||
}
|
||||
QType::Score => {
|
||||
let Some(Json::Array(levels)) = crit else {
|
||||
unreachable!("validated")
|
||||
};
|
||||
let options = levels
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| format!("level {i}: {}", render_criterion(c)))
|
||||
.collect();
|
||||
(options, levels.clone())
|
||||
}
|
||||
QType::Noul => {
|
||||
// Keys are matched case-insensitively; a repeat after
|
||||
// lower-casing keeps the later value, as a dict would.
|
||||
let mut desc: HashMap<String, Json> = HashMap::new();
|
||||
if let Some(Json::Object(m)) = crit {
|
||||
for (k, v) in m {
|
||||
desc.insert(k.to_lowercase(), v.clone());
|
||||
}
|
||||
}
|
||||
let (fl, tl) =
|
||||
resolve_noul_labels(def.get("labels").unwrap_or(&Json::Null)).expect("validated");
|
||||
let render = |key: &str, default: &str| match desc.get(key) {
|
||||
Some(v) if !is_blank(v) => render_criterion(v),
|
||||
_ => default.to_string(),
|
||||
};
|
||||
(
|
||||
vec![
|
||||
format!("{fl}: {}", render("false", DEFAULT_NOUL_FALSE)),
|
||||
format!("{tl}: {}", render("true", DEFAULT_NOUL_TRUE)),
|
||||
],
|
||||
Vec::new(),
|
||||
)
|
||||
}
|
||||
};
|
||||
Question {
|
||||
id: id.to_string(),
|
||||
kind,
|
||||
instructions,
|
||||
options,
|
||||
labels,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn def(s: &str) -> Json {
|
||||
Json::parse(s.as_bytes()).unwrap()
|
||||
}
|
||||
|
||||
fn err(s: &str) -> String {
|
||||
Question::from_definition("q", &def(s)).unwrap_err().detail
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn options_render_as_the_reference_renders_them() {
|
||||
let c = Question::from_definition(
|
||||
"q",
|
||||
&def(r#"{"type": "choice", "instructions": {"a": 1},
|
||||
"criteria": {"x": "", "y": null, "z": 0, "w": {"d": [1, 2.0]}, "v": "text"}}"#),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(c.instructions, r#"{"a": 1}"#);
|
||||
assert_eq!(
|
||||
c.options,
|
||||
["x", "y", "z: 0", r#"w: {"d": [1, 2.0]}"#, "v: text"]
|
||||
);
|
||||
|
||||
let l = Question::from_definition(
|
||||
"q",
|
||||
&def(r#"{"type": "choice", "instructions": "i", "criteria": [1, 2.5, false, "s"]}"#),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(l.options, ["1", "2.5", "False", "s"]);
|
||||
|
||||
let s = Question::from_definition(
|
||||
"q",
|
||||
&def(r#"{"type": "score", "instructions": "i", "criteria": ["calm", {"k": false}]}"#),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(s.options, ["level 0: calm", r#"level 1: {"k": false}"#]);
|
||||
|
||||
let n = Question::from_definition(
|
||||
"q",
|
||||
&def(
|
||||
r#"{"type": "noul", "instructions": "i", "criteria": {"TRUE": "held", "false": ""},
|
||||
"labels": {"false": " no ", "true": "yes"}}"#,
|
||||
),
|
||||
)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
n.options,
|
||||
["no: no, the statement does not hold", "yes: held"]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_questions_name_the_question_and_the_fix() {
|
||||
assert_eq!(
|
||||
err("[]"),
|
||||
"question 'q': definition must be a dict, got list"
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "pick"}"#),
|
||||
"question 'q': unknown type 'pick'; use one of ['choice', 'noul', 'score']"
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "noul"}"#),
|
||||
"question 'q': no 'instructions'; add the text the model should answer"
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "choice", "instructions": "i", "criteria": []}"#)
|
||||
.ends_with("needs at least one criterion")
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "choice", "instructions": "i", "criteria": [[1]]}"#)
|
||||
.starts_with("question 'q': choice label 0 is a list;")
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "choice", "instructions": "i", "criteria": ["a", null]}"#)
|
||||
.starts_with("question 'q': choice label 1 is null;")
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "choice", "instructions": "i", "criteria": [1, "x", true]}"#)
|
||||
.starts_with("question 'q': choice label 2 (True) repeats label 0;")
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "score", "instructions": "i", "criteria": ["a", null]}"#),
|
||||
"question 'q': score level 1 is null; give every level a description, index 0 first"
|
||||
);
|
||||
assert!(
|
||||
err(r#"{"type": "noul", "instructions": "i", "criteria": {"yes": "x"}}"#)
|
||||
.contains("got ['yes'].")
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "score", "instructions": "i", "criteria": ["a"], "labels": {}}"#),
|
||||
"question 'q': 'labels' is only supported for noul questions"
|
||||
);
|
||||
assert_eq!(
|
||||
err(r#"{"type": "noul", "instructions": "i", "labels": {"false": "x", "true": " x"}}"#),
|
||||
"question 'q': noul labels must map exactly 'false' and 'true' to distinct non-empty strings"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn nan_labels_never_collide() {
|
||||
assert!(
|
||||
Question::from_definition(
|
||||
"q",
|
||||
&def(r#"{"type": "choice", "instructions": "i", "criteria": [NaN, NaN]}"#),
|
||||
)
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,301 @@
|
||||
//! Choosing which checkpoint of a Laya family answers a request (#339).
|
||||
//!
|
||||
//! A released Laya repo carries three checkpoints: the English one at
|
||||
//! its root, `multilingual/` and `typed-decisions/`. The English
|
||||
//! checkpoint cannot read other scripts — upstream measured Khmer at
|
||||
//! 0.000 accuracy with 0.952 confidence — so an unpinned request is
|
||||
//! routed by the language of its state. This follows the reference's
|
||||
//! `Router._route` (`laya/router.py`) and its HTTP server's
|
||||
//! `_resolve_model` (`laya/serve.py`), including the `reason` strings.
|
||||
//!
|
||||
//! Precedence, as upstream: a checkpoint the caller names in `model`,
|
||||
//! then the detected script and language, then the default (English).
|
||||
//! Upstream's opt-in typed-decisions workflow routing and caller
|
||||
//! language hints are not exposed over HTTP by the reference, so they
|
||||
//! are not here either; the matched workflow is still reported.
|
||||
|
||||
use cortex_core::decisions::{DecisionRouting, Json};
|
||||
|
||||
use crate::harness::lang_detect;
|
||||
|
||||
/// The root checkpoint's routing name.
|
||||
pub const ENGLISH: &str = "english";
|
||||
pub const MULTILINGUAL: &str = "multilingual";
|
||||
pub const TYPED_DECISIONS: &str = "typed-decisions";
|
||||
|
||||
/// Where an unpinned request with no usable language signal goes.
|
||||
pub const DEFAULT_CHECKPOINT: &str = ENGLISH;
|
||||
|
||||
/// Every checkpoint a released family carries, root first.
|
||||
pub const FAMILY: [&str; 3] = [ENGLISH, MULTILINGUAL, TYPED_DECISIONS];
|
||||
|
||||
/// The reference's `_ALIASES`: names people are likely to type.
|
||||
const ALIASES: [(&str, &str); 10] = [
|
||||
("en", ENGLISH),
|
||||
("laya", ENGLISH),
|
||||
("default", ENGLISH),
|
||||
("multi", MULTILINGUAL),
|
||||
("ml", MULTILINGUAL),
|
||||
("laya-multilingual", MULTILINGUAL),
|
||||
("typed", TYPED_DECISIONS),
|
||||
("typed_decisions", TYPED_DECISIONS),
|
||||
("laya-typed-decisions", TYPED_DECISIONS),
|
||||
("decisions", TYPED_DECISIONS),
|
||||
];
|
||||
|
||||
/// The reference server's `_PUBLISHED_MODEL_IDS`: the Hugging Face ids
|
||||
/// the non-root checkpoints are published under.
|
||||
const PUBLISHED: [(&str, &str); 2] = [
|
||||
("convaiinnovations/laya-multilingual", MULTILINGUAL),
|
||||
("convaiinnovations/laya-typed-decisions", TYPED_DECISIONS),
|
||||
];
|
||||
|
||||
/// Question-id signatures of the four typed-decisions workflows.
|
||||
const WORKFLOWS: [(&str, [&str; 5]); 4] = [
|
||||
(
|
||||
"agent_trace_observability",
|
||||
["action", "needs_review", "outcome", "risk", "urgency"],
|
||||
),
|
||||
(
|
||||
"customer_service",
|
||||
["action", "category", "churn_risk", "needs_human", "urgency"],
|
||||
),
|
||||
(
|
||||
"invoice_processing",
|
||||
[
|
||||
"discrepancy_severity",
|
||||
"disposition",
|
||||
"duplicate",
|
||||
"matches_order",
|
||||
"urgency",
|
||||
],
|
||||
),
|
||||
(
|
||||
"security_incidents",
|
||||
[
|
||||
"credential_compromise",
|
||||
"disposition",
|
||||
"severity",
|
||||
"true_positive",
|
||||
"urgency",
|
||||
],
|
||||
),
|
||||
];
|
||||
|
||||
/// The checkpoint a request's `model` field pins, if it names one.
|
||||
///
|
||||
/// Anything else — a Jev service name like `jev-1`, the family's own
|
||||
/// repo id, an empty string — pins nothing and the request is routed.
|
||||
pub fn resolve_pin(model: Option<&str>) -> Option<&'static str> {
|
||||
let model = model?;
|
||||
if model.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let key = py_strip(model).to_lowercase();
|
||||
if let Some((_, ckpt)) = PUBLISHED.iter().find(|(id, _)| *id == key) {
|
||||
return Some(ckpt);
|
||||
}
|
||||
let key = ALIASES
|
||||
.iter()
|
||||
.find(|(alias, _)| *alias == key)
|
||||
.map(|(_, ckpt)| *ckpt)
|
||||
.unwrap_or(key.as_str());
|
||||
FAMILY.iter().copied().find(|c| *c == key)
|
||||
}
|
||||
|
||||
/// The typed-decisions workflow whose question ids these are exactly.
|
||||
pub fn match_workflow(questions: &[(String, Json)]) -> Option<&'static str> {
|
||||
let ids: std::collections::BTreeSet<&str> = questions.iter().map(|(q, _)| q.as_str()).collect();
|
||||
WORKFLOWS.iter().find_map(|(name, sig)| {
|
||||
let sig: std::collections::BTreeSet<&str> = sig.iter().copied().collect();
|
||||
(ids == sig).then_some(*name)
|
||||
})
|
||||
}
|
||||
|
||||
/// Decide which checkpoint answers, and say why.
|
||||
///
|
||||
/// `model_id` is the served family's repo id, used for the `repo` field.
|
||||
/// The result names a checkpoint whether or not this node has it
|
||||
/// loaded; the caller decides what to do about one it lacks.
|
||||
pub fn route(
|
||||
model_id: &str,
|
||||
state: &Json,
|
||||
questions: &[(String, Json)],
|
||||
pin: Option<&'static str>,
|
||||
) -> DecisionRouting {
|
||||
if let Some(key) = pin {
|
||||
return DecisionRouting {
|
||||
model: key.to_string(),
|
||||
repo: repo_of(model_id, key),
|
||||
reason: format!("explicit model={}", py_repr(key)),
|
||||
detection: None,
|
||||
workflow: None,
|
||||
};
|
||||
}
|
||||
let workflow = match_workflow(questions).map(str::to_string);
|
||||
let det = lang_detect::analyse(state);
|
||||
let (key, reason) = if det.script == "unknown" {
|
||||
(
|
||||
DEFAULT_CHECKPOINT,
|
||||
format!("no letters detected in state; using default ({DEFAULT_CHECKPOINT})"),
|
||||
)
|
||||
} else if det.script != "latin" {
|
||||
(
|
||||
MULTILINGUAL,
|
||||
format!(
|
||||
"non-Latin script ({}, {:.0}% of letters); the English checkpoint cannot read it",
|
||||
det.script,
|
||||
100.0 * det.non_latin_fraction
|
||||
),
|
||||
)
|
||||
} else if !det.is_english {
|
||||
let reason = if let Some(segment) = &det.mixed_segment {
|
||||
format!(
|
||||
"Latin script, mostly English, but a line or field reads as {} ({}); \
|
||||
the English checkpoint cannot read it",
|
||||
det.language.map(py_repr).unwrap_or_else(|| "None".into()),
|
||||
py_repr(&segment.chars().take(60).collect::<String>()),
|
||||
)
|
||||
} else if let Some(lang) = det.language {
|
||||
format!(
|
||||
"Latin script but language looks like {}, not English",
|
||||
py_repr(lang)
|
||||
)
|
||||
} else {
|
||||
format!(
|
||||
"Latin script, language not identified but {:.0}% non-English letters; \
|
||||
not safe for the English checkpoint",
|
||||
100.0 * det.diacritic_rate
|
||||
)
|
||||
};
|
||||
(MULTILINGUAL, reason)
|
||||
} else if det.language_undecided {
|
||||
(
|
||||
DEFAULT_CHECKPOINT,
|
||||
format!(
|
||||
"Latin script, language not identified and no non-English letters; \
|
||||
using default ({DEFAULT_CHECKPOINT})"
|
||||
),
|
||||
)
|
||||
} else {
|
||||
(ENGLISH, "English Latin text".to_string())
|
||||
};
|
||||
DecisionRouting {
|
||||
model: key.to_string(),
|
||||
repo: repo_of(model_id, key),
|
||||
reason,
|
||||
detection: Some(analysis_json(&det)),
|
||||
workflow,
|
||||
}
|
||||
}
|
||||
|
||||
/// The language per-language calibration keys on: the detected one.
|
||||
pub fn detected_language(routing: &DecisionRouting) -> Option<&str> {
|
||||
routing
|
||||
.detection
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("language"))
|
||||
.and_then(Json::as_str)
|
||||
}
|
||||
|
||||
/// `<repo>` for the root checkpoint, `<repo>/<subfolder>` otherwise.
|
||||
pub fn repo_of(model_id: &str, checkpoint: &str) -> String {
|
||||
if checkpoint == ENGLISH {
|
||||
model_id.to_string()
|
||||
} else {
|
||||
format!("{model_id}/{checkpoint}")
|
||||
}
|
||||
}
|
||||
|
||||
/// The analysis as an order-preserving [`Json`], so the response keeps
|
||||
/// the reference's field order.
|
||||
fn analysis_json(det: &lang_detect::Analysis) -> Json {
|
||||
let text = serde_json::to_string(det).unwrap_or_else(|_| "null".into());
|
||||
Json::parse(text.as_bytes()).unwrap_or(Json::Null)
|
||||
}
|
||||
|
||||
fn py_repr(s: &str) -> String {
|
||||
Json::String(s.to_string()).py_repr()
|
||||
}
|
||||
|
||||
/// Python's `str.strip()`: Unicode whitespace, which includes the
|
||||
/// U+001C–U+001F separators Rust's `trim` leaves alone.
|
||||
fn py_strip(s: &str) -> &str {
|
||||
s.trim_matches(|c: char| c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn pins_follow_the_reference_server() {
|
||||
assert_eq!(resolve_pin(Some("multilingual")), Some(MULTILINGUAL));
|
||||
assert_eq!(resolve_pin(Some(" ML ")), Some(MULTILINGUAL));
|
||||
assert_eq!(
|
||||
resolve_pin(Some("convaiinnovations/laya-typed-decisions")),
|
||||
Some(TYPED_DECISIONS)
|
||||
);
|
||||
assert_eq!(resolve_pin(Some("laya")), Some(ENGLISH));
|
||||
// A Jev service name, the family's own id, and nothing: no pin.
|
||||
assert_eq!(resolve_pin(Some("jev-1")), None);
|
||||
assert_eq!(resolve_pin(Some("convaiinnovations/laya")), None);
|
||||
assert_eq!(resolve_pin(Some("")), None);
|
||||
assert_eq!(resolve_pin(None), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn workflow_needs_the_exact_id_set() {
|
||||
let qs = |ids: &[&str]| -> Vec<(String, Json)> {
|
||||
ids.iter().map(|i| (i.to_string(), Json::Null)).collect()
|
||||
};
|
||||
assert_eq!(
|
||||
match_workflow(&qs(&[
|
||||
"urgency",
|
||||
"action",
|
||||
"category",
|
||||
"churn_risk",
|
||||
"needs_human"
|
||||
])),
|
||||
Some("customer_service")
|
||||
);
|
||||
assert_eq!(match_workflow(&qs(&["urgency"])), None);
|
||||
}
|
||||
|
||||
/// Every case in the reference fixtures routes to the checkpoint, with
|
||||
/// the reason and detection, the reference reported.
|
||||
#[test]
|
||||
fn routes_every_reference_case_like_the_reference() {
|
||||
let text = std::fs::read_to_string(
|
||||
std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
|
||||
.join("src/harness/testdata/laya/reference.json"),
|
||||
)
|
||||
.expect("read reference.json");
|
||||
let reference = Json::parse(text.as_bytes()).expect("parse reference.json");
|
||||
let Some(Json::Array(cases)) = reference.get("cases") else {
|
||||
panic!("no cases");
|
||||
};
|
||||
for case in cases {
|
||||
let name = case.get("name").and_then(Json::as_str).unwrap();
|
||||
let request = case.get("request").unwrap();
|
||||
let body = serde_json::to_vec(request).unwrap();
|
||||
let req = cortex_core::decisions::SystemOneRequest::parse(
|
||||
&body,
|
||||
cortex_core::decisions::limits::DEFAULT_MAX_TOKEN_BUDGET,
|
||||
)
|
||||
.unwrap();
|
||||
let routing = route(
|
||||
"convaiinnovations/laya",
|
||||
&req.state,
|
||||
&req.questions,
|
||||
resolve_pin(req.model.as_deref()),
|
||||
);
|
||||
let want = case.get("response").unwrap().get("routing").unwrap();
|
||||
assert_eq!(
|
||||
serde_json::to_string(&routing).unwrap(),
|
||||
serde_json::to_string(want).unwrap(),
|
||||
"{name}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,316 @@
|
||||
//! Sequence assembly: one token sequence per question, under the model's
|
||||
//! token budgets.
|
||||
//!
|
||||
//! ```text
|
||||
//! [CLS] "<type> question: <instructions>" [SEP] [MASK] opt0 [MASK] opt1 … [SEP] state [SEP]
|
||||
//! ```
|
||||
//!
|
||||
//! The budget rules, in the order they apply:
|
||||
//!
|
||||
//! 1. Each option is `[MASK]` plus at most 48 tokens of its text.
|
||||
//! 2. `head_max_len` bounds the instruction text and all options together.
|
||||
//! If fewer than 16 tokens would remain for the instructions, every
|
||||
//! option is cut to the same `max(4, (head_max_len - 16) / options)`
|
||||
//! tokens (marker included).
|
||||
//! 3. The instructions get what is left of `head_max_len`, but never fewer
|
||||
//! than 8 tokens.
|
||||
//! 4. The state fills the rest of `max_len`. A conversation (a JSON list)
|
||||
//! keeps its newest turns, so it is cut from the left; anything else is
|
||||
//! cut from the right.
|
||||
//! 5. The whole sequence is cut to `max_len`, dropping any marker that no
|
||||
//! longer fits — a question whose markers do not all fit is refused.
|
||||
|
||||
use super::question::{QType, Question, questions_from_request};
|
||||
use super::tokenizer::TextEncoder;
|
||||
use super::{DecisionConfig, DecisionRequestError, unprocessable};
|
||||
use cortex_core::decisions::Json;
|
||||
|
||||
/// Option text tokens kept before the head budget is considered.
|
||||
pub const MAX_OPTION_TOKENS: usize = 48;
|
||||
/// Head tokens the instructions are guaranteed before options are shrunk.
|
||||
const MIN_INSTRUCTION_ROOM: usize = 16;
|
||||
/// Per-option floor (marker included) when options are shrunk.
|
||||
const MIN_TOKENS_PER_OPTION: usize = 4;
|
||||
/// Instruction tokens kept however many the options take.
|
||||
const MIN_INSTRUCTION_TOKENS: usize = 8;
|
||||
|
||||
/// Per-request overrides of the checkpoint's token budgets.
|
||||
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
|
||||
pub struct Budget {
|
||||
pub max_len: Option<usize>,
|
||||
pub head_max_len: Option<usize>,
|
||||
}
|
||||
|
||||
/// One model input row: one question over the state.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct Row {
|
||||
pub input_ids: Vec<u32>,
|
||||
/// Position of each option's `[MASK]` marker, in option order.
|
||||
pub markers: Vec<usize>,
|
||||
pub qtype: QType,
|
||||
}
|
||||
|
||||
/// What the head budget did to a question's options.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct OptionStats {
|
||||
pub options: usize,
|
||||
/// Options whose token span is still unique after the cut. Fewer than
|
||||
/// `options` means the model can no longer tell some of them apart.
|
||||
pub options_distinct: usize,
|
||||
/// The per-option cap applied, when the options had to be shrunk.
|
||||
pub tokens_per_option: Option<usize>,
|
||||
}
|
||||
|
||||
/// A request turned into model input: one [`Row`] per question, in request
|
||||
/// order. With no questions there are no rows and nothing to run.
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct Prepared {
|
||||
pub questions: Vec<Question>,
|
||||
pub rows: Vec<Row>,
|
||||
pub stats: Vec<OptionStats>,
|
||||
/// Every token the forward pass encodes: the metered input.
|
||||
pub input_tokens: usize,
|
||||
/// The id batched rows are padded with.
|
||||
pub pad_id: u32,
|
||||
}
|
||||
|
||||
/// The text a state is tokenized from: a string as-is, anything else as
|
||||
/// `json.dumps(state, ensure_ascii=False)`.
|
||||
pub fn serialize_state(state: &Json) -> String {
|
||||
match state {
|
||||
Json::String(s) => s.clone(),
|
||||
other => other.dumps(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Validate a request's questions and build their model input.
|
||||
///
|
||||
/// `questions` are the request's definitions in request order; `cfg` and
|
||||
/// `tok` belong to the checkpoint the request was routed to.
|
||||
pub fn prepare(
|
||||
state: &Json,
|
||||
questions: &[(String, Json)],
|
||||
cfg: &DecisionConfig,
|
||||
tok: &dyn TextEncoder,
|
||||
budget: Budget,
|
||||
) -> Result<Prepared, DecisionRequestError> {
|
||||
let special = tok.special();
|
||||
if questions.is_empty() {
|
||||
return Ok(Prepared {
|
||||
questions: Vec::new(),
|
||||
rows: Vec::new(),
|
||||
stats: Vec::new(),
|
||||
input_tokens: 0,
|
||||
pad_id: special.pad_id,
|
||||
});
|
||||
}
|
||||
let questions = questions_from_request(questions)?;
|
||||
let max_len = budget.max_len.unwrap_or(cfg.max_len);
|
||||
let head_max_len = budget.head_max_len.unwrap_or(cfg.head_max_len);
|
||||
let encode = |text: &str| {
|
||||
tok.encode(text)
|
||||
.map_err(|e| DecisionRequestError::new(500, format!("tokenizer failed: {e}")))
|
||||
};
|
||||
// Tokenized once and sliced per question.
|
||||
let state_ids = encode(&serialize_state(state).replace(&special.mask_token, " "))?;
|
||||
let truncate_left = matches!(state, Json::Array(_));
|
||||
let frame = Frame {
|
||||
cls: special.cls_id,
|
||||
sep: special.sep_id,
|
||||
max_len,
|
||||
head_max_len,
|
||||
};
|
||||
|
||||
let mut rows = Vec::with_capacity(questions.len());
|
||||
let mut stats = Vec::with_capacity(questions.len());
|
||||
for question in &questions {
|
||||
let instructions = question.instructions.replace(&special.mask_token, " ");
|
||||
let head_ids = encode(&format!(
|
||||
"{} question: {instructions}",
|
||||
question.kind.name()
|
||||
))?;
|
||||
let mut option_ids = Vec::with_capacity(question.options.len());
|
||||
for option in &question.options {
|
||||
let mut ids = vec![special.mask_id];
|
||||
let text = encode(&format!(" {}", option.replace(&special.mask_token, " ")))?;
|
||||
ids.extend(text.into_iter().take(MAX_OPTION_TOKENS));
|
||||
option_ids.push(ids);
|
||||
}
|
||||
let (row, s) = frame.assemble(head_ids, option_ids, &state_ids, truncate_left);
|
||||
if row.markers.len() != question.options.len() {
|
||||
return Err(unprocessable(format!(
|
||||
"question {} options exceed head_max_len={head_max_len}",
|
||||
Json::String(question.id.clone()).py_repr()
|
||||
)));
|
||||
}
|
||||
rows.push(Row {
|
||||
qtype: question.kind,
|
||||
..row
|
||||
});
|
||||
stats.push(s);
|
||||
}
|
||||
let input_tokens = rows.iter().map(|r| r.input_ids.len()).sum();
|
||||
Ok(Prepared {
|
||||
questions,
|
||||
rows,
|
||||
stats,
|
||||
input_tokens,
|
||||
pad_id: special.pad_id,
|
||||
})
|
||||
}
|
||||
|
||||
/// The fixed parts of every sequence for one request.
|
||||
struct Frame {
|
||||
cls: u32,
|
||||
sep: u32,
|
||||
max_len: usize,
|
||||
head_max_len: usize,
|
||||
}
|
||||
|
||||
impl Frame {
|
||||
/// The pure budget arithmetic of one sequence, over already-tokenized parts.
|
||||
fn assemble(
|
||||
&self,
|
||||
mut head_ids: Vec<u32>,
|
||||
mut option_ids: Vec<Vec<u32>>,
|
||||
state_ids: &[u32],
|
||||
truncate_left: bool,
|
||||
) -> (Row, OptionStats) {
|
||||
let Frame {
|
||||
cls,
|
||||
sep,
|
||||
max_len,
|
||||
head_max_len,
|
||||
} = *self;
|
||||
let used = |o: &[Vec<u32>]| o.iter().map(Vec::len).sum::<usize>();
|
||||
// Signed: options can overrun the head budget.
|
||||
let mut opt_budget = head_max_len as isize - used(&option_ids) as isize;
|
||||
let mut tokens_per_option = None;
|
||||
if opt_budget < MIN_INSTRUCTION_ROOM as isize {
|
||||
let per = MIN_TOKENS_PER_OPTION
|
||||
.max(head_max_len.saturating_sub(MIN_INSTRUCTION_ROOM) / option_ids.len().max(1));
|
||||
tokens_per_option = Some(per);
|
||||
for o in &mut option_ids {
|
||||
o.truncate(per);
|
||||
}
|
||||
opt_budget = head_max_len as isize - used(&option_ids) as isize;
|
||||
}
|
||||
head_ids.truncate(opt_budget.max(MIN_INSTRUCTION_TOKENS as isize) as usize);
|
||||
|
||||
let mut ids = Vec::with_capacity(max_len);
|
||||
ids.push(cls);
|
||||
ids.extend(head_ids);
|
||||
ids.push(sep);
|
||||
let mut markers = Vec::with_capacity(option_ids.len());
|
||||
for o in &option_ids {
|
||||
markers.push(ids.len());
|
||||
ids.extend_from_slice(o);
|
||||
}
|
||||
ids.push(sep);
|
||||
let room = max_len.saturating_sub(ids.len() + 1);
|
||||
let state = if truncate_left {
|
||||
&state_ids[state_ids.len().saturating_sub(room)..]
|
||||
} else {
|
||||
&state_ids[..room.min(state_ids.len())]
|
||||
};
|
||||
ids.extend_from_slice(state);
|
||||
ids.push(sep);
|
||||
ids.truncate(max_len);
|
||||
markers.retain(|&m| m < max_len);
|
||||
|
||||
let mut distinct = option_ids.clone();
|
||||
distinct.sort();
|
||||
distinct.dedup();
|
||||
(
|
||||
Row {
|
||||
input_ids: ids,
|
||||
markers,
|
||||
qtype: QType::Choice,
|
||||
},
|
||||
OptionStats {
|
||||
options: option_ids.len(),
|
||||
options_distinct: distinct.len(),
|
||||
tokens_per_option,
|
||||
},
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
const CLS: u32 = 1;
|
||||
const SEP: u32 = 2;
|
||||
const M: u32 = 9;
|
||||
|
||||
fn frame(max_len: usize, head_max_len: usize) -> Frame {
|
||||
Frame {
|
||||
cls: CLS,
|
||||
sep: SEP,
|
||||
max_len,
|
||||
head_max_len,
|
||||
}
|
||||
}
|
||||
|
||||
fn opts(lens: &[usize]) -> Vec<Vec<u32>> {
|
||||
lens.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &n)| {
|
||||
let mut o = vec![M];
|
||||
o.extend((0..n - 1).map(|j| 100 + (i * 50 + j) as u32));
|
||||
o
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn roomy_budget_keeps_everything() {
|
||||
let (row, s) = frame(64, 32).assemble(vec![10, 11], opts(&[2, 3]), &[50, 51], false);
|
||||
assert_eq!(
|
||||
row.input_ids,
|
||||
[CLS, 10, 11, SEP, M, 100, M, 150, 151, SEP, 50, 51, SEP]
|
||||
);
|
||||
assert_eq!(row.markers, [4, 6]);
|
||||
assert_eq!(
|
||||
(s.options, s.options_distinct, s.tokens_per_option),
|
||||
(2, 2, None)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn crowded_head_shrinks_options_evenly_and_keeps_eight_instruction_tokens() {
|
||||
// 3 options of 10 = 30 tokens against a 40-token head: 10 left < 16.
|
||||
let head: Vec<u32> = (10..40).collect();
|
||||
let (row, s) = frame(128, 40).assemble(head, opts(&[10, 10, 10]), &[], false);
|
||||
// per = max(4, (40 - 16) / 3) = 8; 40 - 24 = 16 instruction tokens.
|
||||
assert_eq!(s.tokens_per_option, Some(8));
|
||||
assert_eq!(row.markers, [18, 26, 34]);
|
||||
let (row, s) = frame(256, 40).assemble((10..40).collect(), opts(&[10; 12]), &[], false);
|
||||
// per = max(4, 24 / 12 = 2) = 4; the budget goes negative, 8 tokens survive.
|
||||
assert_eq!(s.tokens_per_option, Some(4));
|
||||
assert_eq!(row.markers[0], 1 + 8 + 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn state_is_cut_from_the_side_the_state_shape_asks_for() {
|
||||
let state: Vec<u32> = (50..60).collect();
|
||||
let (right, _) = frame(8, 32).assemble(vec![], opts(&[1]), &state, false);
|
||||
let (left, _) = frame(8, 32).assemble(vec![], opts(&[1]), &state, true);
|
||||
// [CLS] [SEP] [MASK] [SEP] = 4 tokens, 1 reserved for the last [SEP]: room 3.
|
||||
assert_eq!(right.input_ids, [CLS, SEP, M, SEP, 50, 51, 52, SEP]);
|
||||
assert_eq!(left.input_ids, [CLS, SEP, M, SEP, 57, 58, 59, SEP]);
|
||||
// No room at all: none of the state, not all of it.
|
||||
let (none, _) = frame(5, 32).assemble(vec![], opts(&[1]), &state, true);
|
||||
assert_eq!(none.input_ids, [CLS, SEP, M, SEP, SEP]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn markers_past_max_len_are_dropped_and_collisions_counted() {
|
||||
let same = vec![vec![M, 7, 7], vec![M, 7, 7]];
|
||||
let (row, s) = frame(5, 64).assemble(vec![], same, &[], false);
|
||||
// [CLS] [SEP] then markers at 2 and 5; only the first fits in 5.
|
||||
assert_eq!(row.markers, [2]);
|
||||
assert_eq!((s.options, s.options_distinct), (2, 1));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,318 @@
|
||||
//! Per-option logits → calibrated, typed answers.
|
||||
//!
|
||||
//! The reference computes this in numpy on the float32 logits the model
|
||||
//! returns, and publishes each probability rounded to 4 decimal places.
|
||||
//! Reproducing the rounded values exactly means reproducing the float32
|
||||
//! arithmetic, in numpy's order: division by a float32 temperature, a
|
||||
//! max-subtracted `exp`, and numpy's *pairwise* summation (not a running
|
||||
//! sum) for every reduction. The expected score is the exception: numpy
|
||||
//! promotes `arange(k) * p` to float64.
|
||||
|
||||
use super::question::QType;
|
||||
use super::sequence::Prepared;
|
||||
use super::{DECISION_MODEL_NAME, DecisionConfig};
|
||||
use cortex_core::decisions::{
|
||||
CollapsedOptions, DecisionAction, DecisionAnswer, DecisionUsage, OrderedMap, SystemOneResponse,
|
||||
};
|
||||
|
||||
/// The model's output for one [`super::Row`].
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct RowOutput {
|
||||
/// The scorer's raw logit at each option marker, before calibration.
|
||||
/// Entries past the row's marker count (batch padding) are ignored.
|
||||
pub logits: Vec<f32>,
|
||||
/// The act head's probability of acting: `softmax(act_logits)[0]`.
|
||||
pub act_probability: f32,
|
||||
}
|
||||
|
||||
/// Build the answers and usage for a prepared request from the model's
|
||||
/// per-row outputs (`outputs[i]` belongs to `prepared.rows[i]`).
|
||||
///
|
||||
/// `lang` selects per-language calibration where the checkpoint has it:
|
||||
/// the language the request was declared or detected as.
|
||||
pub fn shape(
|
||||
prepared: &Prepared,
|
||||
outputs: &[RowOutput],
|
||||
cfg: &DecisionConfig,
|
||||
lang: Option<&str>,
|
||||
) -> (OrderedMap<DecisionAnswer>, DecisionUsage) {
|
||||
assert_eq!(
|
||||
outputs.len(),
|
||||
prepared.rows.len(),
|
||||
"one model output per prepared row"
|
||||
);
|
||||
let temps = cfg.temperatures_for(lang);
|
||||
let mut answers = Vec::with_capacity(outputs.len());
|
||||
let mut collapsed = Vec::new();
|
||||
for ((question, row), (out, stats)) in prepared
|
||||
.questions
|
||||
.iter()
|
||||
.zip(&prepared.rows)
|
||||
.zip(outputs.iter().zip(&prepared.stats))
|
||||
{
|
||||
let k = row.markers.len();
|
||||
let p = calibrate(
|
||||
&out.logits[..k],
|
||||
temps.for_question(question.kind as usize, k),
|
||||
);
|
||||
let answer_confidence = round4(
|
||||
p.iter()
|
||||
.copied()
|
||||
.fold(f32::NEG_INFINITY, f32::max)
|
||||
.clamp(0.0, 1.0) as f64,
|
||||
);
|
||||
let action = DecisionAction {
|
||||
act_probability: round4(out.act_probability as f64),
|
||||
};
|
||||
let probabilities = |key: &dyn Fn(usize) -> String| {
|
||||
OrderedMap(
|
||||
p.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| (key(i), round4(v as f64)))
|
||||
.collect(),
|
||||
)
|
||||
};
|
||||
let answer = match question.kind {
|
||||
QType::Choice => DecisionAnswer::Choice {
|
||||
choice: question.labels[argmax(&p)].clone(),
|
||||
probabilities: probabilities(&|i| {
|
||||
question.labels[i]
|
||||
.json_key()
|
||||
.expect("choice labels are scalars")
|
||||
}),
|
||||
confidence: round4(entropy_confidence(&p) as f64),
|
||||
answer_confidence,
|
||||
action,
|
||||
},
|
||||
QType::Score => {
|
||||
// int64 * float32 promotes to float64 in numpy.
|
||||
let expected: Vec<f64> = p
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, &v)| i as f64 * v as f64)
|
||||
.collect();
|
||||
DecisionAnswer::Score {
|
||||
score: round4(pairwise_sum(&expected)),
|
||||
legend: OrderedMap(
|
||||
question
|
||||
.labels
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, c)| (i.to_string(), c.clone()))
|
||||
.collect(),
|
||||
),
|
||||
probabilities: probabilities(&|i| i.to_string()),
|
||||
confidence: round4(entropy_confidence(&p) as f64),
|
||||
answer_confidence,
|
||||
action,
|
||||
}
|
||||
}
|
||||
QType::Noul => {
|
||||
let yes = p[1] as f64;
|
||||
DecisionAnswer::Noul {
|
||||
noul: round4(yes),
|
||||
confidence: round4(yes.max(1.0 - yes)),
|
||||
answer_confidence,
|
||||
action,
|
||||
}
|
||||
}
|
||||
};
|
||||
answers.push((question.id.clone(), answer));
|
||||
if stats.options_distinct < stats.options {
|
||||
collapsed.push((
|
||||
question.id.clone(),
|
||||
CollapsedOptions {
|
||||
total: stats.options,
|
||||
distinct: stats.options_distinct,
|
||||
tokens_per_option: stats.tokens_per_option,
|
||||
},
|
||||
));
|
||||
}
|
||||
}
|
||||
let usage = DecisionUsage {
|
||||
input_tokens: prepared.input_tokens,
|
||||
output_tokens: 0,
|
||||
options: (!collapsed.is_empty()).then_some(OrderedMap(collapsed)),
|
||||
};
|
||||
(OrderedMap(answers), usage)
|
||||
}
|
||||
|
||||
/// Assemble the full response once routing is known.
|
||||
pub fn response(
|
||||
answers: OrderedMap<DecisionAnswer>,
|
||||
usage: DecisionUsage,
|
||||
routing: Option<cortex_core::decisions::DecisionRouting>,
|
||||
) -> SystemOneResponse {
|
||||
SystemOneResponse {
|
||||
model: DECISION_MODEL_NAME.to_string(),
|
||||
answers,
|
||||
usage,
|
||||
routing,
|
||||
}
|
||||
}
|
||||
|
||||
/// Temperature-scaled softmax, as the reference evaluates it in numpy:
|
||||
/// `z = logits / t`, `e = exp(z - z.max())`, `p = e / e.sum()`.
|
||||
///
|
||||
/// numpy's SIMD `expf` and libm's disagree by one ulp on about a third of
|
||||
/// inputs. That is invisible after the 4-decimal rounding on every
|
||||
/// reference fixture, and far below the engine's own distance from the
|
||||
/// reference's logits, so it is not reproduced; the steps either side of
|
||||
/// the `exp` are.
|
||||
fn calibrate(logits: &[f32], t: f64) -> Vec<f32> {
|
||||
let z = scale(logits, t);
|
||||
let max = z.iter().copied().fold(f32::NEG_INFINITY, f32::max);
|
||||
let e: Vec<f32> = z.iter().map(|&v| (v - max).exp()).collect();
|
||||
normalize(&e)
|
||||
}
|
||||
|
||||
/// `logits / t` in float32: numpy casts the Python-float temperature to the
|
||||
/// array's dtype rather than promoting the array.
|
||||
fn scale(logits: &[f32], t: f64) -> Vec<f32> {
|
||||
let t = t as f32;
|
||||
logits.iter().map(|&l| l / t).collect()
|
||||
}
|
||||
|
||||
/// `e / e.sum()` with numpy's pairwise float32 sum.
|
||||
fn normalize(e: &[f32]) -> Vec<f32> {
|
||||
let sum = pairwise_sum(e);
|
||||
e.iter().map(|&v| v / sum).collect()
|
||||
}
|
||||
|
||||
/// `1 - H(p) / log(k)`, clipped to [0, 1]: how concentrated the whole
|
||||
/// distribution is. Not calibrated; see `answer_confidence` for that.
|
||||
fn entropy_confidence(p: &[f32]) -> f32 {
|
||||
let k = p.len();
|
||||
if k < 2 {
|
||||
return 1.0;
|
||||
}
|
||||
let terms: Vec<f32> = p.iter().map(|&v| v * v.clamp(1e-12, 1.0).ln()).collect();
|
||||
let ent = -pairwise_sum(&terms);
|
||||
(1.0 - ent / (k as f64).ln() as f32).clamp(0.0, 1.0)
|
||||
}
|
||||
|
||||
/// The first index of the largest value, as `argmax` picks it.
|
||||
fn argmax(p: &[f32]) -> usize {
|
||||
let mut best = 0;
|
||||
for (i, &v) in p.iter().enumerate() {
|
||||
if v > p[best] {
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best
|
||||
}
|
||||
|
||||
/// Python's `round(x, 4)`: correctly rounded, ties to even on the exact
|
||||
/// binary value. Rust's fixed-precision formatting rounds the same way.
|
||||
fn round4(x: f64) -> f64 {
|
||||
format!("{x:.4}").parse().expect("formatted float parses")
|
||||
}
|
||||
|
||||
/// numpy's pairwise summation over a contiguous array, which the rounded
|
||||
/// outputs depend on: a plain loop below 8 elements, eight interleaved
|
||||
/// accumulators up to 128, and recursive halving (at a multiple of 8)
|
||||
/// beyond.
|
||||
fn pairwise_sum<T>(a: &[T]) -> T
|
||||
where
|
||||
T: Copy + std::ops::Add<Output = T> + Default,
|
||||
{
|
||||
const BLOCK: usize = 128;
|
||||
let n = a.len();
|
||||
if n < 8 {
|
||||
let mut res = T::default();
|
||||
for &v in a {
|
||||
res = res + v;
|
||||
}
|
||||
res
|
||||
} else if n <= BLOCK {
|
||||
let mut r = [a[0], a[1], a[2], a[3], a[4], a[5], a[6], a[7]];
|
||||
let mut i = 8;
|
||||
while i < n - n % 8 {
|
||||
for j in 0..8 {
|
||||
r[j] = r[j] + a[i + j];
|
||||
}
|
||||
i += 8;
|
||||
}
|
||||
let mut res = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7]));
|
||||
while i < n {
|
||||
res = res + a[i];
|
||||
i += 1;
|
||||
}
|
||||
res
|
||||
} else {
|
||||
let mut n2 = n / 2;
|
||||
n2 -= n2 % 8;
|
||||
pairwise_sum(&a[..n2]) + pairwise_sum(&a[n2..])
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn pairwise_sum_differs_from_a_running_sum_where_numpy_does() {
|
||||
// 1e8 swamps every 1.0 added to it one at a time in f32; pairwise
|
||||
// grouping keeps them. `int(np.float32 array.sum())` of this array is 100000008.
|
||||
let mut a = vec![1.0f32; 16];
|
||||
a[0] = 1e8;
|
||||
let running = a.iter().fold(0.0f32, |s, &v| s + v);
|
||||
assert_eq!(running, 1e8);
|
||||
assert_eq!(pairwise_sum(&a), 100_000_008.0);
|
||||
}
|
||||
|
||||
/// Rows of `testdata/laya/numpy_softmax.txt`: the reference's
|
||||
/// calibration run in numpy on random float32 logits, with the bits of
|
||||
/// every intermediate. Bit equality pins the evaluation order — a
|
||||
/// running sum instead of numpy's pairwise one, or dividing by the
|
||||
/// temperature in float64, changes low bits that the 4-decimal rounding
|
||||
/// usually hides.
|
||||
fn numpy_rows() -> Vec<(f64, [Vec<f32>; 4])> {
|
||||
include_str!("../testdata/laya/numpy_softmax.txt")
|
||||
.lines()
|
||||
.filter(|l| !l.starts_with('#') && !l.trim().is_empty())
|
||||
.map(|l| {
|
||||
let mut f = l.split('|');
|
||||
let t: f64 = f.next().unwrap().parse().unwrap();
|
||||
let bits = |s: &str| -> Vec<f32> {
|
||||
s.split(',')
|
||||
.map(|h| f32::from_bits(u32::from_str_radix(&h[2..], 16).unwrap()))
|
||||
.collect()
|
||||
};
|
||||
let v: Vec<Vec<f32>> = f.map(bits).collect();
|
||||
(t, [v[0].clone(), v[1].clone(), v[2].clone(), v[3].clone()])
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn bits(v: &[f32]) -> Vec<u32> {
|
||||
v.iter().map(|x| x.to_bits()).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn calibration_matches_numpy_either_side_of_exp() {
|
||||
let rows = numpy_rows();
|
||||
assert_eq!(rows.len(), 4);
|
||||
for (t, [logits, z, e, p]) in rows {
|
||||
assert_eq!(
|
||||
bits(&scale(&logits, t)),
|
||||
bits(&z),
|
||||
"k={} t={t}: z",
|
||||
logits.len()
|
||||
);
|
||||
assert_eq!(
|
||||
bits(&normalize(&e)),
|
||||
bits(&p),
|
||||
"k={} t={t}: p",
|
||||
logits.len()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn round4_ties_to_even_on_the_binary_value() {
|
||||
assert_eq!(round4(0.03125), 0.0312);
|
||||
assert_eq!(round4(1.00005), 1.0001); // 1.00005 is just above the tie in binary
|
||||
assert_eq!(round4(0.99995), 1.0);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
//! The tokenizer a decision checkpoint ships, and its special tokens.
|
||||
|
||||
use std::path::Path;
|
||||
|
||||
/// The special tokens the sequence layout uses, resolved to ids.
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct SpecialTokens {
|
||||
/// The option marker's text (`[MASK]`, `<mask>` for mmBERT). Literal
|
||||
/// occurrences in caller text are replaced with a space before
|
||||
/// tokenizing, so no caller can forge a marker.
|
||||
pub mask_token: String,
|
||||
pub cls_id: u32,
|
||||
pub sep_id: u32,
|
||||
pub mask_id: u32,
|
||||
/// Padding for batched rows; not part of any sequence.
|
||||
pub pad_id: u32,
|
||||
}
|
||||
|
||||
/// Text → token ids, with no special tokens added.
|
||||
pub trait TextEncoder {
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>, String>;
|
||||
fn special(&self) -> &SpecialTokens;
|
||||
}
|
||||
|
||||
/// A checkpoint's `tokenizer/` directory: `tokenizer.json` plus the special
|
||||
/// token names in `tokenizer_config.json`.
|
||||
pub struct DecisionTokenizer {
|
||||
tok: tokenizers::Tokenizer,
|
||||
special: SpecialTokens,
|
||||
}
|
||||
|
||||
impl DecisionTokenizer {
|
||||
pub fn from_dir(dir: &Path) -> anyhow::Result<Self> {
|
||||
let mut tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
|
||||
.map_err(|e| anyhow::anyhow!("{}: {e}", dir.join("tokenizer.json").display()))?;
|
||||
// The Python tokenizer call the reference makes pads and truncates
|
||||
// nothing unless asked; a tokenizer.json that carries either setting
|
||||
// would otherwise apply it here.
|
||||
tok.with_truncation(None)
|
||||
.map_err(|e| anyhow::anyhow!("disabling truncation: {e}"))?;
|
||||
tok.with_padding(None);
|
||||
let cfg_path = dir.join("tokenizer_config.json");
|
||||
let cfg: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&cfg_path)?)
|
||||
.map_err(|e| anyhow::anyhow!("{}: {e}", cfg_path.display()))?;
|
||||
let name = |key: &str| -> anyhow::Result<String> {
|
||||
cfg.get(key)
|
||||
.and_then(|v| {
|
||||
v.as_str()
|
||||
.or_else(|| v.get("content").and_then(|c| c.as_str()))
|
||||
})
|
||||
.map(str::to_owned)
|
||||
.ok_or_else(|| anyhow::anyhow!("{}: no {key}", cfg_path.display()))
|
||||
};
|
||||
let id = |token: &str| -> anyhow::Result<u32> {
|
||||
tok.token_to_id(token)
|
||||
.ok_or_else(|| anyhow::anyhow!("special token {token:?} is not in the vocabulary"))
|
||||
};
|
||||
let mask_token = name("mask_token")?;
|
||||
let special = SpecialTokens {
|
||||
cls_id: id(&name("cls_token")?)?,
|
||||
sep_id: id(&name("sep_token")?)?,
|
||||
mask_id: id(&mask_token)?,
|
||||
pad_id: id(&name("pad_token")?)?,
|
||||
mask_token,
|
||||
};
|
||||
Ok(Self { tok, special })
|
||||
}
|
||||
}
|
||||
|
||||
impl TextEncoder for DecisionTokenizer {
|
||||
fn encode(&self, text: &str) -> Result<Vec<u32>, String> {
|
||||
self.tok
|
||||
.encode(text, false)
|
||||
.map(|e| e.get_ids().to_vec())
|
||||
.map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
fn special(&self) -> &SpecialTokens {
|
||||
&self.special
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
//! A loaded Laya decision model, as the rest of neuron sees it (#336).
|
||||
//!
|
||||
//! The weights live in a device worker's decision slab (CUDA) or in an
|
||||
//! `Arc` scored on the blocking pool (CPU); either way the caller holds
|
||||
//! only this struct and asks it to [`DecisionCheckpoint::decide`] a
|
||||
//! batch of already-assembled questions. Turning a `/v1/systemone`
|
||||
//! request into that batch — rendering options, budgets, tokenisation —
|
||||
//! and turning the scores back into calibrated answers happen outside
|
||||
//! this module.
|
||||
//!
|
||||
//! One Hugging Face repo can carry several checkpoints (the released
|
||||
//! `convaiinnovations/laya` has the English one at its root and
|
||||
//! `multilingual/` and `typed-decisions/` beside it), so a loaded model
|
||||
//! is a list of checkpoints, each named by its subfolder.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicBool;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use candle_core::DType;
|
||||
|
||||
use super::arch::laya::{CheckpointConfig, DecisionBatch, DecisionRow, LayaModel};
|
||||
use super::decision::{DecisionConfig, DecisionTokenizer};
|
||||
use super::device_worker::{DecisionHandle, DeviceWorkerHandle};
|
||||
use cortex_core::harness::ModelSpec;
|
||||
|
||||
/// Routing name of the checkpoint at the repo root. A released Laya
|
||||
/// family's root is its English checkpoint, and upstream routes to it
|
||||
/// under that name.
|
||||
pub const ROOT_CHECKPOINT: &str = super::decision::route::ENGLISH;
|
||||
|
||||
/// Where one checkpoint's weights live and how to reach them.
|
||||
enum Backend {
|
||||
/// Owned by a CUDA device worker; scored through its channel.
|
||||
Worker {
|
||||
worker: Arc<DeviceWorkerHandle>,
|
||||
handle: DecisionHandle,
|
||||
},
|
||||
/// CPU fallback. `LayaModel::forward` takes `&self`, so concurrent
|
||||
/// requests share it without a lock.
|
||||
Cpu(Arc<LayaModel>),
|
||||
}
|
||||
|
||||
/// The files one checkpoint needs, already on local disk.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct CheckpointFiles {
|
||||
/// The checkpoint directory: holds `rl_agent_config.json`,
|
||||
/// `encoder/config.json` and `model.safetensors`.
|
||||
pub dir: PathBuf,
|
||||
pub tokenizer: PathBuf,
|
||||
}
|
||||
|
||||
pub struct DecisionCheckpoint {
|
||||
/// The checkpoint's routing name: [`ROOT_CHECKPOINT`] for the repo
|
||||
/// root, else its subfolder (`multilingual`, `typed-decisions`).
|
||||
pub name: String,
|
||||
pub config: CheckpointConfig,
|
||||
/// Budgets and fitted temperatures from `rl_agent_config.json`, as
|
||||
/// the prompt builder and calibration read them.
|
||||
pub calibration: DecisionConfig,
|
||||
pub tokenizer: DecisionTokenizer,
|
||||
/// Precision the weights were loaded in.
|
||||
pub dtype: DType,
|
||||
backend: Backend,
|
||||
}
|
||||
|
||||
impl DecisionCheckpoint {
|
||||
/// Load one checkpoint onto `worker`'s device, or onto the CPU when
|
||||
/// `worker` is `None`.
|
||||
pub async fn load(
|
||||
name: String,
|
||||
files: CheckpointFiles,
|
||||
model_id: &str,
|
||||
worker: Option<Arc<DeviceWorkerHandle>>,
|
||||
) -> Result<Self> {
|
||||
let config = CheckpointConfig::from_dir(&files.dir)?;
|
||||
let calibration_path = files.dir.join("rl_agent_config.json");
|
||||
let calibration = DecisionConfig::from_json(
|
||||
&std::fs::read_to_string(&calibration_path)
|
||||
.with_context(|| format!("read {}", calibration_path.display()))?,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let tokenizer_dir = files
|
||||
.tokenizer
|
||||
.parent()
|
||||
.context("tokenizer.json has no parent directory")?;
|
||||
let tokenizer = DecisionTokenizer::from_dir(tokenizer_dir)?;
|
||||
let (dtype, backend) = match worker {
|
||||
Some(worker) => {
|
||||
let wanted = device_dtype(config.laya.amp_dtype.as_deref());
|
||||
let (handle, dtype) = worker
|
||||
.load_decision(files.dir.clone(), model_id.to_string(), wanted)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("worker load_decision: {e:#}"))?;
|
||||
(dtype, Backend::Worker { worker, handle })
|
||||
}
|
||||
None => {
|
||||
let dir = files.dir.clone();
|
||||
let model = tokio::task::spawn_blocking(move || {
|
||||
LayaModel::load(&dir, &candle_core::Device::Cpu, DType::F32)
|
||||
})
|
||||
.await
|
||||
.context("decision load task")??;
|
||||
(DType::F32, Backend::Cpu(Arc::new(model)))
|
||||
}
|
||||
};
|
||||
Ok(Self {
|
||||
name,
|
||||
config,
|
||||
calibration,
|
||||
tokenizer,
|
||||
dtype,
|
||||
backend,
|
||||
})
|
||||
}
|
||||
|
||||
/// Score one batch: one [`DecisionRow`] per sequence, in order.
|
||||
pub async fn decide(&self, batch: DecisionBatch) -> Result<Vec<DecisionRow>> {
|
||||
match &self.backend {
|
||||
Backend::Worker { worker, handle } => worker
|
||||
.decide(*handle, batch)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("{e:#}")),
|
||||
Backend::Cpu(model) => {
|
||||
let model = Arc::clone(model);
|
||||
tokio::task::spawn_blocking(move || model.forward(&batch))
|
||||
.await
|
||||
.context("decision forward task")?
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Release the weights. On a worker this drops them on the thread
|
||||
/// that owns the device context.
|
||||
pub async fn release(&self) {
|
||||
if let Backend::Worker { worker, handle } = &self.backend
|
||||
&& let Err(e) = worker.drop_decision(*handle).await
|
||||
{
|
||||
tracing::warn!(
|
||||
checkpoint = %self.name,
|
||||
error = %e,
|
||||
"decision unload: DropDecision RPC failed (weights may leak in worker slab)"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The precision to serve in on a GPU: the autocast dtype the reference
|
||||
/// itself runs there (`amp_dtype` in `rl_agent_config.json`), else bf16.
|
||||
fn device_dtype(amp_dtype: Option<&str>) -> DType {
|
||||
match amp_dtype {
|
||||
Some("fp16") | Some("float16") => DType::F16,
|
||||
Some("fp32") | Some("float32") => DType::F32,
|
||||
_ => DType::BF16,
|
||||
}
|
||||
}
|
||||
|
||||
pub struct LoadedDecisionModel {
|
||||
pub model_id: String,
|
||||
pub spec: ModelSpec,
|
||||
pub devices: Vec<u32>,
|
||||
/// At least one; the first is the default.
|
||||
pub checkpoints: Vec<DecisionCheckpoint>,
|
||||
/// Set when a forward returns a driver-level error, mirroring the
|
||||
/// other model kinds: `/models` reports it and requests fast-reject.
|
||||
pub poisoned: Arc<AtomicBool>,
|
||||
pub admission: super::admission::AdmissionController,
|
||||
}
|
||||
|
||||
impl LoadedDecisionModel {
|
||||
/// The checkpoint called `name`, or the default one for `None`.
|
||||
pub fn checkpoint(&self, name: Option<&str>) -> Option<&DecisionCheckpoint> {
|
||||
match name {
|
||||
None => self.checkpoints.first(),
|
||||
Some(n) => self.checkpoints.iter().find(|c| c.name == n),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn release(&self) {
|
||||
for c in &self.checkpoints {
|
||||
c.release().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn tiny() -> CheckpointFiles {
|
||||
let dir =
|
||||
std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src/harness/testdata/laya/tiny");
|
||||
CheckpointFiles {
|
||||
tokenizer: dir.join("tokenizer.json"),
|
||||
dir,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gpu_dtype_follows_the_reference_autocast() {
|
||||
assert_eq!(device_dtype(Some("bf16")), DType::BF16);
|
||||
assert_eq!(device_dtype(Some("fp16")), DType::F16);
|
||||
assert_eq!(device_dtype(None), DType::BF16);
|
||||
}
|
||||
|
||||
/// The whole async path on the CPU backend, and the device-worker
|
||||
/// path on a CPU build (the worker runs, with a CPU device), must
|
||||
/// give the same scores as the model called directly.
|
||||
#[tokio::test]
|
||||
async fn both_backends_score_like_the_model() {
|
||||
let files = tiny();
|
||||
let direct = LayaModel::load(&files.dir, &candle_core::Device::Cpu, DType::F32).unwrap();
|
||||
let batch = DecisionBatch {
|
||||
seqs: vec![super::super::arch::laya::DecisionSeq {
|
||||
input_ids: vec![5, 9, 11, 3, 7, 2],
|
||||
markers: vec![1, 3],
|
||||
qtype: 0,
|
||||
}],
|
||||
};
|
||||
let want = direct.forward(&batch).unwrap();
|
||||
|
||||
let cpu = DecisionCheckpoint::load("root".into(), files.clone(), "tiny", None)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(cpu.decide(batch.clone()).await.unwrap(), want);
|
||||
|
||||
// A CPU build's worker holds a CPU device, so this exercises the
|
||||
// job round-trip (load, decide, drop) without a GPU — and the
|
||||
// worker must fall back to f32 there, reporting that it did.
|
||||
let worker = DeviceWorkerHandle::spawn(0).unwrap();
|
||||
let on_worker =
|
||||
DecisionCheckpoint::load("root".into(), files, "tiny", Some(worker.clone()))
|
||||
.await
|
||||
.unwrap();
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
{
|
||||
assert_eq!(on_worker.dtype, DType::F32);
|
||||
assert_eq!(on_worker.decide(batch).await.unwrap(), want);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
{
|
||||
let got = on_worker.decide(batch).await.unwrap();
|
||||
for (g, w) in got[0].logits.iter().zip(&want[0].logits) {
|
||||
assert!((g - w).abs() < 5e-2, "{g} vs {w}");
|
||||
}
|
||||
}
|
||||
on_worker.release().await;
|
||||
assert!(on_worker.decide(DecisionBatch::default()).await.is_err());
|
||||
worker.shutdown().unwrap();
|
||||
}
|
||||
}
|
||||
@@ -13,12 +13,13 @@
|
||||
//! ARCH model state in this state slab will gain a companion
|
||||
//! `tp_models: HashMap<TpHandle, Box<TpLeaderModel>>`.
|
||||
|
||||
use crate::harness::arch::laya::LayaModel;
|
||||
use crate::harness::arch::snapshot::KvCacheSnapshot;
|
||||
use crate::harness::candle::ModelArch;
|
||||
#[cfg(feature = "cuda")]
|
||||
use crate::harness::device_worker::jobs::TpHandle;
|
||||
use crate::harness::device_worker::jobs::{
|
||||
ArchHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
|
||||
ArchHandle, DecisionHandle, DenseLoad, ImageHandle, ImageInput, Job, KvSnapshotId,
|
||||
};
|
||||
use crate::harness::image::ZImagePipeline;
|
||||
#[cfg(feature = "cuda")]
|
||||
@@ -65,6 +66,11 @@ struct DeviceWorkerState {
|
||||
image_models: HashMap<ImageHandle, Box<ZImagePipeline>>,
|
||||
/// Counter for minting fresh `ImageHandle`s.
|
||||
next_image_handle: u64,
|
||||
/// Laya decision-model slab (#336). Same lifecycle discipline as
|
||||
/// `image_models`, in its own namespace.
|
||||
decision_models: HashMap<DecisionHandle, Box<LayaModel>>,
|
||||
/// Counter for minting fresh `DecisionHandle`s.
|
||||
next_decision_handle: u64,
|
||||
/// Leader's NCCL state. Populated by `Job::NcclInit`; the
|
||||
/// underlying `Comm`'s libnccl handle lives bound to this thread
|
||||
/// for its entire lifetime. Subprocess workers maintain their own
|
||||
@@ -608,6 +614,61 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpLoadMtpHead {
|
||||
handle,
|
||||
config_json,
|
||||
safetensors_paths,
|
||||
reply,
|
||||
} => {
|
||||
let result = tp_load_mtp_head(&mut state, handle, &config_json, &safetensors_paths);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpMtpPrefillChunk {
|
||||
handle,
|
||||
shifted,
|
||||
start_pos,
|
||||
reply,
|
||||
} => {
|
||||
let result = tp_mtp(&mut state, handle, |m| {
|
||||
m.mtp_prefill_chunk(&shifted, start_pos)
|
||||
});
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpMtpDraft {
|
||||
handle,
|
||||
first_token,
|
||||
start_pos,
|
||||
k,
|
||||
reply,
|
||||
} => {
|
||||
let result = tp_mtp(&mut state, handle, |m| {
|
||||
m.mtp_draft(first_token, start_pos, k)
|
||||
});
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpMtpAdvance {
|
||||
handle,
|
||||
token,
|
||||
pos,
|
||||
reply,
|
||||
} => {
|
||||
let result = tp_mtp(&mut state, handle, |m| m.mtp_advance(token, pos));
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpForwardLogitsMulti {
|
||||
handle,
|
||||
tokens,
|
||||
offset,
|
||||
reply,
|
||||
} => {
|
||||
let result = tp_forward_logits_multi(&mut state, handle, &tokens, offset);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpForwardLogitsWithImages {
|
||||
handle,
|
||||
tokens,
|
||||
@@ -696,6 +757,62 @@ pub(crate) fn run(device_index: u32, rx: Receiver<Job>, poisoned: Arc<AtomicBool
|
||||
trim_device_pool(&state);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::LoadDecision {
|
||||
dir,
|
||||
model_id,
|
||||
dtype,
|
||||
reply,
|
||||
} => {
|
||||
let dtype = if state.device.is_cpu() {
|
||||
candle_core::DType::F32
|
||||
} else {
|
||||
dtype
|
||||
};
|
||||
let result = LayaModel::load(&dir, &state.device, dtype)
|
||||
.with_context(|| format!("load laya checkpoint {}", dir.display()))
|
||||
.map(|model| {
|
||||
let handle = DecisionHandle(state.next_decision_handle);
|
||||
state.next_decision_handle = state.next_decision_handle.wrapping_add(1);
|
||||
state.decision_models.insert(handle, Box::new(model));
|
||||
tracing::info!(
|
||||
device_index,
|
||||
model_id,
|
||||
checkpoint = %dir.display(),
|
||||
handle = handle.0,
|
||||
"device worker: decision model loaded"
|
||||
);
|
||||
(handle, dtype)
|
||||
});
|
||||
trim_device_pool(&state);
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::DropDecision { handle, reply } => {
|
||||
let removed = state.decision_models.remove(&handle);
|
||||
let was_present = removed.is_some();
|
||||
// Explicit drop on this thread, as for DropArch/DropImage.
|
||||
drop(removed);
|
||||
if was_present {
|
||||
trim_device_pool(&state);
|
||||
}
|
||||
tracing::debug!(
|
||||
device_index,
|
||||
handle = handle.0,
|
||||
was_present,
|
||||
"device worker: decision model dropped"
|
||||
);
|
||||
let _ = reply.send(());
|
||||
}
|
||||
Job::Decide {
|
||||
handle,
|
||||
batch,
|
||||
reply,
|
||||
} => {
|
||||
let result = match state.decision_models.get(&handle) {
|
||||
Some(model) => model.forward(&batch),
|
||||
None => Err(anyhow::anyhow!("no decision model for handle {}", handle.0)),
|
||||
};
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
Job::Shutdown => unreachable!("Shutdown should break above"),
|
||||
}
|
||||
}
|
||||
@@ -768,6 +885,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
|
||||
next_kv_snapshot_id: 1,
|
||||
image_models: HashMap::new(),
|
||||
next_image_handle: 1,
|
||||
decision_models: HashMap::new(),
|
||||
next_decision_handle: 1,
|
||||
nccl: NcclState::new(),
|
||||
tp_models: HashMap::new(),
|
||||
next_tp_handle: 1,
|
||||
@@ -786,6 +905,8 @@ fn init_state(device_index: u32) -> DeviceWorkerState {
|
||||
next_kv_snapshot_id: 1,
|
||||
image_models: HashMap::new(),
|
||||
next_image_handle: 1,
|
||||
decision_models: HashMap::new(),
|
||||
next_decision_handle: 1,
|
||||
nccl: NcclState::new(),
|
||||
}
|
||||
}
|
||||
@@ -1314,6 +1435,85 @@ fn tp_forward_logits(
|
||||
Ok(values)
|
||||
}
|
||||
|
||||
/// Run something on the leader's TP draft head (#96).
|
||||
///
|
||||
/// The head only exists on the `qwen3_5` arch and only on a model that
|
||||
/// was asked to load one, so both misses are reported rather than
|
||||
/// silently skipped.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn tp_mtp<T>(
|
||||
state: &mut DeviceWorkerState,
|
||||
handle: TpHandle,
|
||||
f: impl FnOnce(&mut super::super::tp::tp_qwen3_5::TpQwen3_5ForCausalLM) -> anyhow::Result<T>,
|
||||
) -> anyhow::Result<T> {
|
||||
let model = state
|
||||
.tp_models
|
||||
.get_mut(&handle)
|
||||
.ok_or_else(|| anyhow::anyhow!("tp mtp: no model for handle {}", handle.0))?;
|
||||
match &mut **model {
|
||||
super::super::tp::TpLeaderModel::Qwen3_5(m) => f(m),
|
||||
super::super::tp::TpLeaderModel::Qwen3(_) => {
|
||||
anyhow::bail!("tp mtp: the qwen3 arch ships no MTP head")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Load the draft head onto the leader's TP model (#96).
|
||||
#[cfg(feature = "cuda")]
|
||||
fn tp_load_mtp_head(
|
||||
state: &mut DeviceWorkerState,
|
||||
handle: TpHandle,
|
||||
config_json: &str,
|
||||
safetensors_paths: &[String],
|
||||
) -> anyhow::Result<()> {
|
||||
use candle_core::DType;
|
||||
|
||||
let cfg = super::super::tp::tp_qwen3_5::Config::from_config_json(config_json)
|
||||
.context("tp mtp: parse config.json")?;
|
||||
let paths: Vec<std::path::PathBuf> = safetensors_paths.iter().map(Into::into).collect();
|
||||
let device = state.device.clone();
|
||||
// SAFETY: the same checkpoint files the shard load already mmapped.
|
||||
let vb = unsafe {
|
||||
candle_nn::var_builder::ShardedSafeTensors::var_builder(&paths, DType::BF16, &device)?
|
||||
};
|
||||
tp_mtp(state, handle, |m| m.load_mtp_head(&cfg.text_config, &vb))
|
||||
}
|
||||
|
||||
/// TP mirror of [`forward_logits_multi`] on the leader's shard (#96).
|
||||
/// The caller has already fanned the matching `GenerateStepMulti` out
|
||||
/// to the subprocess ranks; the `AllReduce` collectives inside this
|
||||
/// forward block until they arrive.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn tp_forward_logits_multi(
|
||||
state: &mut DeviceWorkerState,
|
||||
handle: TpHandle,
|
||||
tokens: &[u32],
|
||||
offset: usize,
|
||||
) -> anyhow::Result<Vec<Vec<f32>>> {
|
||||
use candle_core::{DType, Tensor};
|
||||
|
||||
anyhow::ensure!(
|
||||
!tokens.is_empty(),
|
||||
"TpForwardLogitsMulti: empty token slice"
|
||||
);
|
||||
let input = Tensor::new(tokens, &state.device)?.unsqueeze(0)?;
|
||||
let model = state
|
||||
.tp_models
|
||||
.get_mut(&handle)
|
||||
.ok_or_else(|| anyhow::anyhow!("TpForwardLogitsMulti: no model for handle {}", handle.0))?;
|
||||
|
||||
// (1, L, vocab) → one CPU row per position.
|
||||
let logits = model.forward_multi(&input, offset)?.to_dtype(DType::F32)?;
|
||||
let (_, positions, vocab) = logits.dims3()?;
|
||||
anyhow::ensure!(
|
||||
positions == tokens.len(),
|
||||
"TpForwardLogitsMulti: model returned {positions} rows for {} tokens",
|
||||
tokens.len()
|
||||
);
|
||||
let flat: Vec<f32> = logits.flatten_all()?.to_vec1()?;
|
||||
Ok(flat.chunks(vocab).map(<[f32]>::to_vec).collect())
|
||||
}
|
||||
|
||||
/// TP-equivalent of [`forward_logits_batch`] on the leader's shard
|
||||
/// (#98). The caller has already fanned the matching
|
||||
/// `GenerateStepBatch` out to the subprocess ranks.
|
||||
@@ -1835,6 +2035,15 @@ fn drain_poisoned(job: Job, device_index: u32) {
|
||||
Job::GenerateImage { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::LoadDecision { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::DropDecision { reply, .. } => {
|
||||
let _ = reply.send(());
|
||||
}
|
||||
Job::Decide { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
Job::ForwardLogitsWithImages { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
@@ -1893,6 +2102,26 @@ fn drain_poisoned(job: Job, device_index: u32) {
|
||||
let _ = reply.send(());
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpLoadMtpHead { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpMtpPrefillChunk { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpMtpDraft { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpMtpAdvance { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpForwardLogitsMulti { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
#[cfg(feature = "cuda")]
|
||||
Job::TpForwardLogits { reply, .. } => {
|
||||
let _ = reply.send(Err(err()));
|
||||
}
|
||||
|
||||
@@ -59,6 +59,15 @@ pub struct TpHandle(pub u64);
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct ImageHandle(pub u64);
|
||||
|
||||
/// Opaque handle to a Laya decision model ([`LayaModel`]) stored in the
|
||||
/// worker thread's decision slab (#336). Same discipline as
|
||||
/// [`ImageHandle`]: construction, every forward and `Drop` happen on the
|
||||
/// worker thread, and `Job::DropDecision` is the only way to release it.
|
||||
///
|
||||
/// [`LayaModel`]: crate::harness::arch::laya::LayaModel
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
|
||||
pub struct DecisionHandle(pub u64);
|
||||
|
||||
/// Opaque handle to a prefix-cache snapshot (#11) stored worker-side
|
||||
/// next to the model slab. Scoped to the `ArchHandle` it was captured
|
||||
/// from — `Job::DropArch` drops every snapshot under its handle. The
|
||||
@@ -394,6 +403,59 @@ pub enum Job {
|
||||
/// without holding a device tensor. The caller is also
|
||||
/// responsible for fan-out to subprocess ranks and drain — only
|
||||
/// the leader's forward moves into the worker thread.
|
||||
/// Load the MTP draft head onto the leader's TP model (#96).
|
||||
///
|
||||
/// Leader-only and deliberately not part of `TpLoadShard`: the
|
||||
/// subprocess ranks have no use for a head they never run, and an
|
||||
/// unsharded head on every rank would cost 810 MB each to sit idle.
|
||||
#[cfg(feature = "cuda")]
|
||||
TpLoadMtpHead {
|
||||
handle: TpHandle,
|
||||
config_json: String,
|
||||
safetensors_paths: Vec<String>,
|
||||
reply: oneshot::Sender<Result<()>>,
|
||||
},
|
||||
/// Walk the leader's draft head over a prompt chunk (#96), using
|
||||
/// the hidden state the target's forward just produced.
|
||||
#[cfg(feature = "cuda")]
|
||||
TpMtpPrefillChunk {
|
||||
handle: TpHandle,
|
||||
/// The chunk's tokens shifted one left — the head trails the
|
||||
/// target by a token.
|
||||
shifted: Vec<u32>,
|
||||
start_pos: usize,
|
||||
reply: oneshot::Sender<Result<()>>,
|
||||
},
|
||||
/// Draft `k` tokens on the leader and put the head's cache back
|
||||
/// (#96). Commits nothing: the snapshot/restore pair lives inside
|
||||
/// the job so a caller cannot forget it.
|
||||
#[cfg(feature = "cuda")]
|
||||
TpMtpDraft {
|
||||
handle: TpHandle,
|
||||
first_token: u32,
|
||||
start_pos: usize,
|
||||
k: usize,
|
||||
reply: oneshot::Sender<Result<Vec<u32>>>,
|
||||
},
|
||||
/// Advance the leader's draft head over one committed token (#96).
|
||||
#[cfg(feature = "cuda")]
|
||||
TpMtpAdvance {
|
||||
handle: TpHandle,
|
||||
token: u32,
|
||||
pos: usize,
|
||||
reply: oneshot::Sender<Result<()>>,
|
||||
},
|
||||
/// TP mirror of `ForwardLogitsMulti` (#96): logits at every
|
||||
/// position of a multi-token forward on the leader's shard. The
|
||||
/// caller has already fanned the matching `GenerateStepMulti` out
|
||||
/// to the subprocess ranks, whose collectives this one waits on.
|
||||
#[cfg(feature = "cuda")]
|
||||
TpForwardLogitsMulti {
|
||||
handle: TpHandle,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
reply: oneshot::Sender<Result<Vec<Vec<f32>>>>,
|
||||
},
|
||||
#[cfg(feature = "cuda")]
|
||||
TpForwardLogits {
|
||||
handle: TpHandle,
|
||||
@@ -496,6 +558,33 @@ pub enum Job {
|
||||
max_dim: usize,
|
||||
reply: oneshot::Sender<Result<crate::harness::image::ImageGenResult>>,
|
||||
},
|
||||
/// Load one Laya checkpoint directory (#336) — the repo root or one
|
||||
/// of its subfolders, already on local disk — into the decision
|
||||
/// slab. File resolution (hf-hub) happens on the async side.
|
||||
///
|
||||
/// `dtype` is the precision wanted on a GPU; a worker whose device
|
||||
/// is the CPU loads in f32 instead (the CPU backend has no bf16
|
||||
/// matmul). The reply carries the dtype actually used.
|
||||
LoadDecision {
|
||||
dir: PathBuf,
|
||||
model_id: String,
|
||||
dtype: candle_core::DType,
|
||||
reply: oneshot::Sender<Result<(DecisionHandle, candle_core::DType)>>,
|
||||
},
|
||||
/// Remove a decision model from the slab and drop it on the worker
|
||||
/// thread. Mirrors `DropArch` / `DropImage`.
|
||||
DropDecision {
|
||||
handle: DecisionHandle,
|
||||
reply: oneshot::Sender<()>,
|
||||
},
|
||||
/// Score one batch of assembled questions (#336). One job per
|
||||
/// request: the whole batch is a single encoder forward, and the
|
||||
/// reply is CPU-side scores — no device tensor crosses back.
|
||||
Decide {
|
||||
handle: DecisionHandle,
|
||||
batch: crate::harness::arch::laya::DecisionBatch,
|
||||
reply: oneshot::Sender<Result<Vec<crate::harness::arch::laya::DecisionRow>>>,
|
||||
},
|
||||
/// Tell the worker to break its dispatch loop and exit. Any jobs
|
||||
/// queued after this in the channel reply `Err` to their oneshot
|
||||
/// senders (the senders are dropped on the worker's exit, which
|
||||
|
||||
@@ -51,7 +51,7 @@ use tokio::sync::oneshot;
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
pub use jobs::TpHandle;
|
||||
pub use jobs::{ArchHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
|
||||
pub use jobs::{ArchHandle, DecisionHandle, DenseLoad, ImageHandle, Job, KvSnapshotId};
|
||||
|
||||
/// Errors returned by `DeviceWorkerHandle` submit methods.
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
@@ -337,6 +337,85 @@ impl DeviceWorkerHandle {
|
||||
}
|
||||
}
|
||||
|
||||
/// Load one Laya checkpoint directory (#336) on the worker thread.
|
||||
/// Returns the handle and the dtype the weights were loaded in.
|
||||
pub async fn load_decision(
|
||||
&self,
|
||||
dir: std::path::PathBuf,
|
||||
model_id: String,
|
||||
dtype: candle_core::DType,
|
||||
) -> Result<(DecisionHandle, candle_core::DType), WorkerError> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
return Err(WorkerError::Poisoned {
|
||||
device_index: self.device_index,
|
||||
});
|
||||
}
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(Job::LoadDecision {
|
||||
dir,
|
||||
model_id,
|
||||
dtype,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})?;
|
||||
match reply_rx.await {
|
||||
Ok(result) => result.map_err(WorkerError::from),
|
||||
Err(_) => Err(WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Drop the decision model for `handle` on the worker thread. Same
|
||||
/// poison-tolerant contract as [`Self::drop_arch`].
|
||||
pub async fn drop_decision(&self, handle: DecisionHandle) -> Result<(), WorkerError> {
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(Job::DropDecision {
|
||||
handle,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})?;
|
||||
reply_rx.await.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})
|
||||
}
|
||||
|
||||
/// Score one batch of assembled questions (#336). The reply is
|
||||
/// CPU-side scores; no device tensor crosses back.
|
||||
pub async fn decide(
|
||||
&self,
|
||||
handle: DecisionHandle,
|
||||
batch: crate::harness::arch::laya::DecisionBatch,
|
||||
) -> Result<Vec<crate::harness::arch::laya::DecisionRow>, WorkerError> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
return Err(WorkerError::Poisoned {
|
||||
device_index: self.device_index,
|
||||
});
|
||||
}
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(Job::Decide {
|
||||
handle,
|
||||
batch,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})?;
|
||||
match reply_rx.await {
|
||||
Ok(result) => result.map_err(WorkerError::from),
|
||||
Err(_) => Err(WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// Tell the worker to drop the `ModelArch` for `handle` on the
|
||||
/// worker thread (so CUDA tensors release on the right context).
|
||||
/// Returns `Ok(())` even if the handle wasn't in the slab — Drop
|
||||
@@ -968,6 +1047,137 @@ impl DeviceWorkerHandle {
|
||||
/// logits as `Vec<f32>` ready for sampling. The caller is
|
||||
/// responsible for fan-out / drain of the subprocess workers
|
||||
/// concurrently with this call.
|
||||
/// Load the MTP draft head onto the leader's TP model (#96).
|
||||
#[cfg(feature = "cuda")]
|
||||
pub async fn tp_load_mtp_head(
|
||||
&self,
|
||||
handle: TpHandle,
|
||||
config_json: String,
|
||||
safetensors_paths: Vec<String>,
|
||||
) -> Result<(), WorkerError> {
|
||||
self.tp_mtp_job(|reply| Job::TpLoadMtpHead {
|
||||
handle,
|
||||
config_json,
|
||||
safetensors_paths,
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Walk the leader's draft head over a prompt chunk (#96).
|
||||
#[cfg(feature = "cuda")]
|
||||
pub async fn tp_mtp_prefill_chunk(
|
||||
&self,
|
||||
handle: TpHandle,
|
||||
shifted: Vec<u32>,
|
||||
start_pos: usize,
|
||||
) -> Result<(), WorkerError> {
|
||||
self.tp_mtp_job(|reply| Job::TpMtpPrefillChunk {
|
||||
handle,
|
||||
shifted,
|
||||
start_pos,
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Draft `k` tokens on the leader, rewinding the head afterwards
|
||||
/// (#96). Commits nothing.
|
||||
#[cfg(feature = "cuda")]
|
||||
pub async fn tp_mtp_draft(
|
||||
&self,
|
||||
handle: TpHandle,
|
||||
first_token: u32,
|
||||
start_pos: usize,
|
||||
k: usize,
|
||||
) -> Result<Vec<u32>, WorkerError> {
|
||||
self.tp_mtp_job(|reply| Job::TpMtpDraft {
|
||||
handle,
|
||||
first_token,
|
||||
start_pos,
|
||||
k,
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Advance the leader's draft head over one committed token (#96).
|
||||
#[cfg(feature = "cuda")]
|
||||
pub async fn tp_mtp_advance(
|
||||
&self,
|
||||
handle: TpHandle,
|
||||
token: u32,
|
||||
pos: usize,
|
||||
) -> Result<(), WorkerError> {
|
||||
self.tp_mtp_job(|reply| Job::TpMtpAdvance {
|
||||
handle,
|
||||
token,
|
||||
pos,
|
||||
reply,
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
/// Shared submit/await for the leader-only draft-head jobs (#96):
|
||||
/// same poison check and channel handling as every other job, in
|
||||
/// one place because these four differ only in their payload.
|
||||
#[cfg(feature = "cuda")]
|
||||
async fn tp_mtp_job<T>(
|
||||
&self,
|
||||
build: impl FnOnce(oneshot::Sender<anyhow::Result<T>>) -> Job,
|
||||
) -> Result<T, WorkerError> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
return Err(WorkerError::Poisoned {
|
||||
device_index: self.device_index,
|
||||
});
|
||||
}
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(build(reply_tx))
|
||||
.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})?;
|
||||
match reply_rx.await {
|
||||
Ok(result) => result.map_err(WorkerError::from),
|
||||
Err(_) => Err(WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// TP mirror of [`Self::forward_logits_multi`] (#96): one CPU
|
||||
/// `[vocab]` row per position from the leader's shard.
|
||||
#[cfg(feature = "cuda")]
|
||||
pub async fn tp_forward_logits_multi(
|
||||
&self,
|
||||
handle: TpHandle,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
) -> Result<Vec<Vec<f32>>, WorkerError> {
|
||||
if self.poisoned.load(Ordering::Acquire) {
|
||||
return Err(WorkerError::Poisoned {
|
||||
device_index: self.device_index,
|
||||
});
|
||||
}
|
||||
let (reply_tx, reply_rx) = oneshot::channel();
|
||||
self.tx
|
||||
.send(Job::TpForwardLogitsMulti {
|
||||
handle,
|
||||
tokens,
|
||||
offset,
|
||||
reply: reply_tx,
|
||||
})
|
||||
.map_err(|_| WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
})?;
|
||||
match reply_rx.await {
|
||||
Ok(result) => result.map_err(WorkerError::from),
|
||||
Err(_) => Err(WorkerError::Gone {
|
||||
device_index: self.device_index,
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "cuda")]
|
||||
pub async fn tp_forward_logits(
|
||||
&self,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -6,10 +6,13 @@ pub mod candle;
|
||||
pub mod capability;
|
||||
pub mod chat_template;
|
||||
pub mod context_limit;
|
||||
pub mod decision;
|
||||
pub mod decision_model;
|
||||
pub mod device_worker;
|
||||
pub mod engine;
|
||||
pub mod image;
|
||||
pub mod isq_cache;
|
||||
pub mod lang_detect;
|
||||
pub mod prefix_cache;
|
||||
pub mod preflight;
|
||||
pub mod preprocess;
|
||||
|
||||
@@ -54,6 +54,11 @@ pub enum SourceFormat {
|
||||
/// Takes precedence over the safetensors sniff because diffusers
|
||||
/// repos also contain `.safetensors` shards (in subdirs).
|
||||
Diffusers,
|
||||
/// A Laya decision-model checkpoint (#336): `rl_agent_config.json`
|
||||
/// plus `encoder/config.json` at the root, no top-level
|
||||
/// `config.json`. Served by the decision path — single device,
|
||||
/// unquantized. Takes precedence over the safetensors sniff.
|
||||
DecisionModel,
|
||||
}
|
||||
|
||||
/// Output of `preflight` for a load that can proceed. Carries the
|
||||
@@ -115,6 +120,12 @@ pub enum PreflightError {
|
||||
available: Vec<String>,
|
||||
nearest: Option<String>,
|
||||
},
|
||||
|
||||
/// The repo's format cannot be served with an option the spec
|
||||
/// asked for (e.g. `tensor_parallel` or `quant` on a decision
|
||||
/// model). `reason` says which and what to set instead.
|
||||
#[error("cannot load '{model_id}' as requested: {reason}")]
|
||||
Unsupported { model_id: String, reason: String },
|
||||
}
|
||||
|
||||
/// Run the placement check.
|
||||
@@ -230,6 +241,24 @@ pub async fn preflight(
|
||||
// The architecture compatibility check stays where it is —
|
||||
// `check_dense_config_supported` runs once `config.json` is
|
||||
// on disk, since it needs the parsed JSON.
|
||||
// Decision models (#336) load unquantized on one device.
|
||||
(SourceFormat::DecisionModel, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!(
|
||||
"decision models load on a single device; tensor_parallel={tp} is not supported"
|
||||
),
|
||||
}),
|
||||
(SourceFormat::DecisionModel, _, Some(q)) => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("decision models load unquantized; remove quant={q:?}"),
|
||||
}),
|
||||
(SourceFormat::DecisionModel, _, None) => Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
}),
|
||||
|
||||
(SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
|
||||
Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
@@ -295,6 +324,9 @@ pub fn classify(filenames: &[&str]) -> SourceFormat {
|
||||
if crate::harness::image::is_diffusers_layout(filenames) {
|
||||
return SourceFormat::Diffusers;
|
||||
}
|
||||
if is_decision_layout(filenames) {
|
||||
return SourceFormat::DecisionModel;
|
||||
}
|
||||
let mut gguf_quants: Vec<String> = filenames
|
||||
.iter()
|
||||
.filter(|f| f.to_lowercase().ends_with(".gguf"))
|
||||
@@ -318,6 +350,16 @@ pub fn classify(filenames: &[&str]) -> SourceFormat {
|
||||
}
|
||||
}
|
||||
|
||||
/// A Laya decision checkpoint at the repo root: its decision config and
|
||||
/// its encoder config, with no top-level `config.json` (a causal LM
|
||||
/// that happened to ship a file with the same name still classifies as
|
||||
/// a text model).
|
||||
pub fn is_decision_layout(filenames: &[&str]) -> bool {
|
||||
filenames.contains(&"rl_agent_config.json")
|
||||
&& filenames.contains(&"encoder/config.json")
|
||||
&& !filenames.contains(&"config.json")
|
||||
}
|
||||
|
||||
/// Mirror of the quant-matching logic in `candle.rs::resolve_files` so
|
||||
/// preflight picks the same file the downstream loader would. Empty
|
||||
/// quant returns the first `.gguf` (any quant). Lowercased substring
|
||||
@@ -651,6 +693,20 @@ mod tests {
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
}),
|
||||
(SourceFormat::DecisionModel, tp, _) if tp > 1 => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("decision models load on a single device; tensor_parallel={tp}"),
|
||||
}),
|
||||
(SourceFormat::DecisionModel, _, Some(q)) => Err(PreflightError::Unsupported {
|
||||
model_id: source_id.to_string(),
|
||||
reason: format!("decision models load unquantized; remove quant={q:?}"),
|
||||
}),
|
||||
(SourceFormat::DecisionModel, _, None) => Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
format: format.clone(),
|
||||
tp_size,
|
||||
picked_quant_file: None,
|
||||
}),
|
||||
(SourceFormat::DenseSafetensors { .. } | SourceFormat::Mixed { .. }, _, _) => {
|
||||
Ok(PlacementPlan {
|
||||
model_id: source_id.to_string(),
|
||||
@@ -674,6 +730,66 @@ mod tests {
|
||||
assert_eq!(classify(&files), SourceFormat::Diffusers);
|
||||
}
|
||||
|
||||
/// The released `convaiinnovations/laya` listing: the English
|
||||
/// checkpoint at the root, two more in subfolders, and no root
|
||||
/// `config.json`.
|
||||
const LAYA_FILES: [&str; 12] = [
|
||||
"README.md",
|
||||
"rl_agent_config.json",
|
||||
"encoder/config.json",
|
||||
"model.safetensors",
|
||||
"tokenizer/tokenizer.json",
|
||||
"tokenizer/tokenizer_config.json",
|
||||
"multilingual/rl_agent_config.json",
|
||||
"multilingual/encoder/config.json",
|
||||
"multilingual/model.safetensors",
|
||||
"typed-decisions/rl_agent_config.json",
|
||||
"typed-decisions/encoder/config.json",
|
||||
"typed-decisions/model.safetensors",
|
||||
];
|
||||
|
||||
#[test]
|
||||
fn classify_laya_layout() {
|
||||
assert_eq!(classify(&LAYA_FILES), SourceFormat::DecisionModel);
|
||||
// A causal LM that also ships an rl config is still a text model.
|
||||
let lm = [
|
||||
"config.json",
|
||||
"rl_agent_config.json",
|
||||
"encoder/config.json",
|
||||
"model.safetensors",
|
||||
];
|
||||
assert!(matches!(
|
||||
classify(&lm),
|
||||
SourceFormat::DenseSafetensors { .. }
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_decision_model() {
|
||||
let fmt = classify(&LAYA_FILES);
|
||||
let plan = decide(
|
||||
&spec("convaiinnovations/laya", None, None),
|
||||
&fmt,
|
||||
&LAYA_FILES,
|
||||
)
|
||||
.expect("single-device unquantized decision load should pass");
|
||||
assert_eq!(plan.format, SourceFormat::DecisionModel);
|
||||
let err = decide(
|
||||
&spec("convaiinnovations/laya", Some(2), None),
|
||||
&fmt,
|
||||
&LAYA_FILES,
|
||||
)
|
||||
.expect_err("tp>1 must be refused");
|
||||
assert!(matches!(err, PreflightError::Unsupported { .. }));
|
||||
let err = decide(
|
||||
&spec("convaiinnovations/laya", None, Some("q8_0")),
|
||||
&fmt,
|
||||
&LAYA_FILES,
|
||||
)
|
||||
.expect_err("quant must be refused");
|
||||
assert!(matches!(err, PreflightError::Unsupported { .. }));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn feasibility_diffusers_single_gpu_ok() {
|
||||
let files = ["model_index.json", "transformer/config.json"];
|
||||
|
||||
+347
@@ -0,0 +1,347 @@
|
||||
{
|
||||
"header": {
|
||||
"laya_version": "0.3.21",
|
||||
"torch_version": "2.14.0+cu130",
|
||||
"device": "cpu",
|
||||
"dtype": "float32",
|
||||
"checkpoints": {
|
||||
"english": {
|
||||
"repo": "convaiinnovations/laya",
|
||||
"revision": "55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851",
|
||||
"cfg": {
|
||||
"max_len": 512,
|
||||
"head_max_len": 192,
|
||||
"temperature": [
|
||||
1.6369030475616455,
|
||||
1.2514300346374512,
|
||||
1.983399510383606
|
||||
],
|
||||
"temperature_by_options": {
|
||||
"choice:3-5": 1.7601518630981445,
|
||||
"choice:6-10": 1.0000158548355103,
|
||||
"score:3-5": 1.2514300346374512,
|
||||
"noul:2": 1.983399510383606,
|
||||
"choice:11+": 0.10058280825614929,
|
||||
"choice:2": 1.9063563346862793
|
||||
}
|
||||
},
|
||||
"special_ids": {
|
||||
"cls": 50281,
|
||||
"sep": 50282,
|
||||
"mask": 50284,
|
||||
"pad": 50283
|
||||
}
|
||||
},
|
||||
"multilingual": {
|
||||
"repo": "convaiinnovations/laya/multilingual",
|
||||
"revision": "55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851",
|
||||
"cfg": {
|
||||
"max_len": 1024,
|
||||
"head_max_len": 256,
|
||||
"temperature": [
|
||||
1.0,
|
||||
1.0,
|
||||
1.0
|
||||
],
|
||||
"temperature_by_options": {}
|
||||
},
|
||||
"special_ids": {
|
||||
"cls": 2,
|
||||
"sep": 1,
|
||||
"mask": 4,
|
||||
"pad": 0
|
||||
}
|
||||
},
|
||||
"typed-decisions": {
|
||||
"repo": "convaiinnovations/laya/typed-decisions",
|
||||
"revision": "55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851",
|
||||
"cfg": {
|
||||
"max_len": 1024,
|
||||
"head_max_len": 256,
|
||||
"temperature": [
|
||||
1.0148024559020996,
|
||||
1.0374259948730469,
|
||||
1.0575125217437744
|
||||
],
|
||||
"temperature_by_options": {
|
||||
"choice:3-5": 1.7601518630981445,
|
||||
"choice:6-10": 1.0000158548355103,
|
||||
"score:3-5": 1.2514300346374512,
|
||||
"noul:2": 1.983399510383606,
|
||||
"choice:11+": 0.10058280825614929,
|
||||
"choice:2": 1.9063563346862793
|
||||
}
|
||||
},
|
||||
"special_ids": {
|
||||
"cls": 50281,
|
||||
"sep": 50282,
|
||||
"mask": 50284,
|
||||
"pad": 50283
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"cases": [
|
||||
{
|
||||
"state": "I was charged twice this month, I want my money back",
|
||||
"analysis": {
|
||||
"script": "latin",
|
||||
"script_profile": {
|
||||
"latin": 1.0
|
||||
},
|
||||
"language": "en",
|
||||
"is_english": true,
|
||||
"language_undecided": false,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "Me han cobrado dos veces este mes",
|
||||
"analysis": {
|
||||
"script": "latin",
|
||||
"script_profile": {
|
||||
"latin": 1.0
|
||||
},
|
||||
"language": "es",
|
||||
"is_english": false,
|
||||
"language_undecided": false,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "Die Rechnung ist falsch, bitte korrigieren",
|
||||
"analysis": {
|
||||
"script": "latin",
|
||||
"script_profile": {
|
||||
"latin": 1.0
|
||||
},
|
||||
"language": "de",
|
||||
"is_english": false,
|
||||
"language_undecided": false,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "Je voudrais annuler ma commande",
|
||||
"analysis": {
|
||||
"script": "latin",
|
||||
"script_profile": {
|
||||
"latin": 1.0
|
||||
},
|
||||
"language": "fr",
|
||||
"is_english": false,
|
||||
"language_undecided": false,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "मुझे पैसे वापस चाहिए",
|
||||
"analysis": {
|
||||
"script": "devanagari",
|
||||
"script_profile": {
|
||||
"devanagari": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": false,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 1.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "ខ្ញុំចង់បានប្រាក់វិញ",
|
||||
"analysis": {
|
||||
"script": "khmer",
|
||||
"script_profile": {
|
||||
"khmer": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": false,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 1.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "这个月我被扣了两次钱",
|
||||
"analysis": {
|
||||
"script": "han",
|
||||
"script_profile": {
|
||||
"han": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": false,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 1.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "注文をキャンセルしたいです",
|
||||
"analysis": {
|
||||
"script": "kana",
|
||||
"script_profile": {
|
||||
"han": 0.15384615384615385,
|
||||
"kana": 0.8461538461538461
|
||||
},
|
||||
"language": null,
|
||||
"is_english": false,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 1.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "تم خصم المبلغ مرتين",
|
||||
"analysis": {
|
||||
"script": "arabic",
|
||||
"script_profile": {
|
||||
"arabic": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": false,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 1.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "Мне дважды списали деньги",
|
||||
"analysis": {
|
||||
"script": "cyrillic",
|
||||
"script_profile": {
|
||||
"cyrillic": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": false,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 1.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "Order #5521 — 注文をキャンセルしたいです. Please help.",
|
||||
"analysis": {
|
||||
"script": "kana",
|
||||
"script_profile": {
|
||||
"latin": 0.5357142857142857,
|
||||
"han": 0.07142857142857142,
|
||||
"kana": 0.39285714285714285
|
||||
},
|
||||
"language": null,
|
||||
"is_english": false,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.4643,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "ok",
|
||||
"analysis": {
|
||||
"script": "latin",
|
||||
"script_profile": {
|
||||
"latin": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": true,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "",
|
||||
"analysis": {
|
||||
"script": "unknown",
|
||||
"script_profile": {},
|
||||
"language": null,
|
||||
"is_english": true,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "12345 67890",
|
||||
"analysis": {
|
||||
"script": "unknown",
|
||||
"script_profile": {},
|
||||
"language": null,
|
||||
"is_english": true,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": "café résumé naïve",
|
||||
"analysis": {
|
||||
"script": "latin",
|
||||
"script_profile": {
|
||||
"latin": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": false,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.2353,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": {
|
||||
"subject": "Invoice",
|
||||
"body": "Bitte zahlen Sie bis Freitag"
|
||||
},
|
||||
"analysis": {
|
||||
"script": "latin",
|
||||
"script_profile": {
|
||||
"latin": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": true,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
},
|
||||
{
|
||||
"state": [
|
||||
"hello",
|
||||
"bonjour",
|
||||
"hola"
|
||||
],
|
||||
"analysis": {
|
||||
"script": "latin",
|
||||
"script_profile": {
|
||||
"latin": 1.0
|
||||
},
|
||||
"language": null,
|
||||
"is_english": true,
|
||||
"language_undecided": true,
|
||||
"diacritic_rate": 0.0,
|
||||
"non_latin_fraction": 0.0,
|
||||
"mixed_segment": null
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
+46824
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
+5610
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,52 @@
|
||||
{
|
||||
"attention_bias": false,
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 1,
|
||||
"classifier_activation": "gelu",
|
||||
"classifier_bias": false,
|
||||
"classifier_dropout": 0.0,
|
||||
"classifier_pooling": "cls",
|
||||
"cls_token_id": 1,
|
||||
"decoder_bias": true,
|
||||
"deterministic_flash_attn": false,
|
||||
"embedding_dropout": 0.0,
|
||||
"eos_token_id": 2,
|
||||
"global_attn_every_n_layers": 3,
|
||||
"hidden_activation": "gelu",
|
||||
"hidden_size": 32,
|
||||
"initializer_cutoff_factor": 2.0,
|
||||
"initializer_range": 0.3,
|
||||
"intermediate_size": 48,
|
||||
"layer_types": [
|
||||
"full_attention",
|
||||
"sliding_attention",
|
||||
"sliding_attention",
|
||||
"full_attention"
|
||||
],
|
||||
"local_attention": 8,
|
||||
"max_position_embeddings": 128,
|
||||
"mlp_bias": false,
|
||||
"mlp_dropout": 0.0,
|
||||
"model_type": "modernbert",
|
||||
"norm_bias": false,
|
||||
"norm_eps": 1e-05,
|
||||
"num_attention_heads": 2,
|
||||
"num_hidden_layers": 4,
|
||||
"pad_token_id": 0,
|
||||
"rope_parameters": {
|
||||
"full_attention": {
|
||||
"rope_theta": 1000.0,
|
||||
"rope_type": "default"
|
||||
},
|
||||
"sliding_attention": {
|
||||
"rope_theta": 300.0,
|
||||
"rope_type": "default"
|
||||
}
|
||||
},
|
||||
"sep_token_id": 2,
|
||||
"sparse_pred_ignore_index": -100,
|
||||
"sparse_prediction": false,
|
||||
"tie_word_embeddings": true,
|
||||
"transformers_version": "5.17.0",
|
||||
"vocab_size": 97
|
||||
}
|
||||
@@ -0,0 +1,182 @@
|
||||
{
|
||||
"seed": 7,
|
||||
"seqs": [
|
||||
{
|
||||
"input_ids": [
|
||||
1,
|
||||
93,
|
||||
82,
|
||||
87,
|
||||
51,
|
||||
46,
|
||||
89,
|
||||
6,
|
||||
68,
|
||||
94,
|
||||
86,
|
||||
43,
|
||||
23,
|
||||
70,
|
||||
42,
|
||||
60,
|
||||
35,
|
||||
3,
|
||||
86,
|
||||
7,
|
||||
20,
|
||||
18,
|
||||
90,
|
||||
84,
|
||||
62,
|
||||
30,
|
||||
42,
|
||||
65,
|
||||
34,
|
||||
37,
|
||||
30,
|
||||
67,
|
||||
2,
|
||||
34,
|
||||
77,
|
||||
96,
|
||||
24,
|
||||
89,
|
||||
19,
|
||||
33
|
||||
],
|
||||
"markers": [
|
||||
3,
|
||||
9,
|
||||
17
|
||||
],
|
||||
"qtype": 0,
|
||||
"logits": [
|
||||
-0.4152399003505707,
|
||||
-0.5505563020706177,
|
||||
-0.600245475769043
|
||||
],
|
||||
"act": [
|
||||
0.4727683663368225,
|
||||
0.5272316336631775
|
||||
]
|
||||
},
|
||||
{
|
||||
"input_ids": [
|
||||
4,
|
||||
5,
|
||||
43,
|
||||
95,
|
||||
77,
|
||||
35,
|
||||
61,
|
||||
13,
|
||||
42,
|
||||
24,
|
||||
21,
|
||||
7,
|
||||
11,
|
||||
38,
|
||||
80,
|
||||
54,
|
||||
4,
|
||||
29,
|
||||
5,
|
||||
80,
|
||||
54,
|
||||
78,
|
||||
46
|
||||
],
|
||||
"markers": [
|
||||
4,
|
||||
11
|
||||
],
|
||||
"qtype": 1,
|
||||
"logits": [
|
||||
-0.1313231736421585,
|
||||
-0.3265465199947357
|
||||
],
|
||||
"act": [
|
||||
0.40568628907203674,
|
||||
0.5943136811256409
|
||||
]
|
||||
},
|
||||
{
|
||||
"input_ids": [
|
||||
67,
|
||||
40,
|
||||
83,
|
||||
41,
|
||||
77,
|
||||
96,
|
||||
58,
|
||||
39,
|
||||
76
|
||||
],
|
||||
"markers": [
|
||||
2
|
||||
],
|
||||
"qtype": 2,
|
||||
"logits": [
|
||||
0.534652590751648
|
||||
],
|
||||
"act": [
|
||||
0.4226877987384796,
|
||||
0.5773122310638428
|
||||
]
|
||||
},
|
||||
{
|
||||
"input_ids": [
|
||||
96,
|
||||
32,
|
||||
60,
|
||||
44,
|
||||
35,
|
||||
47,
|
||||
45,
|
||||
51,
|
||||
60,
|
||||
51,
|
||||
73,
|
||||
50,
|
||||
25,
|
||||
79,
|
||||
74,
|
||||
24,
|
||||
83,
|
||||
7,
|
||||
9,
|
||||
91,
|
||||
21,
|
||||
87,
|
||||
29,
|
||||
25,
|
||||
74,
|
||||
18,
|
||||
31,
|
||||
46,
|
||||
46,
|
||||
30,
|
||||
52
|
||||
],
|
||||
"markers": [
|
||||
5,
|
||||
7,
|
||||
12,
|
||||
20,
|
||||
26
|
||||
],
|
||||
"qtype": 0,
|
||||
"logits": [
|
||||
-0.4001395106315613,
|
||||
-0.5266168117523193,
|
||||
-0.44161877036094666,
|
||||
-0.43973028659820557,
|
||||
-0.5682165622711182
|
||||
],
|
||||
"act": [
|
||||
0.46282079815864563,
|
||||
0.537179172039032
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
Binary file not shown.
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"encoder": "tiny-modernbert",
|
||||
"head_layers": 2,
|
||||
"max_len": 64,
|
||||
"head_max_len": 24,
|
||||
"act_costs": {
|
||||
"escalate": 0.5
|
||||
},
|
||||
"amp_dtype": "bf16",
|
||||
"temperature": [
|
||||
1.0,
|
||||
1.0,
|
||||
1.0
|
||||
],
|
||||
"temperature_by_options": {}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
{"version": "1.0", "truncation": null, "padding": null, "added_tokens": [], "normalizer": null, "pre_tokenizer": {"type": "Whitespace"}, "post_processor": null, "decoder": null, "model": {"type": "WordLevel", "vocab": {"[PAD]": 0, "[UNK]": 1, "[CLS]": 2, "[SEP]": 3, "[MASK]": 4, "a": 5, "b": 6}, "unk_token": "[UNK]"}}
|
||||
@@ -0,0 +1,7 @@
|
||||
{
|
||||
"cls_token": "[CLS]",
|
||||
"sep_token": "[SEP]",
|
||||
"mask_token": "[MASK]",
|
||||
"pad_token": "[PAD]",
|
||||
"unk_token": "[UNK]"
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"cls_token": "[CLS]",
|
||||
"mask_token": "[MASK]",
|
||||
"model_input_names": [
|
||||
"input_ids",
|
||||
"attention_mask"
|
||||
],
|
||||
"model_max_length": 8192,
|
||||
"pad_token": "[PAD]",
|
||||
"sep_token": "[SEP]",
|
||||
"tokenizer_class": "PreTrainedTokenizerFast",
|
||||
"unk_token": "[UNK]"
|
||||
}
|
||||
@@ -67,6 +67,22 @@ impl TpLeaderModel {
|
||||
}
|
||||
}
|
||||
|
||||
/// Logits at every position (#96). Only the `qwen3_5` arch has a
|
||||
/// multi-position path; the older dense `qwen3` arch says so rather
|
||||
/// than silently returning the last row.
|
||||
pub fn forward_multi(
|
||||
&mut self,
|
||||
input: &candle_core::Tensor,
|
||||
offset: usize,
|
||||
) -> candle_core::Result<candle_core::Tensor> {
|
||||
match self {
|
||||
TpLeaderModel::Qwen3_5(m) => m.forward_multi(input, offset),
|
||||
TpLeaderModel::Qwen3(_) => Err(candle_core::Error::Msg(
|
||||
"forward_multi: the tp qwen3 arch has no multi-position logits path".into(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Chunked image prefill on rank 0. Only the vision-capable
|
||||
/// `qwen3_5` arch supports it; the dense `qwen3` arch has no tower.
|
||||
pub fn prefill_with_images_chunked(
|
||||
@@ -996,6 +1012,112 @@ impl WorkerPool {
|
||||
}
|
||||
}
|
||||
|
||||
/// Multi-position forward across every rank (#96).
|
||||
///
|
||||
/// Same fan-out / leader-forward / always-drain shape as
|
||||
/// [`Self::generate_step`], and for the same reason: the
|
||||
/// row-parallel `AllReduce`s only complete when every rank has
|
||||
/// issued them, and a leader error must still drain the workers or
|
||||
/// the stale replies poison the next request.
|
||||
///
|
||||
/// Returns one `[vocab]` row per position from the leader's rank-0
|
||||
/// shard, which is the whole answer because `lm_head` is
|
||||
/// replicated.
|
||||
#[cfg(feature = "cuda")]
|
||||
pub async fn generate_step_multi(
|
||||
&mut self,
|
||||
model_id: &str,
|
||||
leader_handle: super::device_worker::TpHandle,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
) -> Result<Vec<Vec<f32>>> {
|
||||
let step_start = std::time::Instant::now();
|
||||
let tokens_len = tokens.len();
|
||||
tracing::debug!(
|
||||
model = %model_id,
|
||||
tokens = tokens_len,
|
||||
offset,
|
||||
"WorkerPool::generate_step_multi: fan-out"
|
||||
);
|
||||
for w in &mut self.workers {
|
||||
w.send_only(&WorkerRequest::GenerateStepMulti {
|
||||
model_id: model_id.to_string(),
|
||||
tokens: tokens.clone(),
|
||||
offset,
|
||||
})
|
||||
.await?;
|
||||
}
|
||||
|
||||
let leader_start = std::time::Instant::now();
|
||||
let timeout = tp_step_timeout();
|
||||
let leader_fut = self
|
||||
.leader_worker
|
||||
.tp_forward_logits_multi(leader_handle, tokens, offset);
|
||||
let leader_result = match tokio::time::timeout(timeout, leader_fut).await {
|
||||
Ok(r) => r,
|
||||
Err(_elapsed) => {
|
||||
self.watchdog_abort_leader_comm(model_id, timeout.as_secs());
|
||||
anyhow::bail!(
|
||||
"tp watchdog: leader multi-position forward exceeded {}s deadline; \
|
||||
aborted wedged NCCL comm — model will auto-recover",
|
||||
timeout.as_secs()
|
||||
);
|
||||
}
|
||||
};
|
||||
let leader_ok = leader_result.is_ok();
|
||||
if !leader_ok {
|
||||
let detail = leader_result
|
||||
.as_ref()
|
||||
.err()
|
||||
.map(|e| format!("{e:#}"))
|
||||
.unwrap_or_default();
|
||||
tracing::warn!(
|
||||
model = %model_id,
|
||||
tokens = tokens_len,
|
||||
offset,
|
||||
leader_ms = leader_start.elapsed().as_millis(),
|
||||
error = %detail,
|
||||
"WorkerPool::generate_step_multi: leader forward failed"
|
||||
);
|
||||
}
|
||||
|
||||
let worker_errors = drain_workers(&mut self.workers, |r| match r {
|
||||
WorkerResponse::GenerateStepOk => Ok(()),
|
||||
WorkerResponse::Error { kind, message } => Err(format!("[{kind}]: {message}")),
|
||||
other => Err(format!("expected GenerateStepOk, got {other:?}")),
|
||||
})
|
||||
.await;
|
||||
tracing::debug!(
|
||||
model = %model_id,
|
||||
errors = worker_errors.len(),
|
||||
total_ms = step_start.elapsed().as_millis(),
|
||||
"WorkerPool::generate_step_multi: workers drained"
|
||||
);
|
||||
|
||||
match leader_result {
|
||||
Ok(values) => {
|
||||
if worker_errors.is_empty() {
|
||||
Ok(values)
|
||||
} else {
|
||||
anyhow::bail!(
|
||||
"GenerateStepMulti: leader succeeded but workers failed: {}",
|
||||
worker_errors.join("; ")
|
||||
)
|
||||
}
|
||||
}
|
||||
Err(e) => {
|
||||
if worker_errors.is_empty() {
|
||||
Err(anyhow::Error::new(e).context("GenerateStepMulti: leader forward failed"))
|
||||
} else {
|
||||
Err(anyhow::Error::new(e).context(format!(
|
||||
"GenerateStepMulti: leader forward failed and workers also failed: {}",
|
||||
worker_errors.join("; ")
|
||||
)))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One lockstep batched decode step across every rank (#98).
|
||||
/// Same fan-out / leader-forward / always-drain shape as
|
||||
/// [`Self::generate_step`]; every rank derives positions + mask
|
||||
|
||||
@@ -78,6 +78,18 @@ pub enum WorkerRequest {
|
||||
/// same op. The leader does *not* receive logits back over RPC; it
|
||||
/// runs its own rank-0 forward in parallel and uses its own logits
|
||||
/// for sampling.
|
||||
/// Multi-position forward (#96): like `GenerateStep`, but every
|
||||
/// rank keeps the logits at *all* `tokens.len()` positions rather
|
||||
/// than the last. The collectives are identical, so the ranks stay
|
||||
/// in lockstep; only the leader's logits are read. Worker replies
|
||||
/// with `GenerateStepOk`.
|
||||
GenerateStepMulti {
|
||||
model_id: String,
|
||||
/// Identical on every rank, as for `GenerateStep`.
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
},
|
||||
|
||||
GenerateStep {
|
||||
model_id: String,
|
||||
/// Input token ids for this step. For prefill, the whole prompt;
|
||||
|
||||
@@ -1881,6 +1881,27 @@ impl TpQwen3_5Model {
|
||||
}
|
||||
}
|
||||
|
||||
/// Additive causal mask for the draft head, `(1, 1, l, l + offset)`.
|
||||
/// The head keeps its own KV, so its mask is built from its own
|
||||
/// position count rather than the target's.
|
||||
fn mtp_causal_mask(
|
||||
l: usize,
|
||||
offset: usize,
|
||||
dtype: DType,
|
||||
dev: &Device,
|
||||
) -> candle_core::Result<Tensor> {
|
||||
let total = l + offset;
|
||||
let mut data = vec![0f32; l * total];
|
||||
for i in 0..l {
|
||||
for j in 0..total {
|
||||
if j > i + offset {
|
||||
data[i * total + j] = f32::NEG_INFINITY;
|
||||
}
|
||||
}
|
||||
}
|
||||
Tensor::from_vec(data, (1, 1, l, total), dev)?.to_dtype(dtype)
|
||||
}
|
||||
|
||||
pub struct TpQwen3_5ForCausalLM {
|
||||
base: TpQwen3_5Model,
|
||||
lm_head: super::tp_linear::MaybeQuantLinear,
|
||||
@@ -1894,6 +1915,23 @@ pub struct TpQwen3_5ForCausalLM {
|
||||
/// `<|image_pad|>` sentinel id (mirrors `Config::image_token_id`);
|
||||
/// the splice target for `forward_with_vision`.
|
||||
image_token_id: Option<u32>,
|
||||
/// The MTP draft head (#96), on the **leader only**.
|
||||
///
|
||||
/// Unsharded, unlike every other weight here, and that is correct
|
||||
/// rather than lazy: after each row-parallel `AllReduce` every rank
|
||||
/// holds the *same* hidden state, and `lm_head` is replicated. So a
|
||||
/// head fed the leader's hidden produces exactly what a sharded one
|
||||
/// would, with no collectives — which also means drafting cannot
|
||||
/// wedge the ranks, since the other ranks are not in it.
|
||||
///
|
||||
/// `None` on the subprocess ranks and on any load that did not ask
|
||||
/// for it. Loaded by [`Self::load_mtp_head`], never by `load`.
|
||||
mtp: Option<crate::harness::arch::qwen3_5::mtp::MtpHead>,
|
||||
/// The last forward's hidden state, `(1, L, hidden)` — what the
|
||||
/// draft head consumes. Kept because the head runs *after* the
|
||||
/// target's step, outside the collective, so it cannot recompute
|
||||
/// it. One `Arc` clone per forward; no copy.
|
||||
last_hidden: Option<Tensor>,
|
||||
}
|
||||
|
||||
/// Load the replicated vision tower from the unsharded `model.visual.*`
|
||||
@@ -1944,6 +1982,8 @@ impl TpQwen3_5ForCausalLM {
|
||||
lm_head,
|
||||
vision,
|
||||
image_token_id,
|
||||
mtp: None,
|
||||
last_hidden: None,
|
||||
};
|
||||
log_construction_complete(cfg, rank, world_size, quant, model.device());
|
||||
Ok(model)
|
||||
@@ -1969,6 +2009,8 @@ impl TpQwen3_5ForCausalLM {
|
||||
lm_head,
|
||||
vision,
|
||||
image_token_id,
|
||||
mtp: None,
|
||||
last_hidden: None,
|
||||
};
|
||||
log_construction_complete(cfg, rank, world_size, quant, model.device());
|
||||
Ok(model)
|
||||
@@ -2003,9 +2045,201 @@ impl TpQwen3_5ForCausalLM {
|
||||
pub fn forward(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
|
||||
let (_, l) = input.dims2()?;
|
||||
let hidden = self.base.forward(input, offset)?;
|
||||
self.last_hidden = Some(hidden.clone());
|
||||
hidden.i((.., l - 1.., ..))?.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
/// Logits at **every** position: `(B, L, vocab)`.
|
||||
///
|
||||
/// The TP mirror of `Qwen3_5ForCausalLM::forward_multi` (#96). A
|
||||
/// speculative verify pass forwards the drafted block in one pass
|
||||
/// and needs the target's own token at each drafted position;
|
||||
/// [`Self::forward`] slices the last one off before the LM head.
|
||||
///
|
||||
/// Every rank runs this, because the row-parallel `o_proj` and the
|
||||
/// MLP's `down_proj` issue `AllReduce` collectives that only
|
||||
/// complete when all ranks arrive. The leader's copy is the one
|
||||
/// whose logits are read — `lm_head` is replicated, so rank 0's
|
||||
/// answer is the whole answer.
|
||||
/// Load the MTP draft head onto this rank, unsharded (#96).
|
||||
///
|
||||
/// Only the leader should call this. See the `mtp` field for why an
|
||||
/// unsharded head is numerically right; the practical consequence
|
||||
/// is that drafting issues no collectives and therefore cannot
|
||||
/// leave the other ranks waiting.
|
||||
pub fn load_mtp_head(&mut self, cfg: &TextConfig, vb: &ShardedVarBuilder) -> Result<()> {
|
||||
let head = crate::harness::arch::qwen3_5::mtp::MtpHead::load(
|
||||
cfg,
|
||||
Arc::clone(&self.base.rotary),
|
||||
vb,
|
||||
)
|
||||
.context("load mtp head on the tp leader")?;
|
||||
self.mtp = Some(head);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Whether this rank carries a draft head.
|
||||
pub fn has_mtp(&self) -> bool {
|
||||
self.mtp.is_some()
|
||||
}
|
||||
|
||||
/// Walk the draft head over a prompt chunk, using the hidden state
|
||||
/// the target's forward just produced.
|
||||
///
|
||||
/// The head trails the target by one token: at position `i` it
|
||||
/// consumes the embedding of token `i+1` with the target's hidden
|
||||
/// state at `i`. `shifted` is therefore the chunk's tokens moved
|
||||
/// one left, and must be one shorter than the retained hidden block
|
||||
/// unless the caller has the next chunk's first token.
|
||||
pub fn mtp_prefill_chunk(&mut self, shifted: &[u32], start_pos: usize) -> Result<()> {
|
||||
let hidden = self
|
||||
.last_hidden
|
||||
.as_ref()
|
||||
.context("mtp_prefill_chunk: no hidden state retained; run the target forward first")?
|
||||
.clone();
|
||||
let l = shifted.len();
|
||||
anyhow::ensure!(l > 0, "mtp_prefill_chunk: empty chunk");
|
||||
let (_, have, _) = hidden.dims3()?;
|
||||
anyhow::ensure!(
|
||||
have >= l,
|
||||
"mtp_prefill_chunk: {l} tokens against {have} retained hidden positions"
|
||||
);
|
||||
let device = self.base.device.clone();
|
||||
let dtype = self.base.dtype;
|
||||
let rotary = Arc::clone(&self.base.rotary);
|
||||
let tokens = Tensor::new(shifted, &device)?.unsqueeze(0)?;
|
||||
let embeds = self.base.embed_tokens.forward(&tokens)?;
|
||||
let hid = hidden.i((.., ..l, ..))?;
|
||||
let mask = mtp_causal_mask(l, start_pos, dtype, &device)?;
|
||||
let positions: Vec<usize> = (start_pos..start_pos + l).collect();
|
||||
let (cos, sin) = rotary.cos_sin_at(&positions)?;
|
||||
let head = self
|
||||
.mtp
|
||||
.as_mut()
|
||||
.context("mtp_prefill_chunk: no draft head on this rank")?;
|
||||
head.forward(&embeds, &hid, Some(&mask), &cos, &sin)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Draft `k` tokens, then put the head's cache back.
|
||||
///
|
||||
/// The head is walked over tokens that may be rejected, so the
|
||||
/// snapshot/restore pair around the chain is what keeps this an
|
||||
/// observation rather than a commitment. Returns the drafted token
|
||||
/// ids, argmax-sampled through the target's own (replicated)
|
||||
/// `lm_head`.
|
||||
pub fn mtp_draft(&mut self, first_token: u32, start_pos: usize, k: usize) -> Result<Vec<u32>> {
|
||||
anyhow::ensure!(k > 0, "mtp_draft: k must be positive");
|
||||
let mut hidden = self
|
||||
.last_hidden
|
||||
.as_ref()
|
||||
.context("mtp_draft: no hidden state retained; run the target forward first")?
|
||||
.clone();
|
||||
let (_, have, _) = hidden.dims3()?;
|
||||
if have > 1 {
|
||||
hidden = hidden.i((.., have - 1.., ..))?;
|
||||
}
|
||||
let device = self.base.device.clone();
|
||||
let rotary = Arc::clone(&self.base.rotary);
|
||||
|
||||
let snap = {
|
||||
let head = self
|
||||
.mtp
|
||||
.as_ref()
|
||||
.context("mtp_draft: no draft head on this rank")?;
|
||||
head.snapshot_kv()?
|
||||
};
|
||||
|
||||
let mut drafted = Vec::with_capacity(k);
|
||||
let mut token = first_token;
|
||||
let mut cur = hidden;
|
||||
let mut err: Option<anyhow::Error> = None;
|
||||
for step in 0..k {
|
||||
let ids = match Tensor::new(&[token], &device).and_then(|t| t.unsqueeze(0)) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
err = Some(e.into());
|
||||
break;
|
||||
}
|
||||
};
|
||||
let out = (|| -> Result<Tensor> {
|
||||
let embeds = self.base.embed_tokens.forward(&ids)?;
|
||||
let (cos, sin) = rotary.cos_sin_at(&[start_pos + step])?;
|
||||
let head = self
|
||||
.mtp
|
||||
.as_mut()
|
||||
.context("mtp_draft: no draft head on this rank")?;
|
||||
head.forward(&embeds, &cur, None, &cos, &sin)
|
||||
})();
|
||||
let out = match out {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
err = Some(e);
|
||||
break;
|
||||
}
|
||||
};
|
||||
let next = match self.lm_head.forward(&out).and_then(|logits| {
|
||||
logits
|
||||
.i((0, 0, ..))?
|
||||
.to_dtype(DType::F32)?
|
||||
.argmax(candle_core::D::Minus1)?
|
||||
.to_scalar::<u32>()
|
||||
}) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
err = Some(e.into());
|
||||
break;
|
||||
}
|
||||
};
|
||||
drafted.push(next);
|
||||
token = next;
|
||||
cur = out;
|
||||
}
|
||||
|
||||
// Restore even when a step failed: a half-walked head is worse
|
||||
// than no draft, because every later round inherits it.
|
||||
if let Some(head) = self.mtp.as_mut() {
|
||||
head.restore_kv(&snap)?;
|
||||
}
|
||||
match err {
|
||||
Some(e) => Err(e),
|
||||
None => Ok(drafted),
|
||||
}
|
||||
}
|
||||
|
||||
/// Walk the head over one token the target actually committed, so
|
||||
/// its cache stays level with the target's.
|
||||
pub fn mtp_advance(&mut self, token: u32, pos: usize) -> Result<()> {
|
||||
let hidden = self
|
||||
.last_hidden
|
||||
.as_ref()
|
||||
.context("mtp_advance: no hidden state retained")?
|
||||
.clone();
|
||||
let (_, have, _) = hidden.dims3()?;
|
||||
let hidden = if have > 1 {
|
||||
hidden.i((.., have - 1.., ..))?
|
||||
} else {
|
||||
hidden
|
||||
};
|
||||
let device = self.base.device.clone();
|
||||
let rotary = Arc::clone(&self.base.rotary);
|
||||
let ids = Tensor::new(&[token], &device)?.unsqueeze(0)?;
|
||||
let embeds = self.base.embed_tokens.forward(&ids)?;
|
||||
let (cos, sin) = rotary.cos_sin_at(&[pos])?;
|
||||
let head = self
|
||||
.mtp
|
||||
.as_mut()
|
||||
.context("mtp_advance: no draft head on this rank")?;
|
||||
head.forward(&embeds, &hidden, None, &cos, &sin)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn forward_multi(&mut self, input: &Tensor, offset: usize) -> candle_core::Result<Tensor> {
|
||||
let hidden = self.base.forward(input, offset)?;
|
||||
self.last_hidden = Some(hidden.clone());
|
||||
hidden.apply(&self.lm_head)
|
||||
}
|
||||
|
||||
/// Lockstep batched decode step (#98): `(B, 1)` input, per-row
|
||||
/// positions, padding mask from
|
||||
/// [`TpQwen3_5Model::batch_decode_mask`]. Returns `(B, 1, vocab)`.
|
||||
|
||||
@@ -52,6 +52,21 @@ impl WorkerModel {
|
||||
}
|
||||
}
|
||||
|
||||
/// Logits at every position (#96). This rank discards them; what
|
||||
/// matters is the identical collective sequence.
|
||||
fn forward_multi(
|
||||
&mut self,
|
||||
input: &candle_core::Tensor,
|
||||
offset: usize,
|
||||
) -> candle_core::Result<candle_core::Tensor> {
|
||||
match self {
|
||||
WorkerModel::Qwen3_5(m) => m.forward_multi(input, offset),
|
||||
WorkerModel::Qwen3(_) => Err(candle_core::Error::Msg(
|
||||
"forward_multi: the tp qwen3 arch has no multi-position logits path".into(),
|
||||
)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Chunked image prefill on this rank. Only the vision-capable
|
||||
/// `qwen3_5` arch has a replicated tower; the dense `qwen3` arch
|
||||
/// errors. The returned logits are discarded by the caller (the
|
||||
@@ -269,6 +284,11 @@ impl WorkerState {
|
||||
tokens,
|
||||
offset,
|
||||
} => self.handle_generate_step(&model_id, tokens, offset),
|
||||
WorkerRequest::GenerateStepMulti {
|
||||
model_id,
|
||||
tokens,
|
||||
offset,
|
||||
} => self.handle_generate_step_multi(&model_id, tokens, offset),
|
||||
WorkerRequest::GenerateStepWithImages {
|
||||
model_id,
|
||||
tokens,
|
||||
@@ -594,6 +614,80 @@ impl WorkerState {
|
||||
}
|
||||
}
|
||||
|
||||
/// Multi-position forward (#96). Identical collectives to
|
||||
/// `handle_generate_step` — the only difference is that the model
|
||||
/// keeps every position's logits instead of the last. This rank
|
||||
/// discards them either way; what matters is that it issues the
|
||||
/// same `AllReduce` sequence, or the leader's forward never
|
||||
/// completes.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn handle_generate_step_multi(
|
||||
&mut self,
|
||||
model_id: &str,
|
||||
tokens: Vec<u32>,
|
||||
offset: usize,
|
||||
) -> WorkerResponse {
|
||||
use candle_core::Tensor;
|
||||
|
||||
let Some(model) = self.models.get_mut(model_id) else {
|
||||
return WorkerResponse::Error {
|
||||
kind: "model_not_loaded".into(),
|
||||
message: format!("model '{model_id}' not loaded on rank {}", self.config.rank),
|
||||
};
|
||||
};
|
||||
let device = model.device().clone();
|
||||
let input = match Tensor::new(tokens.as_slice(), &device).and_then(|t| t.unsqueeze(0)) {
|
||||
Ok(t) => t,
|
||||
Err(e) => {
|
||||
return WorkerResponse::Error {
|
||||
kind: "forward_failed".into(),
|
||||
message: format!("build input tensor: {e}"),
|
||||
};
|
||||
}
|
||||
};
|
||||
let start = std::time::Instant::now();
|
||||
tracing::debug!(
|
||||
rank = self.config.rank,
|
||||
model = %model_id,
|
||||
tokens = tokens.len(),
|
||||
offset,
|
||||
"worker GenerateStepMulti: forward starting"
|
||||
);
|
||||
if let Err(e) = model.forward_multi(&input, offset) {
|
||||
tracing::warn!(
|
||||
rank = self.config.rank,
|
||||
model = %model_id,
|
||||
elapsed_ms = start.elapsed().as_millis(),
|
||||
error = %e,
|
||||
"worker GenerateStepMulti: forward failed"
|
||||
);
|
||||
return WorkerResponse::Error {
|
||||
kind: "forward_failed".into(),
|
||||
message: format!("TP multi-position forward: {e}"),
|
||||
};
|
||||
}
|
||||
tracing::debug!(
|
||||
rank = self.config.rank,
|
||||
model = %model_id,
|
||||
elapsed_ms = start.elapsed().as_millis(),
|
||||
"worker GenerateStepMulti: forward done"
|
||||
);
|
||||
WorkerResponse::GenerateStepOk
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
fn handle_generate_step_multi(
|
||||
&mut self,
|
||||
_model_id: &str,
|
||||
_tokens: Vec<u32>,
|
||||
_offset: usize,
|
||||
) -> WorkerResponse {
|
||||
WorkerResponse::Error {
|
||||
kind: "cuda_feature_not_enabled".into(),
|
||||
message: "GenerateStepMulti requires --features cuda".into(),
|
||||
}
|
||||
}
|
||||
|
||||
/// One lockstep batched decode step on this rank (#98). Mirrors
|
||||
/// the leader's `forward_logits_batch`: derives per-row positions
|
||||
/// and the padding mask locally from the broadcast geometry, runs
|
||||
|
||||
@@ -158,6 +158,7 @@ fn preflight_kind(err: &PreflightError) -> &'static str {
|
||||
PreflightError::EmptyRepo { .. } => "empty_repo",
|
||||
PreflightError::TpRequiresSafetensors { .. } => "tp_requires_safetensors",
|
||||
PreflightError::QuantNotFound { .. } => "quant_not_found",
|
||||
PreflightError::Unsupported { .. } => "unsupported_load_option",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -28,6 +28,9 @@ region = "eu-central"
|
||||
[aliases]
|
||||
"helexa/small" = "Qwen/Qwen3-8B"
|
||||
"helexa/balanced" = "Qwen/Qwen3.8-27B"
|
||||
# System-One decisions (#333). Also the model /v1/systemone falls back
|
||||
# to when a request names none or names a service no operator serves.
|
||||
"helexa/one" = "convaiinnovations/laya"
|
||||
|
||||
# hanzalova operator cortex, plaintext over the WireGuard mesh (the mesh
|
||||
# is the trust boundary; no TLS pin needed for http://).
|
||||
|
||||
+26
@@ -167,6 +167,28 @@ min_device_vram_mb = 4000
|
||||
cost.input = 0.015
|
||||
cost.output = 0.06
|
||||
|
||||
# System-One decisions (#333): Laya, a 421M non-autoregressive decision
|
||||
# model (ModernBERT encoder + option-scoring head) served on
|
||||
# /v1/systemone. One repo, three checkpoints -- English at the root,
|
||||
# multilingual/ and typed-decisions/ beside it -- all resident (~2.3 GB
|
||||
# bf16) so neuron can route each request by language. Not a text model:
|
||||
# chat endpoints answer it with wrong_modality.
|
||||
#
|
||||
# Resident on quadbrat, replacing the Qwen3-1.7B pre-warm (operator
|
||||
# decision 2026-09-29, #341); pinned there so a cold-loaded helexa/small
|
||||
# can share the card but never evict it. Metered like TypeSafe Jev,
|
||||
# input tokens only (every question re-encodes the state, and each
|
||||
# encoded token counts), at Jev's published $0.042 per million.
|
||||
[[models]]
|
||||
id = "convaiinnovations/laya"
|
||||
harness = "candle"
|
||||
vram_mb = 3500
|
||||
min_devices = 1
|
||||
min_device_vram_mb = 4000
|
||||
pinned_on = ["quadbrat"]
|
||||
cost.input = 0.042
|
||||
cost.output = 0.0
|
||||
|
||||
# -- Tier aliases ------------------------------------------------------------
|
||||
# Clients can request inference against an alias and cortex transparently
|
||||
# routes to the concrete model id below. Useful so client code can pick a
|
||||
@@ -177,6 +199,10 @@ cost.output = 0.06
|
||||
"helexa/balanced" = "Qwen/Qwen3-8B"
|
||||
"helexa/large" = "Qwen/Qwen3.8-27B"
|
||||
"helexa/image" = "Tongyi-MAI/Z-Image-Turbo"
|
||||
# The System-One decision class: the default /v1/systemone resolves to
|
||||
# for a request with no model, or one naming a service this fleet does
|
||||
# not serve (a Jev client's own model name).
|
||||
"helexa/one" = "convaiinnovations/laya"
|
||||
|
||||
# Text-to-image (#197/#203): Z-Image-Turbo, 6B S3-DiT + Qwen3-4B text
|
||||
# encoder (runs on host CPU by default — see [harness.candle.image] in
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Generate a tiny Laya parity fixture from the upstream reference (#336).
|
||||
|
||||
Builds the `laya` SDK's own `DecisionModel` (laya/common.py) around a
|
||||
tiny ModernBERT, then records its outputs for a right-padded batch.
|
||||
|
||||
Writes, into --out:
|
||||
rl_agent_config.json the decision-model config, as neuron parses it
|
||||
encoder/config.json the ModernBERT config, saved by transformers
|
||||
model.safetensors the reference's own weights, float32
|
||||
tokenizer.json a placeholder word-level tokenizer (the fixture
|
||||
is scored on token ids; loading needs a file)
|
||||
expected.json the batch (ids, markers, question types) and the
|
||||
reference's per-row scorer logits and act softmax
|
||||
|
||||
The shape is chosen so every mechanism is exercised at toy size:
|
||||
|
||||
- four encoder layers with a global layer on each side of two sliding
|
||||
ones, and a sliding window (local_attention = 8, so +-4 positions)
|
||||
far shorter than the 40-token row, so a window bug changes the result;
|
||||
- two different rotary bases for global and sliding layers, so using
|
||||
the wrong one for either kind is visible;
|
||||
- rows of different lengths, so padding is present and must be masked;
|
||||
- a one-option row, which takes the reference's single-marker branch
|
||||
of the act features;
|
||||
- every question type, so the type embedding is indexed three ways.
|
||||
|
||||
Weights are initialised well above transformers' default scale so
|
||||
attention is sharp rather than near-uniform, and norm weights and biases
|
||||
are perturbed away from their ones/zeros init: at init a missing norm
|
||||
and a present one are indistinguishable.
|
||||
|
||||
Usage, in a venv holding the `laya` SDK and CPU torch:
|
||||
|
||||
python script/generate_laya_fixture.py \
|
||||
--out crates/neuron/src/harness/testdata/laya/tiny
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
import torch
|
||||
from safetensors.torch import save_file
|
||||
from transformers import ModernBertConfig, AutoModel
|
||||
|
||||
from laya.common import DecisionModel
|
||||
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--out", required=True)
|
||||
ap.add_argument("--seed", type=int, default=7)
|
||||
args = ap.parse_args()
|
||||
torch.manual_seed(args.seed)
|
||||
os.makedirs(os.path.join(args.out, "encoder"), exist_ok=True)
|
||||
|
||||
ecfg = ModernBertConfig(
|
||||
vocab_size=97,
|
||||
hidden_size=32,
|
||||
num_hidden_layers=4,
|
||||
num_attention_heads=2,
|
||||
intermediate_size=48,
|
||||
max_position_embeddings=128,
|
||||
local_attention=8,
|
||||
global_attn_every_n_layers=3,
|
||||
pad_token_id=0,
|
||||
bos_token_id=1,
|
||||
eos_token_id=2,
|
||||
cls_token_id=1,
|
||||
sep_token_id=2,
|
||||
norm_eps=1e-5,
|
||||
hidden_activation="gelu",
|
||||
# Well above the 0.02 default: at 0.02 attention is near-uniform and
|
||||
# a window or rotary bug moves the logits only in the sixth decimal.
|
||||
initializer_range=0.3,
|
||||
rope_parameters={
|
||||
"full_attention": {"rope_theta": 1000.0, "rope_type": "default"},
|
||||
"sliding_attention": {"rope_theta": 300.0, "rope_type": "default"},
|
||||
},
|
||||
)
|
||||
encoder = AutoModel.from_config(ecfg, attn_implementation="sdpa")
|
||||
model = DecisionModel(encoder, head_layers=2, n_act=2).eval()
|
||||
|
||||
with torch.no_grad():
|
||||
for name, p in model.named_parameters():
|
||||
if "norm" in name or name.startswith("scorer.0"):
|
||||
p.add_(0.3 * torch.randn_like(p))
|
||||
elif name.startswith("type_emb"):
|
||||
p.mul_(3.0)
|
||||
|
||||
ecfg.save_pretrained(os.path.join(args.out, "encoder"))
|
||||
rl_cfg = {
|
||||
"encoder": "tiny-modernbert",
|
||||
"head_layers": 2,
|
||||
"max_len": 64,
|
||||
"head_max_len": 24,
|
||||
"act_costs": {"escalate": 0.5},
|
||||
"amp_dtype": "bf16",
|
||||
"temperature": [1.0, 1.0, 1.0],
|
||||
"temperature_by_options": {},
|
||||
}
|
||||
with open(os.path.join(args.out, "rl_agent_config.json"), "w") as f:
|
||||
json.dump(rl_cfg, f, indent=1)
|
||||
# A placeholder: the fixture is scored on token ids, but a checkpoint
|
||||
# directory is loaded with its tokenizer, so one must exist. It carries
|
||||
# the special tokens a decision tokenizer is looked up by, with [PAD] at
|
||||
# the encoder's pad_token_id (0).
|
||||
with open(os.path.join(args.out, "tokenizer.json"), "w") as f:
|
||||
json.dump({"version": "1.0", "truncation": None, "padding": None, "added_tokens": [],
|
||||
"normalizer": None, "pre_tokenizer": {"type": "Whitespace"},
|
||||
"post_processor": None, "decoder": None,
|
||||
"model": {"type": "WordLevel",
|
||||
"vocab": {"[PAD]": 0, "[UNK]": 1, "[CLS]": 2, "[SEP]": 3, "[MASK]": 4,
|
||||
"a": 5, "b": 6},
|
||||
"unk_token": "[UNK]"}}, f)
|
||||
with open(os.path.join(args.out, "tokenizer_config.json"), "w") as f:
|
||||
json.dump({"cls_token": "[CLS]", "sep_token": "[SEP]", "mask_token": "[MASK]",
|
||||
"pad_token": "[PAD]", "unk_token": "[UNK]"}, f, indent=1)
|
||||
|
||||
state = {k: v.detach().contiguous().float() for k, v in model.state_dict().items()}
|
||||
save_file(state, os.path.join(args.out, "model.safetensors"))
|
||||
|
||||
# (length, marker positions, qtype). Ids avoid 0 (pad) inside a row.
|
||||
rows = [
|
||||
(40, [3, 9, 17], 0),
|
||||
(23, [4, 11], 1),
|
||||
(9, [2], 2),
|
||||
(31, [5, 7, 12, 20, 26], 0),
|
||||
]
|
||||
seqs = []
|
||||
for n, markers, qtype in rows:
|
||||
ids = torch.randint(1, ecfg.vocab_size, (n,)).tolist()
|
||||
seqs.append({"input_ids": ids, "markers": markers, "qtype": qtype})
|
||||
|
||||
b = len(seqs)
|
||||
L = max(len(s["input_ids"]) for s in seqs)
|
||||
K = max(len(s["markers"]) for s in seqs)
|
||||
ids = torch.zeros((b, L), dtype=torch.long)
|
||||
att = torch.zeros((b, L), dtype=torch.long)
|
||||
mpos = torch.zeros((b, K), dtype=torch.long)
|
||||
mmask = torch.zeros((b, K), dtype=torch.bool)
|
||||
for i, s in enumerate(seqs):
|
||||
n = len(s["input_ids"])
|
||||
ids[i, :n] = torch.tensor(s["input_ids"])
|
||||
att[i, :n] = 1
|
||||
k = len(s["markers"])
|
||||
mpos[i, :k] = torch.tensor(s["markers"])
|
||||
mmask[i, :k] = True
|
||||
qtype = torch.tensor([s["qtype"] for s in seqs])
|
||||
|
||||
with torch.no_grad():
|
||||
logits, act = model(ids, att, mpos, mmask, qtype)
|
||||
act = torch.softmax(act.float(), -1)
|
||||
|
||||
for i, s in enumerate(seqs):
|
||||
k = len(s["markers"])
|
||||
s["logits"] = logits[i, :k].tolist()
|
||||
s["act"] = act[i].tolist()
|
||||
|
||||
with open(os.path.join(args.out, "expected.json"), "w") as f:
|
||||
json.dump({"seed": args.seed, "seqs": seqs}, f, indent=1)
|
||||
print("wrote", args.out, "params:", sum(v.numel() for v in state.values()))
|
||||
@@ -0,0 +1,361 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Record Laya's language/script detector over a broad corpus (#339).
|
||||
|
||||
neuron routes each decision request to the English or multilingual Laya
|
||||
checkpoint with a port of `laya.lang.analyse`. That function is a pile of
|
||||
heuristics sitting on Python's Unicode semantics -- `str.isalpha`, `re`'s
|
||||
`\\w`, `unicodedata.combining`, `str.lower`, `str.split` -- and a port that
|
||||
reads the same can still route differently. This records the reference's
|
||||
answer over a corpus built to reach every rule and every place those
|
||||
semantics could diverge:
|
||||
|
||||
- every script range, real sentences and generated letter runs
|
||||
- each function-word language, accented and ASCII-stripped, long and short
|
||||
- English with loanwords, names, symbols and IPA; code, URLs, e-mails,
|
||||
identifiers, acronyms and slash compounds
|
||||
- mixed-script and mixed-language states, multi-line and structured
|
||||
- the 4000-character budgets, the nesting depth limit, non-string leaves
|
||||
- combining marks, decomposed text, case-mapping oddities (İ, ß, final
|
||||
sigma, titlecase digraphs), fullwidth forms, numerals, emoji, and the
|
||||
whitespace Python and Rust disagree about
|
||||
- a seeded fuzz section mixing all of the above
|
||||
|
||||
Every mapping in the corpus has its keys in sorted order, so a consumer
|
||||
whose JSON maps iterate in key order (serde_json without `preserve_order`)
|
||||
sees the same leaf order Python does.
|
||||
|
||||
Characters are limited to those assigned in the Unicode version of the
|
||||
Python that records the corpus, so a newer Unicode table on the consuming
|
||||
side cannot disagree about a code point Python considers unassigned.
|
||||
|
||||
Usage (any Python 3.10+ with the laya SDK importable; no model, no network):
|
||||
|
||||
python script/laya-lang-reference.py \\
|
||||
--out crates/neuron/src/harness/testdata/laya/lang_corpus.json
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import random
|
||||
import sys
|
||||
import unicodedata
|
||||
|
||||
import laya
|
||||
from laya.lang import _WORD, _iter_text, analyse
|
||||
|
||||
SENTENCES = {
|
||||
"en": "Please refund the duplicate charge, I was billed twice this month and it is not fair.",
|
||||
"en_short": "refund please",
|
||||
"en_loanword": "I would like a café table and my résumé is attached for the role.",
|
||||
"en_name": "José and Zoë have the tickets that were sent to you from the office.",
|
||||
"en_two_loanwords": "The naïve café owner sent the résumé to you and it was not what we wanted.",
|
||||
"en_symbols": "Set α to 0.05 and β to 0.9, then run the job again with the new values.",
|
||||
"en_russian_name": "The report from Дмитрий Петрович Савицкий was sent to the team yesterday.",
|
||||
"en_ipa": "Vladimir is pronounced [vlɐˈdʲimʲɪr] in Russian and that is how we say it.",
|
||||
"en_stress": "The name Влади́мир has a stress mark on the second syllable for you.",
|
||||
"fr": "Bonjour, je voudrais annuler ma commande car elle est arrivée très en retard.",
|
||||
"fr_short": "merci beaucoup",
|
||||
"de": "Die Rechnung ist falsch, bitte korrigieren Sie den Betrag für diesen Monat.",
|
||||
"de_short": "Danke schön",
|
||||
"de_eszett": "Ich möchte die Straße und die Größe meiner Bestellung ändern, bitte.",
|
||||
"es": "Me han cobrado dos veces este mes y quiero que me devuelvan el dinero ya.",
|
||||
"pt": "Não consigo acessar a minha conta, já tentei redefinir a senha várias vezes.",
|
||||
"pt_br": "Voce pode me mandar a nota fiscal? Nao recebi nada ate agora.",
|
||||
"pt_ticket": "Deu erro 500 no endpoint de login depois do update de ontem.",
|
||||
"it": "Vorrei sapere quando arriverà il mio ordine, è già in ritardo di una settimana.",
|
||||
"nl": "Het pakket is niet aangekomen en ik wil graag een terugbetaling voor deze bestelling.",
|
||||
"ro": "Vreau să anulez comanda pentru că produsul este deteriorat și nu funcționează.",
|
||||
"bn": "ami taka ferot chai, amar order ekhono ashe nai, ki korbo bolun",
|
||||
"az": "Mən sifarişi ləğv etmək istəyirəm, çünki məhsul çox gec gəldi.",
|
||||
"tr": "Siparişimi iptal etmek istiyorum çünkü ürün çok geç geldi ve kırık.",
|
||||
"pl": "Chciałbym zwrócić pieniądze, ponieważ paczka dotarła uszkodzona.",
|
||||
"cs": "Chtěl bych vrátit peníze, protože zásilka dorazila poškozená.",
|
||||
"hu": "Szeretném visszakapni a pénzemet, mert a csomag sérülten érkezett.",
|
||||
"vi": "Tôi muốn được hoàn tiền vì gói hàng đã đến bị hư hỏng.",
|
||||
"sv": "Jag vill ha pengarna tillbaka eftersom paketet kom fram skadat.",
|
||||
"da": "Jeg vil gerne have pengene tilbage, fordi pakken kom i stykker.",
|
||||
"sw": "Nataka kurudishiwa pesa zangu kwa sababu kifurushi kiliharibika.",
|
||||
"la_shared_only": "la la e e o o un un",
|
||||
"el": "Θέλω να ακυρώσω την παραγγελία μου γιατί ήρθε πολύ αργά.",
|
||||
"el_upper_sigma": "ΘΕΛΩ ΤΗΝ ΕΠΙΣΤΡΟΦΗ ΧΡΗΜΑΤΩΝ ΤΟΥ ΛΟΓΑΡΙΑΣΜΟΥ ΣΑΣ",
|
||||
"ru": "Мне дважды списали деньги в этом месяце, верните, пожалуйста.",
|
||||
"uk": "Мені двічі списали гроші цього місяця, поверніть, будь ласка.",
|
||||
"hy": "Ես ուզում եմ չեղարկել պատվերս, քանի որ այն ուշ հասավ։",
|
||||
"he": "אני רוצה לבטל את ההזמנה שלי כי היא הגיעה באיחור.",
|
||||
"he_niqqud": "שָׁלוֹם, אֲנִי רוֹצֶה לְבַטֵּל אֶת הַהַזְמָנָה",
|
||||
"ar": "تم خصم المبلغ مرتين هذا الشهر، أريد استرداد أموالي من فضلكم.",
|
||||
"fa": "میخواهم سفارشم را لغو کنم چون خیلی دیر رسید.",
|
||||
"hi": "मुझे इस महीने दो बार शुल्क लिया गया, मुझे पैसे वापस चाहिए।",
|
||||
"bn_script": "এই মাসে আমার কাছ থেকে দুবার টাকা কাটা হয়েছে, টাকা ফেরত চাই।",
|
||||
"pa": "ਮੈਨੂੰ ਇਸ ਮਹੀਨੇ ਦੋ ਵਾਰ ਚਾਰਜ ਕੀਤਾ ਗਿਆ, ਪੈਸੇ ਵਾਪਸ ਚਾਹੀਦੇ ਹਨ।",
|
||||
"gu": "મને આ મહિને બે વાર ચાર્જ કરવામાં આવ્યો, પૈસા પાછા જોઈએ છે.",
|
||||
"or": "ମୋତେ ଏହି ମାସରେ ଦୁଇଥର ଚାର୍ଜ କରାଯାଇଛି।",
|
||||
"ta": "இந்த மாதம் எனக்கு இரண்டு முறை கட்டணம் வசூலிக்கப்பட்டது.",
|
||||
"te": "ఈ నెలలో నాకు రెండుసార్లు ఛార్జ్ చేశారు, డబ్బు తిరిగి ఇవ్వండి.",
|
||||
"kn": "ಈ ತಿಂಗಳು ನನಗೆ ಎರಡು ಬಾರಿ ಶುಲ್ಕ ವಿಧಿಸಲಾಗಿದೆ.",
|
||||
"ml": "ഈ മാസം എനിക്ക് രണ്ടുതവണ ചാർജ് ഈടാക്കി.",
|
||||
"si": "මෙම මාසයේ මට දෙවරක් ගාස්තු අය කළා.",
|
||||
"th": "เดือนนี้ฉันถูกเรียกเก็บเงินสองครั้ง ขอเงินคืนด้วยค่ะ",
|
||||
"lo": "ເດືອນນີ້ຂ້ອຍຖືກເກັບເງິນສອງເທື່ອ",
|
||||
"bo": "ང་ལ་ཟླ་བ་འདིར་ཐེངས་གཉིས་རིན་བསྡུས་སོང་།",
|
||||
"my": "ဒီလမှာ ကျွန်တော့်ကို နှစ်ကြိမ် ငွေတောင်းခံခဲ့တယ်။",
|
||||
"ka": "ამ თვეში ორჯერ ჩამომეჭრა თანხა, გთხოვთ დამიბრუნოთ.",
|
||||
"am": "በዚህ ወር ሁለት ጊዜ ክፍያ ተቀንሶብኛል፣ እባክዎ ይመልሱልኝ።",
|
||||
"km": "ខ្ញុំត្រូវបានគិតប្រាក់ពីរដងក្នុងខែនេះ ខ្ញុំចង់បានប្រាក់វិញ",
|
||||
"ko": "이번 달에 두 번 청구되었습니다. 환불해 주세요.",
|
||||
"ja": "今月二回請求されました。返金をお願いします。",
|
||||
"ja_kana": "ありがとうございます。キャンセルしたいです。",
|
||||
"zh": "这个月我被扣了两次钱,我要退款,请尽快处理。",
|
||||
"zh_trad": "這個月我被扣了兩次錢,我要退款。",
|
||||
"yue_ext_b": "𠮷野家の𩸽を食べたい",
|
||||
"mn_cyr": "Би захиалгаа цуцлахыг хүсч байна.",
|
||||
"bopomofo": "ㄅㄆㄇㄈ ㄉㄊㄋㄌ",
|
||||
"halfwidth_kana": "キャンセルシタイデス",
|
||||
"cherokee": "ᏣᎳᎩ ᎦᏬᏂᎯᏍᏗ",
|
||||
}
|
||||
|
||||
EXTRA = [
|
||||
"", " ", " \n\t ", "12345 67890", "3.14 2.71 1.41", "!!! ??? ...", "ok", "a", "no",
|
||||
"🙂🙂🙂 👍", "refund 🙂 please 👍 thanks", "café", "résumé naïve café",
|
||||
"ß", "Straße", "İstanbul İzmir", "ıi İI", "Džemal Ljubljana Nj", "ⓐⓑⓒ ⒶⒷⒸ", "𝐀𝐁𝐂 𝐚𝐛𝐜 test",
|
||||
"Ⅰ Ⅱ Ⅲ ⅳ ⅴ", "x² y³ ½ ¾", "٣٤٥ ١٢", "१२३ ४५६", "FULLWIDTH text here",
|
||||
"ʰʲʷ ˈˌ", "ɐɑɒ ʃʒ θð", "ƀƁƂ ǝǞ ȀȁȂ", "Ȿ ɀ", "ḀḁḂḃ ẞ ỳỹ",
|
||||
"a\x1cb\x1dc\x1ed\x1ff", "word\u00a0word\u00a0word\u00a0word", "一\u3000二\u3000三",
|
||||
"tab\tseparated\tvalues\there", "zero\u200bwidth\u200bjoiner", "soft\u00adhyphen",
|
||||
"e\u0301 e\u0300 a\u0308 o\u0303", "Tôi muốn được hoàn tiền",
|
||||
unicodedata.normalize("NFD", "Tôi muốn được hoàn tiền vì gói hàng đã đến bị hư hỏng."),
|
||||
unicodedata.normalize("NFD", "Não consigo acessar a minha conta, já tentei várias vezes."),
|
||||
unicodedata.normalize("NFD", "Vreau să anulez comanda pentru că produsul este deteriorat."),
|
||||
"ΟΔΟΣ ΟΔΟΣ ΟΔΟΣ", "Σ", "ΑΣ ΑΣ", "ὈΔΥΣΣΕΎΣ",
|
||||
]
|
||||
|
||||
CODE_AND_IDENTIFIERS = [
|
||||
"see https://github.com/acme/os.path and user@acme.com for details",
|
||||
"email me at joao.silva@empresa.com.br or maria@exemplo.pt please",
|
||||
"version v1.2.3 of the lib, see U.S.A. rules and e.g. the docs",
|
||||
"x = round(el, 2); return non_english",
|
||||
"import os.path\nfrom foo import bar",
|
||||
"def f(x):\n return {'la': 1, 'le': 2}",
|
||||
"Nav/Com OS/2 C:\\DOS\\mode ESA/UN are names not sentences at all",
|
||||
"MON LA EST COM DES are teams, states and radio bands in the listing",
|
||||
"COM COM COM COM radio listing for the English region",
|
||||
"O CLIENTE ESTÁ MUITO IRRITADO COM O ATRASO DO PEDIDO",
|
||||
"call foo(bar) then baz(qux) and see what the output of the call is",
|
||||
"Deu erro (500) no login depois do update de ontem",
|
||||
"a.b.c.d e@f g.h-i j_k.l m-n.o",
|
||||
"-leading.dash and trailing.dot. and @handle and .hidden files here",
|
||||
"..... @@@@ .@. a.@b",
|
||||
"x" * 300 + ".com",
|
||||
"path/to/file.py line 42 in <module> raise ValueError('não funciona')",
|
||||
]
|
||||
|
||||
MULTILINE = [
|
||||
"Hello team,\nPlease see below.\nNão consigo acessar a minha conta, já tentei redefinir a senha.\nThanks",
|
||||
"Traceback (most recent call last):\n File \"app.py\", line 3, in <module>\n x = foo()\nO sistema não "
|
||||
"está funcionando desde ontem e preciso de ajuda urgente com isso",
|
||||
"Line one in English is here.\nLa fattura è sbagliata, per favore correggetela subito.",
|
||||
"Short line\n\n\nAnother short\nDie Rechnung ist falsch und ich bitte um eine Korrektur heute.",
|
||||
"English only line one.\nEnglish only line two.\nAnd the third line of English text.",
|
||||
"José\nMaría\nFrançois\nplease refund the order that was sent to them",
|
||||
"Ticket #1\nこんにちは、注文をキャンセルしたいです\nThanks for the help with this",
|
||||
# Targeted: each of the next cases flips if one rule is ported wrong.
|
||||
# All-caps function words in mixed-case text are acronyms, not French.
|
||||
"Please refund the duplicate charge, I was billed twice this month and it is not fair.\n"
|
||||
"Scores: MON LA EST DES COM LES ET tonight",
|
||||
# U+001C splits tokens for str.split(), so the identifier is dropped
|
||||
# alone rather than taking the Portuguese words with it.
|
||||
"Please refund the duplicate charge, I was billed twice this month.\n"
|
||||
"Não\x1cconsigo\x1cacessar\x1cminha\x1cconta\x1cjá\x1cx.y",
|
||||
# A combining stress mark must not split a capitalised name into a
|
||||
# lowercase "word".
|
||||
"The customer name on the invoice is Влади́мир "
|
||||
"Петро́вич and we need it fixed today",
|
||||
# One loanword: English function words rescue the marginal diacritic rate.
|
||||
"We will meet at the café for you",
|
||||
"I think the naïve plan is fine for all of us here",
|
||||
]
|
||||
|
||||
|
||||
def budget_cases():
|
||||
en = "Please refund the duplicate charge on my card for this month. "
|
||||
pt = "Não consigo acessar a minha conta, já tentei redefinir a senha várias vezes. "
|
||||
out = []
|
||||
for n in (3990, 3999, 4000, 4001, 4010, 8000):
|
||||
out.append((en * (n // len(en) + 1))[:n])
|
||||
out.append((en * 70)[:3995] + " " + pt)
|
||||
out.append((en * 70) + pt)
|
||||
out.append((en * 100) + "\n" + pt)
|
||||
out.append("x" * 3999 + "\n" + pt)
|
||||
out.append("x" * 4000 + "\n" + pt)
|
||||
out.append(("é" * 5000))
|
||||
out.append(("ab " * 1400) + pt)
|
||||
return out
|
||||
|
||||
|
||||
def structured_cases():
|
||||
en = SENTENCES["en"]
|
||||
pt = SENTENCES["pt"]
|
||||
de = SENTENCES["de"]
|
||||
ja = SENTENCES["ja"]
|
||||
deep = "Não consigo acessar a minha conta, já tentei redefinir a senha."
|
||||
nested = deep
|
||||
deeps = []
|
||||
for depth in range(9):
|
||||
nested = [nested]
|
||||
deeps.append(nested)
|
||||
return [
|
||||
{"body": pt, "subject": "Refund request"},
|
||||
{"body": en, "note": pt},
|
||||
{"a_note": en * 60, "b_message": de},
|
||||
{"a_note": en * 60, "b_message": "Hallo, bitte."},
|
||||
{"customer": "José", "message": en},
|
||||
{"name": "Дмитрий Петрович", "text": en},
|
||||
{"message": ja, "subject": "Order"},
|
||||
{"count": 3, "flag": True, "missing": None, "text": en},
|
||||
{"count": 3, "flag": False, "text": pt},
|
||||
{"items": [en, pt, {"deep": de}], "meta": {"lang": "?"}},
|
||||
{"conversation": [{"content": en, "role": "user"}, {"content": pt, "role": "user"}]},
|
||||
{"a": {"b": {"c": {"d": {"e": {"f": {"g": pt}}}}}}},
|
||||
{"a": {"b": {"c": {"d": {"e": {"f": pt}}}}}},
|
||||
{"code": "x = round(el, 2)\nreturn os.path", "text": en},
|
||||
{"body": "Traceback:\n File x.py\nO sistema não está funcionando desde ontem e preciso de ajuda", "id": 7},
|
||||
{"a": "José", "b": "María", "c": "François", "d": en},
|
||||
{"a": en * 70, "b": "Vreau să anulez comanda pentru că produsul este deteriorat."},
|
||||
{"a": en * 70, "b": "Мне дважды списали деньги в этом месяце, верните."},
|
||||
{"a": "ok", "b": "no", "c": "yes"},
|
||||
{},
|
||||
[],
|
||||
[en, pt],
|
||||
[pt, en],
|
||||
[en, en, en, "Die Rechnung ist falsch, bitte korrigieren."],
|
||||
["hello", "bonjour", "hola"],
|
||||
[1, 2, 3],
|
||||
[None, True, 2.5],
|
||||
[[["deep", ["deeper", [pt]]]]],
|
||||
*deeps,
|
||||
12345,
|
||||
3.5,
|
||||
True,
|
||||
None,
|
||||
]
|
||||
|
||||
|
||||
def fuzz_cases(n, seed):
|
||||
rng = random.Random(seed)
|
||||
pools = [v for v in SENTENCES.values()] + EXTRA + CODE_AND_IDENTIFIERS
|
||||
words = []
|
||||
for s in pools:
|
||||
words.extend(s.split())
|
||||
specials = list("\n\t .,;:=()[]{}/\\@_-'\"") + ["\x1c", "\u00a0", "\u3000", "\u0301", "\u0308",
|
||||
"\u200b", "İ", "ı", "ß", "Σ", "Dž", "Ⅲ", "²",
|
||||
"🙂", "1", "A", "a", "ʰ", "ɐ", "é", "ñ"]
|
||||
out = []
|
||||
for _ in range(n):
|
||||
k = rng.randint(1, 40)
|
||||
parts = []
|
||||
for _ in range(k):
|
||||
r = rng.random()
|
||||
if r < 0.75:
|
||||
parts.append(rng.choice(words))
|
||||
else:
|
||||
parts.append(rng.choice(specials))
|
||||
sep = rng.choice([" ", " ", " ", "\n", ""])
|
||||
out.append(sep.join(parts))
|
||||
return out
|
||||
|
||||
|
||||
def generated_script_runs():
|
||||
"""A run of assigned letters from every script range the detector names."""
|
||||
from laya.lang import _SCRIPT_RANGES
|
||||
out = []
|
||||
for name, ranges in _SCRIPT_RANGES:
|
||||
for lo, hi in ranges:
|
||||
letters = [chr(cp) for cp in range(lo, hi + 1)
|
||||
if unicodedata.category(chr(cp)) != "Cn"]
|
||||
if not letters:
|
||||
continue
|
||||
step = max(1, len(letters) // 24)
|
||||
sample = letters[::step][:24]
|
||||
out.append("".join(sample))
|
||||
out.append(" ".join("".join(sample[i:i + 4]) for i in range(0, len(sample), 4)))
|
||||
return out
|
||||
|
||||
|
||||
def assigned(s):
|
||||
return all(unicodedata.category(c) != "Cn" and not (0xD800 <= ord(c) <= 0xDFFF) for c in s)
|
||||
|
||||
|
||||
def sorted_keys(v):
|
||||
if isinstance(v, dict):
|
||||
return list(v) == sorted(v) and all(sorted_keys(x) for x in v.values())
|
||||
if isinstance(v, list):
|
||||
return all(sorted_keys(x) for x in v)
|
||||
return True
|
||||
|
||||
|
||||
def all_strings(v):
|
||||
if isinstance(v, str):
|
||||
yield v
|
||||
elif isinstance(v, dict):
|
||||
for x in v.values():
|
||||
yield from all_strings(x)
|
||||
elif isinstance(v, list):
|
||||
for x in v:
|
||||
yield from all_strings(x)
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--out", required=True)
|
||||
ap.add_argument("--fuzz", type=int, default=400)
|
||||
ap.add_argument("--seed", type=int, default=339)
|
||||
args = ap.parse_args()
|
||||
|
||||
states = []
|
||||
for name, s in SENTENCES.items():
|
||||
states.append(s)
|
||||
ascii_s = "".join(c for c in unicodedata.normalize("NFKD", s) if not unicodedata.combining(c))
|
||||
if ascii_s != s:
|
||||
states.append(ascii_s)
|
||||
states.append(s.upper())
|
||||
states.append(" ".join(s.split()[:3]))
|
||||
states += EXTRA + CODE_AND_IDENTIFIERS + MULTILINE + budget_cases()
|
||||
states += generated_script_runs()
|
||||
states += structured_cases()
|
||||
# Pairwise mixes: one English and one foreign line, both orders.
|
||||
for name, s in SENTENCES.items():
|
||||
if not name.startswith("en"):
|
||||
states.append(SENTENCES["en"] + "\n" + s)
|
||||
states.append(s + " " + SENTENCES["en"])
|
||||
states += fuzz_cases(args.fuzz, args.seed)
|
||||
|
||||
cases, seen = [], set()
|
||||
for st in states:
|
||||
key = json.dumps(st, ensure_ascii=False, sort_keys=False)
|
||||
if key in seen:
|
||||
continue
|
||||
seen.add(key)
|
||||
assert sorted_keys(st), "mapping keys must be sorted: %r" % (st,)
|
||||
if not all(assigned(s) for s in all_strings(st)):
|
||||
continue
|
||||
det = analyse(st)
|
||||
# Recorded separately because a JSON reader that does not keep key
|
||||
# order loses the profile's order, which the reference defines.
|
||||
rec = {"state": st, "leaves": _iter_text(st), "analysis": det,
|
||||
"profile_order": list(det["script_profile"])}
|
||||
if isinstance(st, str):
|
||||
rec["words"] = _WORD.findall(st)
|
||||
cases.append(rec)
|
||||
|
||||
header = {
|
||||
"laya_version": laya.__version__,
|
||||
"python": sys.version.split()[0],
|
||||
"unidata_version": unicodedata.unidata_version,
|
||||
"fuzz_seed": args.seed,
|
||||
}
|
||||
with open(args.out, "w") as f:
|
||||
json.dump({"header": header, "cases": cases}, f, ensure_ascii=False, indent=0)
|
||||
print("%d cases" % len(cases), file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,385 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Record Laya reference outputs for neuron's parity tests (#334).
|
||||
|
||||
Runs the upstream `laya` SDK (the reference implementation, not the
|
||||
`rl_common.py` shipped in the model repo, which has diverged) over a fixed
|
||||
corpus and writes, per case:
|
||||
|
||||
- the request exactly as a `/v1/systemone` client would send it
|
||||
- per question: `input_ids`, `[MASK]` marker positions, question type and
|
||||
the option-collapse stats, straight out of `Agent._encode_state`
|
||||
- the raw scorer logits and act probabilities out of `Agent._forward`,
|
||||
before any temperature is applied
|
||||
- the final Jev-shaped response (`answers`, `usage`, `routing`)
|
||||
|
||||
A separate section records `laya.lang.analyse` on a language corpus, for
|
||||
the router port.
|
||||
|
||||
The reference runs on CPU in float32 so the numbers are deterministic and
|
||||
free of autocast noise; neuron's reduced-precision paths are compared
|
||||
against them with a tolerance.
|
||||
|
||||
Usage (in a venv with `pip install laya==<pinned>` and CPU torch):
|
||||
|
||||
HF_ENDPOINT=... HF_TOKEN=... python script/laya-reference.py \
|
||||
--out crates/neuron/src/harness/testdata/laya
|
||||
|
||||
Point HF_ENDPOINT at whichever Hugging Face-compatible registry you fetch
|
||||
weights through. Python's huggingface_hub may additionally need
|
||||
SSL_CERT_FILE (for a private CA) and HF_HUB_DISABLE_XET=1 (for registries
|
||||
without Xet support).
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
import laya
|
||||
import numpy as np
|
||||
import torch
|
||||
from laya import Router
|
||||
from laya.agent import Agent
|
||||
from laya.lang import analyse
|
||||
from laya.serve import _resolve_model
|
||||
|
||||
LONG_TEXT = (
|
||||
"Hello team, following up on the invoice discrepancy from last quarter. "
|
||||
"The amounts on INV-2291 and INV-2307 do not match the purchase order, "
|
||||
"and our finance department needs a corrected copy before month end. "
|
||||
) * 12
|
||||
|
||||
CASES = [
|
||||
# --- the three primitives, default budgets
|
||||
{
|
||||
"name": "choice_basic",
|
||||
"state": "I was charged twice this month, I want my money back",
|
||||
"questions": {
|
||||
"queue": {"type": "choice", "instructions": "Which team?",
|
||||
"criteria": {"billing": "billing and refunds", "tech": "login and app issues",
|
||||
"other": "everything else"}},
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "score_basic",
|
||||
"state": "This is the THIRD time I've asked. Fix it today or I'm cancelling.",
|
||||
"questions": {
|
||||
"urgency": {"type": "score", "instructions": "How urgent?",
|
||||
"criteria": ["calm", "firm", "angry", "furious"]},
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "noul_default_criteria",
|
||||
"state": "Please refund the duplicate charge on my card.",
|
||||
"questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
|
||||
},
|
||||
{
|
||||
"name": "noul_custom_criteria_and_labels",
|
||||
"state": "Can you reset my password? I'm locked out.",
|
||||
"questions": {
|
||||
"locked": {"type": "noul", "instructions": "Is the user locked out of their account?",
|
||||
"criteria": {"true": "the user cannot sign in", "false": "the user can sign in"},
|
||||
"labels": {"true": "locked", "false": "not locked"}},
|
||||
},
|
||||
},
|
||||
# --- several questions in one request: the shared-state, one-forward path
|
||||
{
|
||||
"name": "multi_question",
|
||||
"state": "I was charged twice this month, I want my money back",
|
||||
"questions": {
|
||||
"queue": {"type": "choice", "instructions": "Which team?",
|
||||
"criteria": {"billing": "billing and refunds", "tech": "login and app issues",
|
||||
"other": "everything else"}},
|
||||
"urgency": {"type": "score", "instructions": "How urgent?",
|
||||
"criteria": ["calm", "firm", "angry", "furious"]},
|
||||
"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"},
|
||||
"tone": {"type": "choice", "instructions": "Tone?", "criteria": ["polite", "rude", "neutral"]},
|
||||
"spam": {"type": "noul", "instructions": "Is this spam?"},
|
||||
},
|
||||
},
|
||||
# --- every temperature bucket
|
||||
{
|
||||
"name": "choice_2",
|
||||
"state": "The app crashes when I open settings.",
|
||||
"questions": {"kind": {"type": "choice", "instructions": "Bug or feature request?",
|
||||
"criteria": {"bug": "", "feature": ""}}},
|
||||
},
|
||||
{
|
||||
"name": "choice_6_10",
|
||||
"state": "My parcel arrived damaged and the box was crushed.",
|
||||
"questions": {"topic": {"type": "choice", "instructions": "Topic?",
|
||||
"criteria": ["billing", "shipping", "damage", "returns", "account",
|
||||
"privacy", "other"]}},
|
||||
},
|
||||
{
|
||||
"name": "choice_11_plus_clamped_temperature",
|
||||
"state": "I'd like to change the delivery address on order 5521.",
|
||||
"questions": {"intent": {"type": "choice", "instructions": "Intent?",
|
||||
"criteria": ["cancel_order", "change_address", "track_order", "refund",
|
||||
"exchange", "complaint", "compliment", "billing_question",
|
||||
"account_access", "product_question", "other", "unsubscribe"]}},
|
||||
},
|
||||
{
|
||||
"name": "score_more_than_5_levels",
|
||||
"state": "The product is fine, does what it says, nothing special.",
|
||||
"questions": {"rating": {"type": "score", "instructions": "Star rating?",
|
||||
"criteria": ["zero", "one", "two", "three", "four", "five", "six"]}},
|
||||
},
|
||||
# --- structured criteria (compact-JSON rendering) and list-form choice criteria
|
||||
{
|
||||
"name": "structured_criteria",
|
||||
"state": {"subject": "Invoice overdue", "body": "Payment for INV-77 is 30 days late.", "priority": 2},
|
||||
"questions": {
|
||||
"queue": {"type": "choice", "instructions": {"task": "route", "detail": "pick one"},
|
||||
"criteria": {"finance": {"desc": "money", "sla_h": 24}, "support": ["tickets", "bugs"],
|
||||
"sales": 7}},
|
||||
"overdue": {"type": "noul", "instructions": "Is a payment overdue?",
|
||||
"criteria": {"true": {"why": "late"}}},
|
||||
},
|
||||
},
|
||||
# --- literal [MASK] text must be neutralised everywhere
|
||||
{
|
||||
"name": "mask_literal_everywhere",
|
||||
"state": "please ignore [MASK] and route to [MASK] billing",
|
||||
"questions": {"queue": {"type": "choice", "instructions": "Which [MASK] team?",
|
||||
"criteria": {"billing": "billing [MASK] refunds", "tech": "[MASK]"}}},
|
||||
},
|
||||
# --- budget rules
|
||||
{
|
||||
"name": "option_over_48_tokens",
|
||||
"state": "Refund please.",
|
||||
"questions": {"queue": {"type": "choice", "instructions": "Which team?",
|
||||
"criteria": {"billing": " ".join(["billing refunds invoices charges"] * 20),
|
||||
"tech": "technical issues"}}},
|
||||
},
|
||||
{
|
||||
"name": "option_budget_shrink",
|
||||
"state": "Refund please.",
|
||||
"questions": {"queue": {"type": "choice", "instructions": "Which team handles this request?",
|
||||
"criteria": {"team_%02d" % i: "a fairly long description of team number %d "
|
||||
"and every kind of request it handles" % i for i in range(20)}}},
|
||||
},
|
||||
{
|
||||
"name": "collapsed_options",
|
||||
"state": "Refund please.",
|
||||
# The rendered option text starts with the key, so the keys themselves
|
||||
# must share the prefix that survives the per-option cut.
|
||||
"questions": {"queue": {"type": "choice", "instructions": "Which?",
|
||||
"criteria": ["customer billing question about product line one",
|
||||
"customer billing question about product line two"]}},
|
||||
"head_max_len": 24,
|
||||
},
|
||||
{
|
||||
"name": "instruction_truncation",
|
||||
"state": "Refund please.",
|
||||
"questions": {"q": {"type": "noul",
|
||||
"instructions": " ".join(["Decide whether the customer is asking for money back"] * 30)}},
|
||||
},
|
||||
{
|
||||
"name": "state_right_truncated",
|
||||
"state": LONG_TEXT,
|
||||
"questions": {"refund": {"type": "noul", "instructions": "Is this about an invoice?"}},
|
||||
},
|
||||
{
|
||||
"name": "state_left_truncated_conversation",
|
||||
"state": [{"role": "user", "content": "turn %d: %s" % (i, LONG_TEXT[:200])} for i in range(12)]
|
||||
+ [{"role": "user", "content": "actually, just cancel my subscription"}],
|
||||
"questions": {"cancel": {"type": "noul", "instructions": "Does the user want to cancel?"}},
|
||||
},
|
||||
{
|
||||
"name": "budget_overrides",
|
||||
"state": LONG_TEXT,
|
||||
"questions": {"queue": {"type": "choice", "instructions": "Which team?",
|
||||
"criteria": ["finance", "support", "sales"]}},
|
||||
"max_len": 128,
|
||||
"head_max_len": 64,
|
||||
},
|
||||
# --- non-ASCII inside a JSON state
|
||||
{
|
||||
"name": "json_state_non_ascii",
|
||||
"state": {"from": "Zoë", "text": "Die Rechnung ist falsch — bitte korrigieren. 请尽快处理"},
|
||||
"questions": {"queue": {"type": "choice", "instructions": "Which team?",
|
||||
"criteria": ["billing", "tech", "other"]}},
|
||||
},
|
||||
# --- routing to the multilingual checkpoint
|
||||
{
|
||||
"name": "ml_devanagari",
|
||||
"state": "मुझे इस महीने दो बार शुल्क लिया गया, मुझे पैसे वापस चाहिए",
|
||||
"questions": {"queue": {"type": "choice", "instructions": "Which team?",
|
||||
"criteria": {"billing": "billing and refunds", "tech": "login and app issues",
|
||||
"other": "everything else"}}},
|
||||
},
|
||||
{
|
||||
"name": "ml_khmer",
|
||||
"state": "ខ្ញុំត្រូវបានគិតប្រាក់ពីរដងក្នុងខែនេះ ខ្ញុំចង់បានប្រាក់វិញ",
|
||||
"questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
|
||||
},
|
||||
{
|
||||
"name": "ml_cjk",
|
||||
"state": "这个月我被扣了两次钱,我要退款",
|
||||
"questions": {"urgency": {"type": "score", "instructions": "How urgent?",
|
||||
"criteria": ["calm", "firm", "angry", "furious"]}},
|
||||
},
|
||||
{
|
||||
"name": "ml_arabic",
|
||||
"state": "تم خصم المبلغ مرتين هذا الشهر، أريد استرداد أموالي",
|
||||
"questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
|
||||
},
|
||||
{
|
||||
"name": "ml_latin_non_english",
|
||||
"state": "Me han cobrado dos veces este mes y quiero que me devuelvan el dinero",
|
||||
"questions": {"queue": {"type": "choice", "instructions": "Which team?",
|
||||
"criteria": ["billing", "tech", "other"]}},
|
||||
},
|
||||
{
|
||||
"name": "ml_mixed_script",
|
||||
"state": "Order #5521 — 注文をキャンセルしたいです. Please help.",
|
||||
"questions": {"cancel": {"type": "noul", "instructions": "Does the user want to cancel?"}},
|
||||
},
|
||||
# --- explicit checkpoint selection through `model`
|
||||
{
|
||||
"name": "model_pinned_multilingual",
|
||||
"state": "I was charged twice this month, I want my money back",
|
||||
"questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
|
||||
"model": "multilingual",
|
||||
},
|
||||
{
|
||||
"name": "model_pinned_typed_decisions",
|
||||
"state": "I was charged twice this month, I want my money back",
|
||||
"questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
|
||||
"model": "typed-decisions",
|
||||
},
|
||||
{
|
||||
"name": "model_foreign_name_routes_by_default",
|
||||
"state": "I was charged twice this month, I want my money back",
|
||||
"questions": {"refund": {"type": "noul", "instructions": "Is the customer asking for a refund?"}},
|
||||
"model": "jev-1",
|
||||
},
|
||||
]
|
||||
|
||||
LANG_CASES = [
|
||||
"I was charged twice this month, I want my money back",
|
||||
"Me han cobrado dos veces este mes",
|
||||
"Die Rechnung ist falsch, bitte korrigieren",
|
||||
"Je voudrais annuler ma commande",
|
||||
"मुझे पैसे वापस चाहिए",
|
||||
"ខ្ញុំចង់បានប្រាក់វិញ",
|
||||
"这个月我被扣了两次钱",
|
||||
"注文をキャンセルしたいです",
|
||||
"تم خصم المبلغ مرتين",
|
||||
"Мне дважды списали деньги",
|
||||
"Order #5521 — 注文をキャンセルしたいです. Please help.",
|
||||
"ok",
|
||||
"",
|
||||
"12345 67890",
|
||||
"café résumé naïve",
|
||||
{"subject": "Invoice", "body": "Bitte zahlen Sie bis Freitag"},
|
||||
["hello", "bonjour", "hola"],
|
||||
]
|
||||
|
||||
|
||||
class Capture:
|
||||
"""Wraps the two Agent internals the fixtures need, per call."""
|
||||
|
||||
def __init__(self):
|
||||
self.items = None
|
||||
self.logits = None
|
||||
self.act = None
|
||||
self.agent = None
|
||||
orig_encode, orig_forward = Agent._encode_state, Agent._forward
|
||||
cap = self
|
||||
|
||||
def encode(agent, *a, **kw):
|
||||
items = orig_encode(agent, *a, **kw)
|
||||
cap.items, cap.agent = items, agent
|
||||
return items
|
||||
|
||||
def forward(agent, b):
|
||||
logits, act = orig_forward(agent, b)
|
||||
cap.logits, cap.act = logits, act
|
||||
return logits, act
|
||||
|
||||
Agent._encode_state, Agent._forward = encode, forward
|
||||
|
||||
def reset(self):
|
||||
self.items = self.logits = self.act = self.agent = None
|
||||
|
||||
|
||||
def to_py(x):
|
||||
if isinstance(x, np.ndarray):
|
||||
return x.tolist()
|
||||
if isinstance(x, (np.floating, np.integer)):
|
||||
return x.item()
|
||||
if isinstance(x, dict):
|
||||
return {k: to_py(v) for k, v in x.items()}
|
||||
if isinstance(x, (list, tuple)):
|
||||
return [to_py(v) for v in x]
|
||||
return x
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--out", required=True)
|
||||
args = ap.parse_args()
|
||||
os.makedirs(args.out, exist_ok=True)
|
||||
torch.manual_seed(0)
|
||||
torch.set_num_threads(max(1, (os.cpu_count() or 2) // 2))
|
||||
|
||||
cap = Capture()
|
||||
router = Router(device="cpu", preload=True, max_loaded=3)
|
||||
header = {
|
||||
"laya_version": laya.__version__,
|
||||
"torch_version": torch.__version__,
|
||||
"device": "cpu",
|
||||
"dtype": "float32",
|
||||
}
|
||||
|
||||
cases, checkpoints = [], {}
|
||||
for case in CASES:
|
||||
req = {"state": case["state"], "questions": case["questions"]}
|
||||
kw = {}
|
||||
for key in ("model", "max_len", "head_max_len"):
|
||||
if key in case:
|
||||
req[key] = case[key]
|
||||
if key != "model":
|
||||
kw[key] = case[key]
|
||||
cap.reset()
|
||||
# The HTTP server, not the Router, maps a foreign model name (a Jev id)
|
||||
# to "let the router choose"; go through the same function.
|
||||
resp = router.predict(case["state"], case["questions"], model=_resolve_model(case.get("model")), **kw)
|
||||
agent = cap.agent
|
||||
rows = []
|
||||
for j, (qid, it) in enumerate(zip(case["questions"], cap.items)):
|
||||
k = len(it["markers"])
|
||||
rows.append({
|
||||
"qid": qid,
|
||||
"input_ids": list(it["ids"]),
|
||||
"markers": list(it["markers"]),
|
||||
"qtype": it["qtype"],
|
||||
"options": it.get("options"),
|
||||
"logits": cap.logits[j, :k].tolist(),
|
||||
"act": cap.act[j].tolist(),
|
||||
})
|
||||
ckpt = resp["routing"]["model"]
|
||||
if ckpt not in checkpoints:
|
||||
checkpoints[ckpt] = {
|
||||
"repo": resp["routing"].get("repo"),
|
||||
"revision": getattr(agent, "revision", None),
|
||||
"cfg": {k: agent.cfg.get(k) for k in ("max_len", "head_max_len", "temperature",
|
||||
"temperature_by_options")},
|
||||
"special_ids": {"cls": agent.tok.cls_token_id, "sep": agent.tok.sep_token_id,
|
||||
"mask": agent.tok.mask_token_id, "pad": agent.tok.pad_token_id},
|
||||
}
|
||||
cases.append({"name": case["name"], "request": req, "items": rows, "response": to_py(resp)})
|
||||
print("%-40s -> %-16s %d question(s)" % (case["name"], ckpt, len(rows)), file=sys.stderr)
|
||||
|
||||
lang = [{"state": s, "analysis": to_py(analyse(s))} for s in LANG_CASES]
|
||||
|
||||
header["checkpoints"] = checkpoints
|
||||
with open(os.path.join(args.out, "reference.json"), "w") as f:
|
||||
json.dump({"header": header, "cases": cases}, f, ensure_ascii=False, indent=1)
|
||||
with open(os.path.join(args.out, "lang.json"), "w") as f:
|
||||
json.dump({"header": header, "cases": lang}, f, ensure_ascii=False, indent=1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+133
@@ -0,0 +1,133 @@
|
||||
#!/usr/bin/env python3
|
||||
"""End-to-end check of /v1/systemone decision serving (#341).
|
||||
|
||||
Replays every case of the recorded Laya reference corpus
|
||||
(`crates/neuron/src/harness/testdata/laya/reference.json`) against a
|
||||
serving endpoint -- a neuron directly, a cortex, or the public router --
|
||||
and compares each response with what the reference server returned:
|
||||
|
||||
- `routing.model` (which checkpoint answered) and `usage` must match
|
||||
exactly;
|
||||
- each answer's type, and its chosen option or score level, must match,
|
||||
except where the reference itself was within `--tie` of a tie (reduced
|
||||
precision may break a near-tie either way);
|
||||
- probabilities and confidences must be within `--tolerance`. The
|
||||
reference is f32; a bf16 GPU path drifts from it by about as much as
|
||||
the reference's own bf16 GPU run does.
|
||||
|
||||
It prints each case's latency and a summary, and exits non-zero on any
|
||||
mismatch. Nothing is loaded or unloaded: point it at an endpoint that
|
||||
already serves a decision model.
|
||||
|
||||
Usage:
|
||||
script/validate-decision.py URL [--model M] [--bearer KEY]
|
||||
|
||||
URL http://quadbrat.hanzalova.internal:13131 (neuron)
|
||||
http://hanzalova.internal:31313 (cortex)
|
||||
https://router.example (public chain)
|
||||
--model override every request's `model` (e.g. helexa/one); by
|
||||
default each case sends what it was recorded with, which
|
||||
includes checkpoint pins and a foreign Jev model name.
|
||||
"""
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
import urllib.error
|
||||
import urllib.request
|
||||
|
||||
FIXTURE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "crates", "neuron", "src",
|
||||
"harness", "testdata", "laya", "reference.json")
|
||||
|
||||
|
||||
def post(url, body, bearer):
|
||||
req = urllib.request.Request(url, data=json.dumps(body).encode(), method="POST",
|
||||
headers={"Content-Type": "application/json"})
|
||||
if bearer:
|
||||
req.add_header("Authorization", "Bearer " + bearer)
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=120) as r:
|
||||
payload = json.load(r)
|
||||
status = r.status
|
||||
except urllib.error.HTTPError as e:
|
||||
payload = json.loads(e.read() or b"{}")
|
||||
status = e.code
|
||||
return status, payload, (time.perf_counter() - started) * 1e3
|
||||
|
||||
|
||||
def top_gap(probs):
|
||||
p = sorted(probs, reverse=True)
|
||||
return p[0] - p[1] if len(p) > 1 else 1.0
|
||||
|
||||
|
||||
def compare(name, got, want, tol, tie):
|
||||
errs = []
|
||||
if got.get("routing", {}).get("model") != want["routing"]["model"]:
|
||||
errs.append("routing.model %r != %r" % (got.get("routing", {}).get("model"), want["routing"]["model"]))
|
||||
if got.get("usage") != want["usage"]:
|
||||
errs.append("usage %r != %r" % (got.get("usage"), want["usage"]))
|
||||
ga, wa = got.get("answers", {}), want["answers"]
|
||||
if list(ga) != list(wa):
|
||||
errs.append("answer ids %r != %r" % (list(ga), list(wa)))
|
||||
return errs, 0.0
|
||||
worst = 0.0
|
||||
for qid, w in wa.items():
|
||||
g = ga[qid]
|
||||
if g.get("type") != w["type"]:
|
||||
errs.append("%s: type %r != %r" % (qid, g.get("type"), w["type"]))
|
||||
continue
|
||||
wp = w.get("probabilities")
|
||||
near_tie = wp is not None and top_gap(list(wp.values())) < tie
|
||||
if w["type"] == "choice" and g.get("choice") != w["choice"] and not near_tie:
|
||||
errs.append("%s: choice %r != %r" % (qid, g.get("choice"), w["choice"]))
|
||||
if wp:
|
||||
for opt, p in wp.items():
|
||||
worst = max(worst, abs(g["probabilities"][opt] - p))
|
||||
for key in ("noul", "confidence", "answer_confidence"):
|
||||
if key in w:
|
||||
worst = max(worst, abs(g[key] - w[key]))
|
||||
if worst > tol:
|
||||
errs.append("max |Δ| %.4f > %.4f" % (worst, tol))
|
||||
return errs, worst
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("url")
|
||||
ap.add_argument("--model")
|
||||
ap.add_argument("--bearer")
|
||||
ap.add_argument("--tolerance", type=float, default=0.1)
|
||||
ap.add_argument("--tie", type=float, default=0.05)
|
||||
args = ap.parse_args()
|
||||
endpoint = args.url.rstrip("/") + "/v1/systemone"
|
||||
cases = json.load(open(FIXTURE))["cases"]
|
||||
|
||||
failures, latencies, worst_all = 0, [], 0.0
|
||||
for case in cases:
|
||||
body = dict(case["request"])
|
||||
if args.model:
|
||||
body["model"] = args.model
|
||||
status, got, ms = post(endpoint, body, args.bearer)
|
||||
latencies.append(ms)
|
||||
if status != 200:
|
||||
failures += 1
|
||||
print("FAIL %-40s HTTP %d %s" % (case["name"], status, json.dumps(got)[:200]))
|
||||
continue
|
||||
errs, worst = compare(case["name"], got, case["response"], args.tolerance, args.tie)
|
||||
worst_all = max(worst_all, worst)
|
||||
mark = "ok " if not errs else "FAIL"
|
||||
failures += bool(errs)
|
||||
print("%s %-40s %7.1f ms %-15s max|Δ| %.4f %s" % (
|
||||
mark, case["name"], ms, got.get("routing", {}).get("model", "?"), worst, "; ".join(errs)))
|
||||
|
||||
print("\n%d/%d cases ok; latency median %.1f ms, p90 %.1f ms; worst |Δ| %.4f" % (
|
||||
len(cases) - failures, len(cases), statistics.median(latencies),
|
||||
sorted(latencies)[int(0.9 * (len(latencies) - 1))], worst_all))
|
||||
sys.exit(1 if failures else 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+44
-27
@@ -9,17 +9,27 @@
|
||||
# image modality.
|
||||
#
|
||||
# Usage:
|
||||
# script/validate-image.sh [host] [model_id]
|
||||
# script/validate-image.sh [host] [model_id] [quant]
|
||||
#
|
||||
# Defaults:
|
||||
# host = benjy.hanzalova.internal (the v1 image placement)
|
||||
# model_id = Tongyi-MAI/Z-Image-Turbo
|
||||
# quant = none (BF16). A 12 GB card needs q8_0.
|
||||
#
|
||||
# A co-resident model evicted for the probe window is restored on every
|
||||
# exit path, failures included.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
HOST="${1:-benjy.hanzalova.internal}"
|
||||
MODEL_ID="${2:-Tongyi-MAI/Z-Image-Turbo}"
|
||||
QUANT="${3:-}"
|
||||
BASE="http://${HOST}:13131"
|
||||
if [[ -n "${QUANT}" ]]; then
|
||||
LOAD_JSON="{\"model_id\": \"${MODEL_ID}\", \"harness\": \"candle\", \"quant\": \"${QUANT}\"}"
|
||||
else
|
||||
LOAD_JSON="{\"model_id\": \"${MODEL_ID}\", \"harness\": \"candle\"}"
|
||||
fi
|
||||
|
||||
echo "==> /version"
|
||||
curl -sf -m 10 "${BASE}/version" | python3 -c 'import json,sys; b=json.load(sys.stdin); print(" ", b.get("git_sha","?")[:12], b.get("profile"), b.get("features"))'
|
||||
@@ -28,26 +38,47 @@ curl -sf -m 10 "${BASE}/version" | python3 -c 'import json,sys; b=json.load(sys.
|
||||
# cortex's job. On a host whose resident text model leaves too little
|
||||
# VRAM for the DiT (benjy: 8B @ 16 GB + 14.6 GB DiT > 24 GB), evict
|
||||
# it for the probe window and restore it afterwards.
|
||||
RESTORE_JSON=""
|
||||
EVICT=$(curl -sf -m 10 "${BASE}/models" | python3 -c '
|
||||
EVICTED=()
|
||||
# Unload the probe model and restore the evicted ones however the script
|
||||
# ends: a failed generation must not leave the host without its models.
|
||||
restore() {
|
||||
local rc=$?
|
||||
curl -s -m 60 -X POST "${BASE}/models/unload" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"model_id\": \"${MODEL_ID}\"}" > /dev/null || true
|
||||
for id in "${EVICTED[@]}"; do
|
||||
echo "==> restoring ${id}"
|
||||
curl -s -m 900 -X POST "${BASE}/models/load" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"model_id\": \"${id}\", \"harness\": \"candle\"}" \
|
||||
|| echo " RESTORE FAILED: reload ${id} by hand"
|
||||
echo
|
||||
done
|
||||
exit "${rc}"
|
||||
}
|
||||
trap restore EXIT
|
||||
|
||||
mapfile -t RESIDENT < <(curl -sf -m 10 "${BASE}/models" | python3 -c '
|
||||
import json, sys
|
||||
models = json.load(sys.stdin)
|
||||
loaded = [m for m in models if m["status"] == "loaded" and "image" not in m.get("capabilities", [])]
|
||||
print(loaded[0]["id"] if loaded else "")
|
||||
for m in json.load(sys.stdin):
|
||||
if m["status"] == "loaded" and "image" not in m.get("capabilities", []):
|
||||
print(m["id"])
|
||||
')
|
||||
if [[ -n "${EVICT}" ]]; then
|
||||
echo "==> evicting co-resident ${EVICT} for the probe window"
|
||||
for id in "${RESIDENT[@]}"; do
|
||||
[[ -z "${id}" ]] && continue
|
||||
echo "==> evicting co-resident ${id} for the probe window"
|
||||
curl -sf -m 120 -X POST "${BASE}/models/unload" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"model_id\": \"${EVICT}\"}"
|
||||
-d "{\"model_id\": \"${id}\"}"
|
||||
echo
|
||||
RESTORE_JSON="{\"model_id\": \"${EVICT}\", \"harness\": \"candle\"}"
|
||||
fi
|
||||
EVICTED+=("${id}")
|
||||
done
|
||||
|
||||
|
||||
echo "==> load ${MODEL_ID}"
|
||||
curl -sf -m 900 -X POST "${BASE}/models/load" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"model_id\": \"${MODEL_ID}\", \"harness\": \"candle\"}"
|
||||
-d "${LOAD_JSON}"
|
||||
echo
|
||||
|
||||
echo "==> capabilities"
|
||||
@@ -80,20 +111,6 @@ t = usage["helexa_timing"]
|
||||
print(" 512x512 ok | units={:.3f} | encode={}ms denoise={}ms decode={}ms".format(units, t["encode_ms"], t["denoise_ms"], t["decode_ms"]))
|
||||
'
|
||||
|
||||
echo "==> unload"
|
||||
curl -sf -m 60 -X POST "${BASE}/models/unload" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "{\"model_id\": \"${MODEL_ID}\"}"
|
||||
echo
|
||||
|
||||
if [[ -n "${RESTORE_JSON}" ]]; then
|
||||
echo "==> restoring ${EVICT}"
|
||||
curl -sf -m 900 -X POST "${BASE}/models/load" \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d "${RESTORE_JSON}"
|
||||
echo
|
||||
fi
|
||||
|
||||
echo "==> health"
|
||||
curl -sf -m 10 "${BASE}/health" > /dev/null && echo " healthy"
|
||||
echo "PASS"
|
||||
echo "PASS (unload and restore follow)"
|
||||
|
||||
Reference in New Issue
Block a user