diff --git a/rust/src/engine-core-client/src/protocol/logprobs.rs b/rust/src/engine-core-client/src/protocol/logprobs.rs index d5bc40e061..646c106a04 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs.rs @@ -25,10 +25,13 @@ use crate::protocol::tensor::{WireArrayData, WireNdArray}; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct TokenLogprob { pub token_id: u32, + /// Preserves the engine's value, including NaN and infinities. pub logprob: f32, /// The sampled/selected token uses its actual vocab rank. Remaining entries /// use 1-based top-k ranks matching the engine's returned candidate /// order. + /// A sampled/selected rank of 0 occurs when its logprob is NaN: the engine's + /// `(logprobs >= selected_logprob).sum(-1)` counts no matching values. pub rank: u32, } @@ -53,10 +56,6 @@ impl PositionLogprobs { logprobs.len() ); } - if sampled_rank == 0 { - bail_ext_value_decode!("token_ranks must be >= 1 for decoded engine-core logprobs"); - } - let mut entries = Vec::with_capacity(token_ids.len()); for (index, (&token_id, &logprob)) in token_ids.iter().zip(logprobs.iter()).enumerate() { let rank = if index == 0 { diff --git a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs index 2470460b52..f69ff22f1a 100644 --- a/rust/src/engine-core-client/src/protocol/logprobs/tests.rs +++ b/rust/src/engine-core-client/src/protocol/logprobs/tests.rs @@ -396,3 +396,121 @@ fn rejects_zero_column_logprobs_with_rows() { "new_logprobs: zero-column logprobs payload with 2 rows" ); } + +fn logprobs_value(ids: &[i64], probs: &[f32], ranks: &[i64]) -> Value { + let rows = ranks.len(); + let cols = ids.len() / rows; + Value::Array(vec![ + ndarray_value( + "