WIP: Adding logits processors to protocol

This commit is contained in:
Noam Gat 2023-11-21 16:56:24 +02:00
parent b226e469c9
commit 4ca2c5c945
4 changed files with 12 additions and 0 deletions

View File

@ -50,6 +50,13 @@ message ClearCacheRequest {
/// Empty response /// Empty response
message ClearCacheResponse {} message ClearCacheResponse {}
message LogitsProcessorParameters {
// The name of the processor to apply
string name = 1;
// The parameters to pass to the processor
repeated string parameters = 2;
}
message NextTokenChooserParameters { message NextTokenChooserParameters {
/// exponential scaling output probability distribution /// exponential scaling output probability distribution
float temperature = 1; float temperature = 1;
@ -67,6 +74,8 @@ message NextTokenChooserParameters {
float repetition_penalty = 7; float repetition_penalty = 7;
/// token watermarking using "A Watermark for Large Language Models" /// token watermarking using "A Watermark for Large Language Models"
bool watermark = 8; bool watermark = 8;
/// Optional Logits Processors definitions
repeated LogitsProcessorParameters logits_processors = 9;
} }
message StoppingCriteriaParameters { message StoppingCriteriaParameters {

View File

@ -125,6 +125,7 @@ impl Client {
seed: 0, seed: 0,
repetition_penalty: 1.2, repetition_penalty: 1.2,
watermark: true, watermark: true,
logits_processors: vec![],
}), }),
stopping_parameters: Some(StoppingCriteriaParameters { stopping_parameters: Some(StoppingCriteriaParameters {
max_new_tokens: max_total_tokens - truncate, max_new_tokens: max_total_tokens - truncate,

View File

@ -44,6 +44,7 @@ impl Health {
seed: 0, seed: 0,
repetition_penalty: 1.0, repetition_penalty: 1.0,
watermark: false, watermark: false,
logits_processors: vec![],
}), }),
stopping_parameters: Some(StoppingCriteriaParameters { stopping_parameters: Some(StoppingCriteriaParameters {
max_new_tokens: 1, max_new_tokens: 1,

View File

@ -279,6 +279,7 @@ impl Validation {
do_sample, do_sample,
seed, seed,
watermark, watermark,
logits_processors: vec![],
}; };
let stopping_parameters = StoppingCriteriaParameters { let stopping_parameters = StoppingCriteriaParameters {
max_new_tokens, max_new_tokens,