From 20b05bc8ba2465c355115d9dee538f64aaf27ffb Mon Sep 17 00:00:00 2001 From: marcusdunn Date: Wed, 9 Aug 2023 16:31:01 -0700 Subject: [PATCH] added logit_bias param to REST and GRPC --- benchmark/src/lib.rs | 1 + proto/generate.proto | 10 ++++++++++ router/client/src/client.rs | 1 + router/client/src/lib.rs | 2 +- router/src/health.rs | 1 + router/src/lib.rs | 7 ++++++- router/src/queue.rs | 1 + router/src/validation.rs | 7 ++++++- 8 files changed, 27 insertions(+), 3 deletions(-) diff --git a/benchmark/src/lib.rs b/benchmark/src/lib.rs index fcad400c..d6433612 100644 --- a/benchmark/src/lib.rs +++ b/benchmark/src/lib.rs @@ -42,6 +42,7 @@ pub async fn run( seed: 0, repetition_penalty: repetition_penalty.unwrap_or(1.0), watermark, + logit_bias: vec![], }; // Initialize terminal properties diff --git a/proto/generate.proto b/proto/generate.proto index 57d79bca..5496a803 100644 --- a/proto/generate.proto +++ b/proto/generate.proto @@ -66,6 +66,16 @@ message NextTokenChooserParameters { float repetition_penalty = 7; /// token watermarking using "A Watermark for Large Language Models" bool watermark = 8; + /// bias towards certain token sequences + repeated LogitBias logit_bias = 9; +} + +/// a token sequence and its bias +message LogitBias { + /// string to bias towards (may be more than one token) + string string = 1; + /// bias for the string + float bias = 2; } message StoppingCriteriaParameters { diff --git a/router/client/src/client.rs b/router/client/src/client.rs index 7753f307..407f3dd1 100644 --- a/router/client/src/client.rs +++ b/router/client/src/client.rs @@ -124,6 +124,7 @@ impl Client { seed: 0, repetition_penalty: 1.2, watermark: true, + logit_bias: vec![], }), stopping_parameters: Some(StoppingCriteriaParameters { max_new_tokens: 2, diff --git a/router/client/src/lib.rs b/router/client/src/lib.rs index f334be21..945f7d60 100644 --- a/router/client/src/lib.rs +++ b/router/client/src/lib.rs @@ -10,7 +10,7 @@ pub use pb::generate::v1::HealthResponse; pub use pb::generate::v1::InfoResponse as ShardInfo; pub use pb::generate::v1::{ Batch, CachedBatch, FinishReason, GeneratedText, Generation, NextTokenChooserParameters, - PrefillTokens, Request, StoppingCriteriaParameters, + PrefillTokens, Request, StoppingCriteriaParameters, LogitBias }; pub use sharded_client::ShardedClient; use thiserror::Error; diff --git a/router/src/health.rs b/router/src/health.rs index a3cacdcd..f6a5cfc4 100644 --- a/router/src/health.rs +++ b/router/src/health.rs @@ -44,6 +44,7 @@ impl Health { seed: 0, repetition_penalty: 1.0, watermark: false, + logit_bias: vec![], }), stopping_parameters: Some(StoppingCriteriaParameters { max_new_tokens: 1, diff --git a/router/src/lib.rs b/router/src/lib.rs index 7dff7a11..9825b9df 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -5,6 +5,7 @@ mod queue; pub mod server; mod validation; +use std::collections::BTreeMap; use infer::Infer; use queue::{Entry, Queue}; use serde::{Deserialize, Serialize}; @@ -113,7 +114,7 @@ pub(crate) struct GenerateParameters { #[schema(nullable = true, default = "null", example = false)] pub return_full_text: Option, #[serde(default)] - #[schema(inline, max_items = 4, example = json ! (["photographer"]))] + #[schema(inline, max_items = 4, example = json!(["photographer"]))] pub stop: Vec, #[serde(default)] #[schema(nullable = true, default = "null", example = "null")] @@ -135,6 +136,9 @@ pub(crate) struct GenerateParameters { example = "null" )] pub seed: Option, + #[serde(default)] + #[schema(default=json!({}), example=json!({"hello": 0.5}))] + pub logit_bias: BTreeMap } fn default_max_new_tokens() -> u32 { @@ -158,6 +162,7 @@ fn default_parameters() -> GenerateParameters { details: false, decoder_input_details: false, seed: None, + logit_bias: BTreeMap::new(), } } diff --git a/router/src/queue.rs b/router/src/queue.rs index 2d8d6d1c..3ef235bb 100644 --- a/router/src/queue.rs +++ b/router/src/queue.rs @@ -322,6 +322,7 @@ mod tests { seed: 0, repetition_penalty: 0.0, watermark: false, + logit_bias: vec![], }, stopping_parameters: StoppingCriteriaParameters { ignore_eos_token: false, diff --git a/router/src/validation.rs b/router/src/validation.rs index be835bf0..041d88b7 100644 --- a/router/src/validation.rs +++ b/router/src/validation.rs @@ -2,7 +2,7 @@ use crate::validation::ValidationError::{BestOfSampling, BestOfSeed, EmptyInput}; use crate::{GenerateParameters, GenerateRequest}; use rand::{thread_rng, Rng}; -use text_generation_client::{NextTokenChooserParameters, StoppingCriteriaParameters}; +use text_generation_client::{NextTokenChooserParameters, StoppingCriteriaParameters, LogitBias}; use thiserror::Error; use tokenizers::tokenizer::Tokenizer; use tokenizers::TruncationDirection; @@ -142,6 +142,7 @@ impl Validation { seed, watermark, decoder_input_details, + logit_bias, .. } = request.parameters; @@ -238,6 +239,9 @@ impl Validation { .validate_input(request.inputs, truncate, max_new_tokens) .await?; + // transform logit_bias (received as a map) into a vector of LogitBias + let logit_bias = logit_bias.into_iter().map(|(string, bias)| LogitBias { string, bias, }).collect(); + let parameters = NextTokenChooserParameters { temperature, repetition_penalty, @@ -247,6 +251,7 @@ impl Validation { do_sample, seed, watermark, + logit_bias, }; let stopping_parameters = StoppingCriteriaParameters { max_new_tokens,