huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 20:34:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: update create_post_processor logic for token type

Browse Source
This commit is contained in:
drbh 2024-06-28 15:07:50 +00:00
parent c326ffdac0
commit 8885688630
2 changed files with 13 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
router/src/lib.rs
View File

@ -90,18 +90,6 @@ impl TokenizerConfigToken {
}
}
impl From<TokenizerConfigToken> for String {
fn from(token: TokenizerConfigToken) -> Self {
token.as_str().to_string()
}
}
impl From<String> for TokenizerConfigToken {
fn from(s: String) -> Self {
TokenizerConfigToken::String(s)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "processor_class")]
pub enum HubPreprocessorConfig {

25
router/src/main.rs
View File

@ -553,11 +553,11 @@ pub fn create_post_processor(
if add_bos_token {
if let Some(bos) = bos_token {
let bos_token_id = tokenizer
.token_to_id(bos)
.token_to_id(bos.as_str())
.expect("Should have found the bos token id");
special_tokens.push((bos.clone(), bos_token_id));
single.push(format!("{}:0", bos));
pair.push(format!("{}:0", bos));
special_tokens.push((bos.as_str(), bos_token_id));
single.push(format!("{}:0", bos.as_str()));
pair.push(format!("{}:0", bos.as_str()));
}
}
@ -567,17 +567,17 @@ pub fn create_post_processor(
if add_eos_token {
if let Some(eos) = eos_token {
let eos_token_id = tokenizer
.token_to_id(eos)
.token_to_id(eos.as_str())
.expect("Should have found the eos token id");
special_tokens.push((eos.clone(), eos_token_id));
single.push(format!("{}:0", eos));
pair.push(format!("{}:0", eos));
special_tokens.push((eos.as_str(), eos_token_id));
single.push(format!("{}:0", eos.as_str()));
pair.push(format!("{}:0", eos.as_str()));
}
}
if add_bos_token {
if let Some(bos) = bos_token {
single.push(format!("{}:1", bos));
single.push(format!("{}:1", bos.as_str()));
}
}
@ -585,7 +585,7 @@ pub fn create_post_processor(
if add_eos_token {
if let Some(eos) = eos_token {
pair.push(format!("{}:1", eos));
pair.push(format!("{}:1", eos.as_str()));
}
}
@ -611,14 +611,15 @@ enum RouterError {
#[cfg(test)]
mod tests {
use super::*;
use text_generation_router::TokenizerConfigToken;
#[test]
fn test_create_post_processor() {
let tokenizer_config = HubTokenizerConfig {
add_bos_token: None,
add_eos_token: None,
bos_token: Some("<s>".to_string()),
eos_token: Some("</s>".to_string()),
bos_token: Some(TokenizerConfigToken::String("<s>".to_string())),
eos_token: Some(TokenizerConfigToken::String("</s>".to_string())),
chat_template: None,
tokenizer_class: None,
completion_template: None,