diff --git a/router/src/infer.rs b/router/src/infer.rs index eb489b5f..ec91b54e 100644 --- a/router/src/infer.rs +++ b/router/src/infer.rs @@ -140,19 +140,9 @@ impl Infer { /// Apply the chat template to the chat request #[instrument(skip_all)] - pub(crate) fn apply_chat_template(&self, chat: ChatRequest) -> Result { - let mut env = minijinja::Environment::new(); - let chat_template = self - .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) + self.tokenizer_config + .apply_chat_template(chat) .map_err(InferError::TemplateError) } diff --git a/router/src/lib.rs b/router/src/lib.rs index 05d5e5a2..3ffc1275 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -26,6 +26,22 @@ pub struct HubTokenizerConfig { pub chat_template: Option, } +impl HubTokenizerConfig { + /// Apply the chat template to the chat request + pub(crate) fn apply_chat_template( + &self, + chat: ChatRequest, + ) -> Result { + 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)] pub struct Info { /// Model info