iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Rust Frontend] Lower out-of-vocab validation to text layer (#45685)

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
This commit is contained in:
Bugen Zhao
2026-06-16 03:37:58 +00:00
committed by GitHub
parent 3f65e21e32
commit f99260d2aa
14 changed files with 278 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/server/src/error.rs
+12
View File
@@ -102,6 +102,7 @@ fn is_request_validation_error(error: &vllm_text::Error) -> bool {
vllm_text::Error::PromptTooLong { .. }
| vllm_text::Error::EmptyPromptTokenIds { .. }
| vllm_text::Error::Logprobs(_)
| vllm_text::Error::OutOfVocab(_)
// An empty tokenized prompt detected later, at request prepare
// time, surfaces through the transparent Llm wrapper.
| vllm_text::Error::Llm(vllm_llm::Error::EmptyPromptTokenIds { .. })
@@ -172,6 +173,17 @@ mod tests {
assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
}
#[test]
fn out_of_vocab_validation_maps_to_invalid_request() {
let error = vllm_text::Error::OutOfVocab(vllm_text::OutOfVocabError {
parameter: "logprob_token_ids",
token_ids: vec![1000],
vocab_size: 1000,
});
let api_error = text_submit_error("failed to submit completion request", error);
assert_eq!(api_error.status_code(), StatusCode::BAD_REQUEST);
}
#[test]
fn other_submit_errors_stay_internal() {
let error = vllm_text::Error::Tokenizer("backend exploded".to_string());
rust/src/server/src/routes/openai/chat_completions.rs
-7
View File
@@ -53,13 +53,6 @@ pub async fn chat_completions(
let request_context = resolve_request_context(&headers, body.request_id.as_deref());
let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
if let Err(err) = validate::validate_token_id_ranges(
&body,
state.tokenizer_vocab_size(),
state.model_vocab_size(),
) {
return err.into_response();
}
let prepared = match prepare_chat_request(body, &lora_resolution, request_context) {
Ok(prepared) => prepared,
Err(error) => return error.into_response(),
rust/src/server/src/routes/openai/chat_completions/validate.rs
+1 -43
View File
@@ -1,6 +1,5 @@
use super::types::ChatCompletionRequest;
use crate::error::{ApiError, bail_invalid_request};
use crate::routes::openai::utils::token_ids::{validate_allowed_token_ids, validate_logit_bias};
use crate::routes::openai::utils::types::{ChatMessage, Tool, ToolChoice, ToolChoiceValue};
/// Enforce the minimal compatibility contract for the Rust OpenAI server.
@@ -154,21 +153,6 @@ fn validate_function_tools(tools: &[Tool], param: &'static str) -> Result<(), Ap
Ok(())
}
/// Reject out-of-vocab token ids, mirroring the Python input processor:
/// `allowed_token_ids` against the tokenizer vocab, `logit_bias` keys against the
/// model vocab (skipped when the model size is unknown).
pub(super) fn validate_token_id_ranges(
request: &ChatCompletionRequest,
tokenizer_vocab_size: usize,
model_vocab_size: Option<usize>,
) -> Result<(), ApiError> {
validate_allowed_token_ids(request.allowed_token_ids.as_deref(), tokenizer_vocab_size)?;
validate_logit_bias(
request.logit_bias.as_ref(),
model_vocab_size.unwrap_or(usize::MAX),
)
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
@@ -176,7 +160,7 @@ mod tests {
use serde_json::json;
use vllm_chat::ReasoningEffort;
use super::{validate_request_compat, validate_token_id_ranges};
use super::validate_request_compat;
use crate::routes::openai::chat_completions::types::ChatCompletionRequest;
use crate::routes::openai::utils::structured_outputs::ResponseFormat;
use crate::routes::openai::utils::types::{
@@ -188,32 +172,6 @@ mod tests {
names.iter().map(|s| s.to_string()).collect()
}
#[test]
fn validate_token_id_ranges_rejects_oob_and_accepts_in_vocab() {
// allowed_token_ids are bounded by the tokenizer vocab
let mut request = base_request();
request.allowed_token_ids = Some(vec![5, 1_000_000]);
assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
// logit_bias is bounded by the larger model vocab: an id between the two
// vocabs is valid and must not be rejected (the parity regression we fix)
let mut request = base_request();
request.logit_bias = Some(HashMap::from([("150".to_string(), 1.0)]));
assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
// logit_bias beyond the model vocab -> reject
let mut request = base_request();
request.logit_bias = Some(HashMap::from([("1000000".to_string(), 1.0)]));
assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
// all in-vocab -> accept
let mut request = base_request();
request.allowed_token_ids = Some(vec![5, 50]);
request.logit_bias = Some(HashMap::from([("50".to_string(), 1.0)]));
assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
// unknown sizes -> skip
let mut request = base_request();
request.allowed_token_ids = Some(vec![1_000_000]);
assert!(validate_token_id_ranges(&request, usize::MAX, None).is_ok());
}
fn base_request() -> ChatCompletionRequest {
ChatCompletionRequest {
model: "Qwen/Qwen1.5-0.5B-Chat".to_string(),
rust/src/server/src/routes/openai/completions.rs
-7
View File
@@ -47,13 +47,6 @@ pub async fn completions(
let request_context = resolve_request_context(&headers, body.request_id.as_deref());
let lora_resolution = state.resolve_model_with_loras(Some(&body.model)).await;
if let Err(err) = validate::validate_token_id_ranges(
&body,
state.tokenizer_vocab_size(),
state.model_vocab_size(),
) {
return err.into_response();
}
let prepared = match prepare_completion_request(body, &lora_resolution, request_context) {
Ok(prepared) => prepared,
Err(error) => return error.into_response(),
rust/src/server/src/routes/openai/completions/validate.rs
+1 -54
View File
@@ -2,9 +2,6 @@ use vllm_text::Prompt;
use super::types::CompletionRequest;
use crate::error::{ApiError, bail_invalid_request};
use crate::routes::openai::utils::token_ids::{
validate_allowed_token_ids, validate_logit_bias, validate_prompt_token_ids,
};
/// Enforce the minimal compatibility contract for the Rust OpenAI server.
pub(super) fn validate_request_compat(
@@ -98,63 +95,13 @@ pub(super) fn validate_request_compat(
Ok(())
}
/// Reject out-of-vocab token ids, mirroring the Python input processor. A token-id
/// prompt may reference ids the engine embeds beyond either vocab alone (Qwen3
/// extra LM tokens, multimodal placeholders), so it is bounded by the union of the
/// tokenizer and model vocabularies; `allowed_token_ids` by the tokenizer vocab;
/// `logit_bias` keys by the model vocab (skipped when the model size is unknown).
pub(super) fn validate_token_id_ranges(
request: &CompletionRequest,
tokenizer_vocab_size: usize,
model_vocab_size: Option<usize>,
) -> Result<(), ApiError> {
let prompt_bound = tokenizer_vocab_size.max(model_vocab_size.unwrap_or(0));
validate_prompt_token_ids(&request.prompt, prompt_bound)?;
validate_allowed_token_ids(request.allowed_token_ids.as_deref(), tokenizer_vocab_size)?;
validate_logit_bias(
request.logit_bias.as_ref(),
model_vocab_size.unwrap_or(usize::MAX),
)
}
#[cfg(test)]
mod tests {
use serde_json::json;
use vllm_text::Prompt;
use super::{validate_request_compat, validate_token_id_ranges};
use super::validate_request_compat;
use crate::routes::openai::completions::types::CompletionRequest;
#[test]
fn validate_token_id_ranges_rejects_oob_prompt_and_params() {
// a token-id prompt below both vocabs is accepted (the engine can embed it)
let mut request = base_request();
request.prompt = Prompt::TokenIds(vec![5, 150]);
assert!(validate_token_id_ranges(&request, 100, Some(200)).is_ok());
// an id at or above the union of the two vocabs is rejected
let mut request = base_request();
request.prompt = Prompt::TokenIds(vec![5, 200]);
assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
// an id beyond the model vocab but within the (larger) tokenizer vocab is
// accepted: the engine embeds added/placeholder ids above the model vocab,
// matching the Python input processor's max(tokenizer, model) bound
let mut request = base_request();
request.prompt = Prompt::TokenIds(vec![150]);
assert!(validate_token_id_ranges(&request, 200, Some(100)).is_ok());
// falls back to the tokenizer vocab when the model size is unknown
let mut request = base_request();
request.prompt = Prompt::TokenIds(vec![150]);
assert!(validate_token_id_ranges(&request, 100, None).is_err());
// allowed_token_ids are bounded by the tokenizer vocab -> reject
let mut request = base_request();
request.allowed_token_ids = Some(vec![150]);
assert!(validate_token_id_ranges(&request, 100, Some(200)).is_err());
// unknown sizes -> skip
let mut request = base_request();
request.prompt = Prompt::TokenIds(vec![1_000_000]);
assert!(validate_token_id_ranges(&request, usize::MAX, None).is_ok());
}
fn base_request() -> CompletionRequest {
serde_json::from_value(json!({
"model": "Qwen/Qwen1.5-0.5B-Chat",
rust/src/server/src/routes/openai/utils/mod.rs
-1
View File
@@ -1,6 +1,5 @@
pub mod logprobs;
pub mod structured_outputs;
pub mod token_ids;
pub mod types;
pub mod usage;
pub mod validated_json;
rust/src/server/src/routes/openai/utils/token_ids.rs
-55
View File
@@ -1,55 +0,0 @@
use std::collections::HashMap;
use vllm_text::Prompt;
use crate::error::{ApiError, bail_invalid_request};
/// Reject token-id prompt entries at or above `bound` (the highest in-vocab id is
/// `bound - 1`).
pub(crate) fn validate_prompt_token_ids(prompt: &Prompt, bound: usize) -> Result<(), ApiError> {
if let Prompt::TokenIds(ids) = prompt
&& let Some(&bad) = ids.iter().find(|&&id| id as usize >= bound)
{
bail_invalid_request!(
param = "prompt",
"prompt contains out-of-vocab token id {bad}; vocabulary size is {bound}."
);
}
Ok(())
}
/// Reject `allowed_token_ids` entries at or above `bound`.
pub(crate) fn validate_allowed_token_ids(
allowed_token_ids: Option<&[u32]>,
bound: usize,
) -> Result<(), ApiError> {
if let Some(ids) = allowed_token_ids
&& let Some(&bad) = ids.iter().find(|&&id| id as usize >= bound)
{
bail_invalid_request!(
param = "allowed_token_ids",
"allowed_token_ids contains out-of-vocab token id {bad}; vocabulary size is {bound}."
);
}
Ok(())
}
/// Reject `logit_bias` keys at or above `bound`.
pub(crate) fn validate_logit_bias(
logit_bias: Option<&HashMap<String, f32>>,
bound: usize,
) -> Result<(), ApiError> {
if let Some(bias) = logit_bias {
for key in bias.keys() {
if let Ok(id) = key.parse::<u32>()
&& id as usize >= bound
{
bail_invalid_request!(
param = "logit_bias",
"logit_bias contains out-of-vocab token id {id}; vocabulary size is {bound}."
);
}
}
}
Ok(())
}
rust/src/server/src/state.rs
-10
View File
@@ -114,16 +114,6 @@ impl AppState {
&self.served_model_names
}
/// Tokenizer vocabulary size.
pub fn tokenizer_vocab_size(&self) -> usize {
self.chat.tokenizer_vocab_size()
}
/// Model vocabulary size, else `None`.
pub fn model_vocab_size(&self) -> Option<usize> {
self.chat.model_vocab_size()
}
/// Return base served model names plus dynamically loaded LoRA adapter
/// names.
pub async fn served_model_names_with_loras(&self) -> Vec<String> {
rust/src/text/src/backend/mod.rs
+15 -3
View File
@@ -29,10 +29,12 @@ pub struct SamplingLimits {
///
/// `-1` means allowing requests up to the model vocabulary size.
pub max_logprobs: i32,
/// Model vocabulary size from the model config.
/// Model vocabulary size from the model config, used to bound
/// `logit_bias` keys when available.
pub model_vocab_size: Option<usize>,
/// Tokenizer vocabulary size, used as a fallback when the model config does
/// not expose a vocabulary size.
/// Tokenizer vocabulary size, used to bound `allowed_token_ids` and
/// token-ID prompts.
pub tokenizer_vocab_size: usize,
}
@@ -48,6 +50,16 @@ impl SamplingLimits {
pub fn logprobs_vocab_size(&self) -> usize {
self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size)
}
/// Return the vocabulary size used to validate generated stop token IDs.
pub fn stop_token_vocab_size(&self) -> usize {
self.model_vocab_size.unwrap_or(self.tokenizer_vocab_size)
}
/// Return the union bound used to validate token-ID prompts.
pub fn prompt_token_vocab_size(&self) -> usize {
self.tokenizer_vocab_size.max(self.model_vocab_size.unwrap_or(0))
}
}
/// Minimal text-processing backend needed by `vllm-text`.
rust/src/text/src/error.rs
+3
View File
@@ -3,6 +3,7 @@ use vllm_engine_core_client::Error as EngineCoreError;
use vllm_llm::Error as LlmError;
pub use crate::lower::logprobs::LogprobsError;
pub use crate::lower::token_ids::OutOfVocabError;
#[derive(Debug, Error)]
pub enum Error {
@@ -17,6 +18,8 @@ pub enum Error {
PromptTooLong { max_model_len: u32, prompt_len: u32 },
#[error(transparent)]
Logprobs(#[from] LogprobsError),
#[error(transparent)]
OutOfVocab(#[from] OutOfVocabError),
#[error("text request stream `{request_id}` closed before terminal output")]
StreamClosedBeforeTerminalOutput { request_id: String },
#[error(transparent)]
rust/src/text/src/lib.rs
+1 -1
View File
@@ -7,7 +7,7 @@
use std::mem::take;
pub use backend::{DynTextBackend, SamplingHints, SamplingLimits, TextBackend};
pub use error::{Error, LogprobsError, Result};
pub use error::{Error, LogprobsError, OutOfVocabError, Result};
use futures::Stream;
pub use lower::{
PreparedTextRequest, lower_sampling_params, lower_text_request, resolve_max_tokens,
rust/src/text/src/lower.rs
+141 -6
View File
@@ -1,6 +1,7 @@
use std::collections::BTreeSet;
pub(crate) mod logprobs;
pub(crate) mod token_ids;
use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
use vllm_llm::GenerateRequest;
@@ -10,6 +11,7 @@ use crate::backend::{SamplingHints, SamplingLimits};
use crate::error::{Error, Result};
use crate::request::{SamplingParams, TextRequest};
use logprobs::validate_logprobs;
use token_ids::{validate_prompt_token_ids, validate_vocab_range};
/// One text request after it has been lowered into the raw generate boundary.
#[derive(Debug)]
@@ -31,6 +33,8 @@ pub fn lower_text_request(
tokenizer: &dyn Tokenizer,
) -> Result<PreparedTextRequest> {
let prompt_len = prompt_token_ids.len() as u32;
validate_prompt_token_ids(&prompt_token_ids, &sampling_limits)?;
let generate_request = GenerateRequest {
request_id: request.request_id.clone(),
prompt_token_ids,
@@ -100,7 +104,6 @@ pub fn lower_sampling_params(
vllm_xargs,
} = sampling_params;
// Validate logprobs-related fields first with runtime sampling limits first.
validate_logprobs(
logprobs,
prompt_logprobs,
@@ -139,7 +142,7 @@ pub fn lower_sampling_params(
merge_unique_token_ids(&mut stop_token_ids, extra_eos_token_ids.iter().copied());
}
Ok(EngineCoreSamplingParams {
let params = EngineCoreSamplingParams {
temperature,
top_p,
top_k,
@@ -162,7 +165,9 @@ pub fn lower_sampling_params(
logprob_token_ids,
skip_reading_prefix_cache,
extra_args: vllm_xargs,
})
};
validate_vocab_range(&params, &sampling_limits)?;
Ok(params)
}
/// Convert bad-word strings into token-ID sequences, following the Python vLLM
@@ -243,14 +248,14 @@ fn merge_unique_token_ids(
#[cfg(test)]
mod tests {
use std::collections::BTreeSet;
use std::collections::{BTreeSet, HashMap};
use serial_test::file_serial;
use super::*;
use crate::backend::hf::HfTextBackend;
use crate::backend::{SamplingHints, TextBackend as _};
use crate::error::LogprobsError;
use crate::error::{LogprobsError, OutOfVocabError};
use crate::request::{Prompt, TextRequest};
/// Stub tokenizer that returns empty token IDs — sufficient for tests that
@@ -430,6 +435,57 @@ mod tests {
.assert_debug_eq(&params);
}
#[test]
fn lower_text_request_uses_union_vocab_for_prompt_token_ids() {
lower_text_request(
sample_request(),
vec![1500],
sample_sampling_hints(),
SamplingLimits {
model_vocab_size: Some(2000),
tokenizer_vocab_size: 1000,
..sample_sampling_limits()
},
&stub_tokenizer(),
)
.expect("model vocab extends prompt token range");
lower_text_request(
sample_request(),
vec![1500],
sample_sampling_hints(),
SamplingLimits {
model_vocab_size: Some(1000),
tokenizer_vocab_size: 2000,
..sample_sampling_limits()
},
&stub_tokenizer(),
)
.expect("tokenizer vocab extends prompt token range");
let error = lower_text_request(
sample_request(),
vec![2000],
sample_sampling_hints(),
SamplingLimits {
model_vocab_size: Some(1000),
tokenizer_vocab_size: 2000,
..sample_sampling_limits()
},
&stub_tokenizer(),
)
.unwrap_err();
assert!(matches!(
error,
Error::OutOfVocab(OutOfVocabError {
parameter: "prompt",
token_ids,
vocab_size: 2000,
}) if token_ids == vec![2000]
));
}
#[tokio::test]
#[file_serial(hf_qwen3)]
async fn lower_text_request_uses_real_qwen_generation_defaults() {
@@ -747,13 +803,92 @@ mod tests {
assert!(matches!(
error,
Error::Logprobs(LogprobsError::InvalidTokenIds {
Error::OutOfVocab(OutOfVocabError {
parameter: "logprob_token_ids",
token_ids,
vocab_size: 1000,
}) if token_ids == vec![1000]
));
}
#[test]
fn lower_sampling_params_rejects_out_of_vocab_stop_token_ids() {
let error = lower_sampling_params_with_limits(
SamplingParams {
stop_token_ids: Some(vec![999, 1000]),
..Default::default()
},
sample_sampling_limits(),
)
.unwrap_err();
assert!(matches!(
error,
Error::OutOfVocab(OutOfVocabError {
parameter: "stop_token_ids",
token_ids,
vocab_size: 1000,
}) if token_ids == vec![1000]
));
}
#[test]
fn lower_sampling_params_rejects_out_of_vocab_allowed_token_ids() {
let error = lower_sampling_params_with_limits(
SamplingParams {
allowed_token_ids: Some(vec![1999, 2000]),
..Default::default()
},
sample_sampling_limits(),
)
.unwrap_err();
assert!(matches!(
error,
Error::OutOfVocab(OutOfVocabError {
parameter: "allowed_token_ids",
token_ids,
vocab_size: 2000,
}) if token_ids == vec![2000]
));
}
#[test]
fn lower_sampling_params_rejects_out_of_vocab_logit_bias() {
let error = lower_sampling_params_with_limits(
SamplingParams {
logit_bias: Some(HashMap::from([(1000, 1.0)])),
..Default::default()
},
sample_sampling_limits(),
)
.unwrap_err();
assert!(matches!(
error,
Error::OutOfVocab(OutOfVocabError {
parameter: "logit_bias",
token_ids,
vocab_size: 1000,
}) if token_ids == vec![1000]
));
}
#[test]
fn lower_sampling_params_skips_logit_bias_range_when_model_vocab_is_unknown() {
lower_sampling_params_with_limits(
SamplingParams {
logit_bias: Some(HashMap::from([(1_000_000, 1.0)])),
..Default::default()
},
SamplingLimits {
model_vocab_size: None,
..sample_sampling_limits()
},
)
.expect("logit_bias range check is skipped without model vocab size");
}
#[test]
fn lower_sampling_params_uses_generation_defaults_when_user_omits_values() {
let params = lower_sampling_params(
rust/src/text/src/lower/logprobs.rs
+3 -25
View File
@@ -26,14 +26,6 @@ pub enum LogprobsError {
requested: usize,
max_allowed: usize,
},
#[error(
"token_id(s) {token_ids:?} in logprob_token_ids contain out-of-vocab token ids. \
Vocabulary size: {vocab_size}"
)]
InvalidTokenIds {
token_ids: Vec<u32>,
vocab_size: usize,
},
#[error(
"when both logprobs and logprob_token_ids are set, logprobs must equal \
len(logprob_token_ids). Got logprobs={logprobs}, len(logprob_token_ids)={num_token_ids}."
@@ -41,8 +33,7 @@ pub enum LogprobsError {
TokenIdsMismatch { logprobs: i32, num_token_ids: usize },
}
/// Validate logprobs-related sampling parameters, returning an error if any
/// parameter is out of bounds or if the combination of parameters is invalid.
/// Validate logprobs count sampling parameters.
pub(super) fn validate_logprobs(
logprobs: Option<i32>,
prompt_logprobs: Option<i32>,
@@ -55,7 +46,7 @@ pub(super) fn validate_logprobs(
validate_logprobs_count(logprobs, max_logprobs, vocab_size, "logprobs")?;
validate_logprobs_count(prompt_logprobs, max_logprobs, vocab_size, "prompt_logprobs")?;
validate_logprob_token_ids(logprobs, logprob_token_ids, vocab_size)
validate_logprob_token_ids(logprobs, logprob_token_ids)
}
fn validate_logprobs_count(
@@ -80,10 +71,9 @@ fn validate_logprobs_count(
Ok(())
}
fn validate_logprob_token_ids(
pub(super) fn validate_logprob_token_ids(
logprobs: Option<i32>,
logprob_token_ids: Option<&[u32]>,
vocab_size: usize,
) -> Result<(), LogprobsError> {
let Some(logprob_token_ids) = logprob_token_ids else {
return Ok(());
@@ -97,18 +87,6 @@ fn validate_logprob_token_ids(
});
}
let invalid_token_ids: Vec<_> = logprob_token_ids
.iter()
.copied()
.filter(|&token_id| token_id as usize >= vocab_size)
.collect();
if !invalid_token_ids.is_empty() {
return Err(LogprobsError::InvalidTokenIds {
token_ids: invalid_token_ids,
vocab_size,
});
}
if let Some(logprobs) = logprobs
&& logprobs != n as i32
{
rust/src/text/src/lower/token_ids.rs
+101
View File
@@ -0,0 +1,101 @@
use std::result::Result;
use thiserror::Error;
use vllm_engine_core_client::protocol::EngineCoreSamplingParams;
use crate::SamplingLimits;
#[derive(Debug, Error)]
#[error(
"token_id(s) {token_ids:?} in {parameter} contain out-of-vocab token ids. \
Vocabulary size: {vocab_size}"
)]
pub struct OutOfVocabError {
pub parameter: &'static str,
pub token_ids: Vec<u32>,
pub vocab_size: usize,
}
fn validate_param(
parameter: &'static str,
token_ids: impl IntoIterator<Item = u32>,
vocab_size: usize,
) -> Result<(), OutOfVocabError> {
let invalid_token_ids: Vec<_> = token_ids
.into_iter()
.filter(|&token_id| token_id as usize >= vocab_size)
.collect();
if invalid_token_ids.is_empty() {
return Ok(());
}
Err(OutOfVocabError {
parameter,
token_ids: invalid_token_ids,
vocab_size,
})
}
/// Validate that pre-tokenized prompt IDs are within the engine-visible prompt
/// vocabulary range.
pub(crate) fn validate_prompt_token_ids(
prompt_token_ids: &[u32],
limits: &SamplingLimits,
) -> Result<(), OutOfVocabError> {
validate_param(
"prompt",
prompt_token_ids.iter().copied(),
limits.prompt_token_vocab_size(),
)
}
/// Validate that token IDs in text sampling parameters are within their
/// parameter-specific vocabulary ranges.
pub(crate) fn validate_vocab_range(
params: &EngineCoreSamplingParams,
limits: &SamplingLimits,
) -> Result<(), OutOfVocabError> {
validate_param(
"stop_token_ids",
params.stop_token_ids.iter().copied(),
limits.stop_token_vocab_size(),
)?;
if let Some(token_ids) = params.allowed_token_ids.as_deref() {
validate_param(
"allowed_token_ids",
token_ids.iter().copied(),
limits.tokenizer_vocab_size,
)?;
}
if let (Some(logit_bias), Some(vocab_size)) =
(params.logit_bias.as_ref(), limits.model_vocab_size)
{
validate_param("logit_bias", logit_bias.keys().copied(), vocab_size)?;
}
if let Some(token_ids) = params.logprob_token_ids.as_deref() {
validate_param(
"logprob_token_ids",
token_ids.iter().copied(),
limits.logprobs_vocab_size(),
)?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn validate_vocab_range_rejects_out_of_vocab_ids() {
let error = validate_param("logprob_token_ids", [5_u32, 1000, 1001], 1000).unwrap_err();
assert_eq!(error.parameter, "logprob_token_ids");
assert_eq!(error.token_ids, vec![1000, 1001]);
assert_eq!(error.vocab_size, 1000);
}
}