fix: prefer apply_chat_template logic in HubTokenizerConfig struct

This commit is contained in:
drbh 2024-01-09 12:27:01 -05:00
parent 446b3b6af7
commit adad67e3d0
2 changed files with 18 additions and 12 deletions

View File

@ -140,19 +140,9 @@ impl Infer {
/// Apply the chat template to the chat request /// Apply the chat template to the chat request
#[instrument(skip_all)] #[instrument(skip_all)]
pub(crate) fn apply_chat_template(&self, chat: ChatRequest) -> Result<String, InferError> { pub(crate) fn apply_chat_template(&self, chat: ChatRequest) -> Result<String, InferError> {
let mut env = minijinja::Environment::new(); self.tokenizer_config
let chat_template = self .apply_chat_template(chat)
.tokenizer_config
.chat_template
.as_ref()
.ok_or_else(|| {
InferError::TemplateError(minijinja::ErrorKind::TemplateNotFound.into())
})?;
env.add_template("_", chat_template)?;
env.get_template("_")?
.render(chat)
.map_err(InferError::TemplateError) .map_err(InferError::TemplateError)
} }

View File

@ -26,6 +26,22 @@ pub struct HubTokenizerConfig {
pub chat_template: Option<String>, pub chat_template: Option<String>,
} }
impl HubTokenizerConfig {
/// Apply the chat template to the chat request
pub(crate) fn apply_chat_template(
&self,
chat: ChatRequest,
) -> Result<String, minijinja::Error> {
let mut env = minijinja::Environment::new();
let chat_template = self
.chat_template
.as_ref()
.ok_or(minijinja::ErrorKind::TemplateNotFound)?;
env.add_template("_", chat_template)?;
env.get_template("_")?.render(chat)
}
}
#[derive(Clone, Debug, Serialize, ToSchema)] #[derive(Clone, Debug, Serialize, ToSchema)]
pub struct Info { pub struct Info {
/// Model info /// Model info