handle missing chat templates in transformers.

microsoft · Nov 26, 2024 · ffb7bdb · ffb7bdb
1 parent e7d71ab
commit ffb7bdb
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/packages/core/src/transformers.ts b/packages/core/src/transformers.ts
@@ -124,8 +124,13 @@ export const TransformersCompletion: ChatCompletionHandler = async (
     )
     const msgs: Chat = chatMessagesToTranformerMessages(messages)
     trace.detailsFenced("messages", msgs, "yaml")
+    const chatTemplate = !!generator.tokenizer.chat_template
+    const texts: Chat | string = chatTemplate
+        ? msgs
+        : msgs.map((msg) => `${msg.role}:\n${msg.content}`).join("\n\n")
+    if (chatTemplate) trace.detailsFenced("texts", texts, "markdown")
     const output = (await generator(
-        msgs,
+        texts,
         deleteUndefinedValues({
             max_new_tokens: max_tokens || 4000,
             temperature,
@@ -134,7 +139,10 @@ export const TransformersCompletion: ChatCompletionHandler = async (
         })
     )) as TextGenerationOutput
     const text = output
-        .map((msg) => (msg.generated_text.at(-1) as Message).content)
+        .map((msg) => msg.generated_text)
+        .map((msg) =>
+            typeof msg === "string" ? msg : (msg.at(-1) as Message).content
+        )
         .join("")
     trace.fence(text, "markdown")
     partialCb?.({

diff --git a/packages/core/src/types/prompt_template.d.ts b/packages/core/src/types/prompt_template.d.ts
@@ -168,6 +168,7 @@ type ModelType = OptionsOrString<
     | "alibaba:qwen2-72b-instruct"
     | "alibaba:qwen2-57b-a14b-instruct"
     | "transformers:onnx-community/Qwen2.5-0.5B-Instruct:q4"
+    | "transformers:HuggingFaceTB/SmolLM2-1.7B-Instruct:q4"
 >
 
 type ModelSmallType = OptionsOrString<