fix: handle no beginToken for token-based reasoning models

huggingface · Feb 12, 2025 · c704918 · c704918
1 parent 0d0952c
commit c704918
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 4 deletions.
diff --git a/chart/env/prod.yaml b/chart/env/prod.yaml
@@ -139,7 +139,7 @@ envVars:
         "description": "The first reasoning model from DeepSeek, distilled into a 32B dense model. Outperforms o1-mini on multiple benchmarks.",
         "reasoning": {
           "type": "tokens",
-          "beginToken": "<think>",
+          "beginToken": "",
           "endToken": "</think>"
         },
         "promptExamples": [

diff --git a/src/lib/server/models.ts b/src/lib/server/models.ts
@@ -24,7 +24,7 @@ const reasoningSchema = z.union([
 	}),
 	z.object({
 		type: z.literal("tokens"), // use beginning and end tokens that define the reasoning portion of the answer
-		beginToken: z.string(),
+		beginToken: z.string(), // empty string means the model starts in reasoning mode
 		endToken: z.string(),
 	}),
 	z.object({

diff --git a/src/lib/server/textGeneration/generate.ts b/src/lib/server/textGeneration/generate.ts
@@ -27,7 +27,10 @@ export async function* generate(
 	const startTime = new Date();
 	if (
 		model.reasoning &&
-		(model.reasoning.type === "regex" || model.reasoning.type === "summarize")
+		// if the beginToken is an empty string, the model starts in reasoning mode
+		(model.reasoning.type === "regex" ||
+			model.reasoning.type === "summarize" ||
+			(model.reasoning.type === "tokens" && model.reasoning.beginToken === ""))
 	) {
 		// if the model has reasoning in regex or summarize mode, it starts in reasoning mode
 		// and we extract the answer from the reasoning
@@ -104,7 +107,11 @@ Do not use prefixes such as Response: or Answer: when answering to the user.`,
 			} else if (model.reasoning && model.reasoning.type === "tokens") {
 				// make sure to remove the content of the reasoning buffer from
 				// the final answer to avoid duplication
-				const beginIndex = reasoningBuffer.indexOf(model.reasoning.beginToken);
+
+				// if the beginToken is an empty string, we don't need to remove anything
+				const beginIndex = model.reasoning.beginToken
+					? reasoningBuffer.indexOf(model.reasoning.beginToken)
+					: 0;
 				const endIndex = reasoningBuffer.lastIndexOf(model.reasoning.endToken);
 
 				if (beginIndex !== -1 && endIndex !== -1) {