Compare commits
22
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 |
@@ -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
|
||||
|
||||
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;
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
@@ -751,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"),
|
||||
}
|
||||
}
|
||||
@@ -823,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,
|
||||
@@ -841,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(),
|
||||
}
|
||||
}
|
||||
@@ -1969,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()));
|
||||
}
|
||||
|
||||
@@ -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
|
||||
@@ -549,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
|
||||
|
||||
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"];
|
||||
|
||||
+46824
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
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]"
|
||||
}
|
||||
@@ -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()
|
||||
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