Skip to content

Commit 7659925

Browse files
committed
fix: improve validation and transform logic
1 parent 465294d commit 7659925

File tree

1 file changed

+29
-27
lines changed

1 file changed

+29
-27
lines changed

router/src/validation.rs

Lines changed: 29 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -424,34 +424,36 @@ impl Validation {
424424
None => None,
425425
};
426426

427-
let logit_bias = match &request.parameters.logit_bias {
428-
Some(bias) if !bias.is_empty() => {
429-
for (token_str, _) in bias.iter() {
430-
let token_id = token_str.parse::<u32>().map_err(|_| {
431-
ValidationError::LogitBiasInvalid(format!(
432-
"Token ID {} is not a valid number.",
433-
token_str
434-
))
435-
})?;
436-
437-
if token_id >= self.vocab_size {
438-
return Err(ValidationError::LogitBiasInvalid(format!(
439-
"Token ID {} is out of range. Must be between 0 and {}.",
440-
token_id,
441-
self.vocab_size - 1
442-
)));
443-
}
444-
}
427+
// Validate logit bias and convert to a vector of (token_id, bias_value)
428+
let logit_bias = request
429+
.parameters
430+
.logit_bias
431+
.as_ref()
432+
.filter(|bias_map| !bias_map.is_empty())
433+
.map(|bias_map| {
434+
bias_map
435+
.iter()
436+
.map(|(token_str, &bias_value)| {
437+
let token_id: u32 = token_str.parse().map_err(|_| {
438+
ValidationError::LogitBiasInvalid(format!(
439+
"Token ID {token_str} is not a valid number."
440+
))
441+
})?;
445442

446-
// Transform into the required format
447-
Some(
448-
bias.iter()
449-
.map(|(k, v)| (k.parse::<u32>().unwrap(), *v as f32))
450-
.collect(),
451-
)
452-
}
453-
_ => None,
454-
};
443+
if token_id >= self.vocab_size {
444+
return Err(ValidationError::LogitBiasInvalid(format!(
445+
"Token ID {token_id} is out of range (0..{}).",
446+
self.vocab_size - 1
447+
)));
448+
}
449+
450+
Ok((token_id, bias_value as f32))
451+
})
452+
.collect::<Result<Vec<_>, _>>()
453+
})
454+
// convert Option<Result<T, E>> to Result<Option<T>, E> to throw
455+
// if any of the token IDs are invalid
456+
.transpose()?;
455457

456458
let parameters = ValidParameters {
457459
temperature,

0 commit comments

Comments
 (0)