diff --git a/llama/addon/AddonContext.cpp b/llama/addon/AddonContext.cpp index 9427e8ff..113a3cf3 100644 --- a/llama/addon/AddonContext.cpp +++ b/llama/addon/AddonContext.cpp @@ -235,7 +235,7 @@ class AddonContextSampleTokenWorker : public Napi::AsyncWorker { llama_token_data_array cur_p; sampler->sample(ctx->ctx, batchLogitIndex, cur_p, returnProbabilities || returnConfidence); - if (!(cur_p.selected >= 0 && cur_p.selected < (int32_t)cur_p.size)) { + if (cur_p.size == 0 || !(cur_p.selected >= 0 && cur_p.selected < (int32_t)cur_p.size)) { no_output = true; return; } diff --git a/llama/addon/AddonSampler.cpp b/llama/addon/AddonSampler.cpp index 35828e23..8389bd9f 100644 --- a/llama/addon/AddonSampler.cpp +++ b/llama/addon/AddonSampler.cpp @@ -173,6 +173,10 @@ void AddonSampler::acceptToken(llama_token token) { void AddonSampler::sample(struct llama_context* llamaContext, int32_t batchLogitIndex, llama_token_data_array& curP, bool forceGrammar) { setTokenCandidates(llamaContext, batchLogitIndex, curP); + if (curP.size == 0) { + return; + } + if (forceGrammar && grammarEvaluationState != nullptr && grammarEvaluationState->sampler != nullptr) { llama_sampler_apply(grammarEvaluationState->sampler, &curP); llama_sampler_apply(chain, &curP); @@ -287,7 +291,7 @@ Napi::Value AddonSampler::ApplyConfig(const Napi::CallbackInfo& info) { temperatureSampler = nullptr; } - if (temperatureSampler_temperature <= 0) { + if (temperatureSampler_temperature <= 0.0f) { greedySampler = llama_sampler_init_greedy(); } else { temperatureSampler = llama_sampler_init_temp(temperatureSampler_temperature); @@ -321,7 +325,7 @@ Napi::Value AddonSampler::ApplyConfig(const Napi::CallbackInfo& info) { minPSampler = nullptr; } - if (minPSampler_minP != 0) { + if (minPSampler_minP > 0.0f) { minPSampler = llama_sampler_init_min_p(minPSampler_minP, min_keep); } } @@ -366,7 +370,7 @@ Napi::Value AddonSampler::ApplyConfig(const Napi::CallbackInfo& info) { topPSampler = nullptr; } - if (topPSampler_topP >= 1) { + if (topPSampler_topP >= 0.0f && topPSampler_topP <= 1.0f) { topPSampler = llama_sampler_init_top_p(topPSampler_topP, min_keep); } } diff --git a/llama/addon/AddonSampler.h b/llama/addon/AddonSampler.h index 77735d04..99cac374 100644 --- a/llama/addon/AddonSampler.h +++ b/llama/addon/AddonSampler.h @@ -24,7 +24,7 @@ class AddonSampler : public Napi::ObjectWrap { int topKSampler_topK = 0; llama_sampler * topPSampler = nullptr; - float topPSampler_topP = 0.0f; // Top p sampling >=1.0 = disabled + float topPSampler_topP = 1.0f; // Top p sampling >=1.0 = disabled llama_sampler * xtcSampler = nullptr; float xtcSampler_probability = 0; diff --git a/package-lock.json b/package-lock.json index 9e7338e3..d26b965c 100644 --- a/package-lock.json +++ b/package-lock.json @@ -10,7 +10,7 @@ "hasInstallScript": true, "license": "MIT", "dependencies": { - "@huggingface/jinja": "^0.5.6", + "@huggingface/jinja": "^0.5.9", "async-retry": "^1.3.3", "bytes": "^3.1.2", "chalk": "^5.6.2", @@ -26,12 +26,12 @@ "lifecycle-utils": "^4.0.1", "log-symbols": "^7.0.1", "nanoid": "^5.1.6", - "node-addon-api": "^8.6.0", + "node-addon-api": "^8.9.1", "ora": "^9.3.0", "pretty-ms": "^9.3.0", "proper-lockfile": "^4.1.2", "semver": "^7.7.1", - "simple-git": "^3.33.0", + "simple-git": "^3.36.0", "slice-ansi": "^8.0.0", "stdout-update": "^4.0.1", "strip-ansi": "^7.2.0", @@ -52,7 +52,7 @@ "@nolebase/vitepress-plugin-og-image": "^2.18.2", "@resvg/resvg-js": "^2.6.2", "@semantic-release/exec": "^7.1.0", - "@semantic-release/github": "^12.0.6", + "@semantic-release/github": "^12.0.9", "@semantic-release/npm": "^13.1.5", "@shikijs/vitepress-twoslash": "^3.22.0", "@stylistic/eslint-plugin": "^5.8.0", @@ -78,7 +78,7 @@ "husky": "^9.1.7", "rehype": "^13.0.2", "rimraf": "^6.1.3", - "semantic-release": "^25.0.3", + "semantic-release": "^25.0.8", "sharp": "^0.34.5", "tslib": "^2.8.1", "typedoc": "^0.28.17", @@ -1248,9 +1248,9 @@ } }, "node_modules/@huggingface/jinja": { - "version": "0.5.6", - "resolved": "https://registry.npmjs.org/@huggingface/jinja/-/jinja-0.5.6.tgz", - "integrity": "sha512-MyMWyLnjqo+KRJYSH7oWNbsOn5onuIvfXYPcc0WOGxU0eHUV7oAYUoQTl2BMdu7ml+ea/bu11UM+EshbeHwtIA==", + "version": "0.5.9", + "resolved": "https://registry.npmjs.org/@huggingface/jinja/-/jinja-0.5.9.tgz", + "integrity": "sha512-uWTG+l3VJRsl7EXxYizuL3P+cCPoc3cRqbWWRcQN0FhejRfbdq0RNhCmbY/YDtnTcz9icdLYuLDjsnz4d8JMuw==", "license": "MIT", "engines": { "node": ">=18" @@ -3522,9 +3522,9 @@ } }, "node_modules/@semantic-release/github": { - "version": "12.0.6", - "resolved": "https://registry.npmjs.org/@semantic-release/github/-/github-12.0.6.tgz", - "integrity": "sha512-aYYFkwHW3c6YtHwQF0t0+lAjlU+87NFOZuH2CvWFD0Ylivc7MwhZMiHOJ0FMpIgPpCVib/VUAcOwvrW0KnxQtA==", + "version": "12.0.9", + "resolved": "https://registry.npmjs.org/@semantic-release/github/-/github-12.0.9.tgz", + "integrity": "sha512-ODIqb0V3QqndipryEEiaBxUQCFjvv7Oese5Dt4omMGa60YRNEW0Sx3K+zri0uac2Y6S9nOlMehciWIzvvRCTGQ==", "dev": true, "license": "MIT", "dependencies": { @@ -3536,8 +3536,8 @@ "aggregate-error": "^5.0.0", "debug": "^4.3.4", "dir-glob": "^3.0.1", - "http-proxy-agent": "^7.0.0", - "https-proxy-agent": "^7.0.0", + "http-proxy-agent": "^9.0.0", + "https-proxy-agent": "^9.0.0", "issue-parser": "^7.0.0", "lodash-es": "^4.17.21", "mime": "^4.0.0", @@ -4008,6 +4008,21 @@ "dev": true, "license": "MIT" }, + "node_modules/@simple-git/args-pathspec": { + "version": "1.0.3", + "resolved": "https://registry.npmjs.org/@simple-git/args-pathspec/-/args-pathspec-1.0.3.tgz", + "integrity": "sha512-ngJMaHlsWDTfjyq9F3VIQ8b7NXbBLq5j9i5bJ6XLYtD6qlDXT7fdKY2KscWWUF8t18xx052Y/PUO1K1TRc9yKA==", + "license": "MIT" + }, + "node_modules/@simple-git/argv-parser": { + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/@simple-git/argv-parser/-/argv-parser-1.1.1.tgz", + "integrity": "sha512-Q9lBcfQ+VQCpQqGJFHe5yooOS5hGdLFFbJ5R+R5aDsnkPCahtn1hSkMcORX65J2Z5lxSkD0lQorMsncuBQxYUw==", + "license": "MIT", + "dependencies": { + "@simple-git/args-pathspec": "^1.0.3" + } + }, "node_modules/@simple-libs/child-process-utils": { "version": "1.0.2", "resolved": "https://registry.npmjs.org/@simple-libs/child-process-utils/-/child-process-utils-1.0.2.tgz", @@ -5513,13 +5528,13 @@ } }, "node_modules/agent-base": { - "version": "7.1.4", - "resolved": "https://registry.npmjs.org/agent-base/-/agent-base-7.1.4.tgz", - "integrity": "sha512-MnA+YT8fwfJPgBx3m60MNqakm30XOkyIoH1y6huTQvC0PwZG7ki8NacLBcrPbNoo8vEZy7Jpuk7+jMO+CUovTQ==", + "version": "9.0.0", + "resolved": "https://registry.npmjs.org/agent-base/-/agent-base-9.0.0.tgz", + "integrity": "sha512-TQf59BsZnytt8GdJKLPfUZ54g/iaUL2OWDSFCCvMOhsHduDQxO8xC4PNeyIkVcA5KwL2phPSv0douC0fgWzmnA==", "dev": true, "license": "MIT", "engines": { - "node": ">= 14" + "node": ">= 20" } }, "node_modules/aggregate-error": { @@ -9642,17 +9657,18 @@ "license": "BSD-2-Clause" }, "node_modules/http-proxy-agent": { - "version": "7.0.2", - "resolved": "https://registry.npmjs.org/http-proxy-agent/-/http-proxy-agent-7.0.2.tgz", - "integrity": "sha512-T1gkAiYYDWYx3V5Bmyu7HcfcvL7mUrTWiM6yOfa3PIphViJ/gFPbvidQ+veqSOHci/PxBcDabeUNCzpOODJZig==", + "version": "9.1.0", + "resolved": "https://registry.npmjs.org/http-proxy-agent/-/http-proxy-agent-9.1.0.tgz", + "integrity": "sha512-2NxoveTT58mjYT4n3RPTEfCZGLMbidoO8XEieXfpSYxu+PQJ1qpx4ypwH6N+uF9twBPIvRRgvkvW5HUTYWENig==", "dev": true, "license": "MIT", "dependencies": { - "agent-base": "^7.1.0", - "debug": "^4.3.4" + "agent-base": "9.0.0", + "debug": "^4.3.4", + "proxy-agent-negotiate": "1.1.0" }, "engines": { - "node": ">= 14" + "node": ">= 20" } }, "node_modules/http2-wrapper": { @@ -9670,17 +9686,18 @@ } }, "node_modules/https-proxy-agent": { - "version": "7.0.6", - "resolved": "https://registry.npmjs.org/https-proxy-agent/-/https-proxy-agent-7.0.6.tgz", - "integrity": "sha512-vK9P5/iUfdl95AI+JVyUuIcVtd4ofvtrOr3HNtM2yxC9bnMbEdp3x01OhQNnjb8IJYi38VlTE3mBXwcfvywuSw==", + "version": "9.1.0", + "resolved": "https://registry.npmjs.org/https-proxy-agent/-/https-proxy-agent-9.1.0.tgz", + "integrity": "sha512-ag87y7cJJ9/3+GxFr8Oy4O5faDsGRGnBGsJj/YjOSsSx/5eadKLYTMPlzuR6obgoCDDm0abAAZitXXQkMOPSpA==", "dev": true, "license": "MIT", "dependencies": { - "agent-base": "^7.1.2", - "debug": "4" + "agent-base": "9.0.0", + "debug": "^4.3.4", + "proxy-agent-negotiate": "1.1.0" }, "engines": { - "node": ">= 14" + "node": ">= 20" } }, "node_modules/human-signals": { @@ -12508,9 +12525,9 @@ "license": "MIT" }, "node_modules/node-addon-api": { - "version": "8.6.0", - "resolved": "https://registry.npmjs.org/node-addon-api/-/node-addon-api-8.6.0.tgz", - "integrity": "sha512-gBVjCaqDlRUk0EwoPNKzIr9KkS9041G/q31IBShPs1Xz6UTA+EXdZADbzqAJQrpDRq71CIMnOP5VMut3SL0z5Q==", + "version": "8.9.1", + "resolved": "https://registry.npmjs.org/node-addon-api/-/node-addon-api-8.9.1.tgz", + "integrity": "sha512-4eUQWVPCUUUiBjLnHS3cXWeC6ryoPUc0U3rP7IuzapoGbzMqd/r6KKO0clr0b+snQhsrueFEhCZDdK+LK7hxKg==", "license": "MIT", "engines": { "node": "^18 || ^20 || >= 21" @@ -15472,6 +15489,24 @@ "dev": true, "license": "ISC" }, + "node_modules/proxy-agent-negotiate": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/proxy-agent-negotiate/-/proxy-agent-negotiate-1.1.0.tgz", + "integrity": "sha512-N8IBcM3UgCVzz2L2Lqv8DVntDnnC8/hiV4nEDUPkqq72TPUgYWjQc+bdZlBPZK9LzPAvOY//gAt0S0DApoOXWQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">= 20" + }, + "peerDependencies": { + "kerberos": "^2.0.0" + }, + "peerDependenciesMeta": { + "kerberos": { + "optional": true + } + } + }, "node_modules/prr": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/prr/-/prr-1.0.1.tgz", @@ -16280,9 +16315,9 @@ } }, "node_modules/semantic-release": { - "version": "25.0.3", - "resolved": "https://registry.npmjs.org/semantic-release/-/semantic-release-25.0.3.tgz", - "integrity": "sha512-WRgl5GcypwramYX4HV+eQGzUbD7UUbljVmS+5G1uMwX/wLgYuJAxGeerXJDMO2xshng4+FXqCgyB5QfClV6WjA==", + "version": "25.0.8", + "resolved": "https://registry.npmjs.org/semantic-release/-/semantic-release-25.0.8.tgz", + "integrity": "sha512-w/iZ0bur36rKffXZYmIUmy068eoBY3Ij1DCCddx2JwWEM5Tg+eU9ld/E9qSInVvPASyyR2Ln/XGfQ9OZrMlhtw==", "dev": true, "license": "MIT", "dependencies": { @@ -16871,13 +16906,15 @@ } }, "node_modules/simple-git": { - "version": "3.33.0", - "resolved": "https://registry.npmjs.org/simple-git/-/simple-git-3.33.0.tgz", - "integrity": "sha512-D4V/tGC2sjsoNhoMybKyGoE+v8A60hRawKQ1iFRA1zwuDgGZCBJ4ByOzZ5J8joBbi4Oam0qiPH+GhzmSBwbJng==", + "version": "3.36.0", + "resolved": "https://registry.npmjs.org/simple-git/-/simple-git-3.36.0.tgz", + "integrity": "sha512-cGQjLjK8bxJw4QuYT7gxHw3/IouVESbhahSsHrX97MzCL1gu2u7oy38W6L2ZIGECEfIBG4BabsWDPjBxJENv9Q==", "license": "MIT", "dependencies": { "@kwsites/file-exists": "^1.1.1", "@kwsites/promise-deferred": "^1.1.1", + "@simple-git/args-pathspec": "^1.0.3", + "@simple-git/argv-parser": "^1.1.0", "debug": "^4.4.0" }, "funding": { diff --git a/package.json b/package.json index c148af0a..c3cf7308 100644 --- a/package.json +++ b/package.json @@ -144,7 +144,7 @@ "@nolebase/vitepress-plugin-og-image": "^2.18.2", "@resvg/resvg-js": "^2.6.2", "@semantic-release/exec": "^7.1.0", - "@semantic-release/github": "^12.0.6", + "@semantic-release/github": "^12.0.9", "@semantic-release/npm": "^13.1.5", "@shikijs/vitepress-twoslash": "^3.22.0", "@stylistic/eslint-plugin": "^5.8.0", @@ -170,7 +170,7 @@ "husky": "^9.1.7", "rehype": "^13.0.2", "rimraf": "^6.1.3", - "semantic-release": "^25.0.3", + "semantic-release": "^25.0.8", "sharp": "^0.34.5", "tslib": "^2.8.1", "typedoc": "^0.28.17", @@ -186,7 +186,7 @@ "zx": "^8.8.5" }, "dependencies": { - "@huggingface/jinja": "^0.5.6", + "@huggingface/jinja": "^0.5.9", "async-retry": "^1.3.3", "bytes": "^3.1.2", "chalk": "^5.6.2", @@ -202,12 +202,12 @@ "lifecycle-utils": "^4.0.1", "log-symbols": "^7.0.1", "nanoid": "^5.1.6", - "node-addon-api": "^8.6.0", + "node-addon-api": "^8.9.1", "ora": "^9.3.0", "pretty-ms": "^9.3.0", "proper-lockfile": "^4.1.2", "semver": "^7.7.1", - "simple-git": "^3.33.0", + "simple-git": "^3.36.0", "slice-ansi": "^8.0.0", "stdout-update": "^4.0.1", "strip-ansi": "^7.2.0", diff --git a/src/ChatWrapper.ts b/src/ChatWrapper.ts index 71d7efd5..543b9f5d 100644 --- a/src/ChatWrapper.ts +++ b/src/ChatWrapper.ts @@ -7,7 +7,8 @@ import {LlamaText, SpecialTokensText} from "./utils/LlamaText.js"; import {ChatModelFunctionsDocumentationGenerator} from "./chatWrappers/utils/ChatModelFunctionsDocumentationGenerator.js"; import {jsonDumps} from "./chatWrappers/utils/jsonDumps.js"; import {defaultChatSystemPrompt} from "./config.js"; -import {getChatWrapperSegmentDefinition} from "./utils/getChatWrapperSegmentDefinition.js"; +import {getStandardizedChatWrapperSegmentDefinition} from "./utils/getStandardizedChatWrapperSegmentDefinition.js"; +import {replaceRegularTextInLlamaText} from "./chatWrappers/utils/replaceRegularTextInLlamaText.js"; import type {JinjaTemplateChatWrapperOptions} from "./chatWrappers/generic/JinjaTemplateChatWrapper.js"; export abstract class ChatWrapper { @@ -19,7 +20,9 @@ export abstract class ChatWrapper { prefix: "||call: ", paramsPrefix: LlamaText(new SpecialTokensText("(")), suffix: LlamaText(new SpecialTokensText(")")), - emptyCallParamsPlaceholder: "" + emptyCallParamsPlaceholder: "", + + prefixAlternateMatches: ["||call:"] }, result: { prefix: LlamaText(new SpecialTokensText("\n"), "||result: "), @@ -117,7 +120,7 @@ export abstract class ChatWrapper { : jsonDumps(emptyCallParamsPlaceholder) : jsonDumps(params) ), - this.settings.functions.call.suffix + replaceRegularTextInLlamaText(this.settings.functions.call.suffix, "{{functionName}}", name) ]); } @@ -174,11 +177,11 @@ export abstract class ChatWrapper { res.push(LlamaText(this.settings.segments.closeAllSegments)); } else if (needsToAddSegmentReminder && segmentStack.length > 0 && this.settings.segments?.reiterateStackAfterFunctionCalls) { for (const segmentType of segmentStack) { - const segmentDefinition = getChatWrapperSegmentDefinition(this.settings, segmentType); - if (segmentDefinition == null) + const standardizedSegmentDefinition = getStandardizedChatWrapperSegmentDefinition(this.settings, segmentType); + if (standardizedSegmentDefinition?.prefix == null) continue; - res.push(LlamaText(segmentDefinition.prefix)); + res.push(LlamaText(standardizedSegmentDefinition.prefix)); } } }; @@ -193,7 +196,7 @@ export abstract class ChatWrapper { } else if (isChatModelResponseSegment(response)) { addFunctionCalls(); - const segmentDefinition = getChatWrapperSegmentDefinition(this.settings, response.segmentType); + const standardizedSegmentDefinition = getStandardizedChatWrapperSegmentDefinition(this.settings, response.segmentType); if (response.raw != null && useRawValues) res.push(LlamaText.fromJSON(response.raw)); else @@ -201,22 +204,24 @@ export abstract class ChatWrapper { LlamaText([ (segmentStack.length > 0 && segmentStack.at(-1) === response.segmentType) ? "" - : segmentDefinition?.prefix ?? "", + : (standardizedSegmentDefinition?.prefix ?? ""), response.text, response.ended - ? (segmentDefinition?.suffix ?? "") + ? (standardizedSegmentDefinition?.suffix ?? "") : "" ]) ); - lastSegmentEndedWithoutSuffix = response.ended && segmentDefinition?.suffix == null; + lastSegmentEndedWithoutSuffix = response.ended && standardizedSegmentDefinition?.suffix == null; if (!response.ended && segmentStack.at(-1) !== response.segmentType) segmentStack.push(response.segmentType); else if (response.ended && segmentStack.at(-1) === response.segmentType) { segmentStack.pop(); - if (segmentStack.length === 0 && segmentDefinition?.suffix == null && this.settings.segments?.closeAllSegments != null) + if (segmentStack.length === 0 && standardizedSegmentDefinition?.suffix == null && + this.settings.segments?.closeAllSegments != null + ) res.push(LlamaText(this.settings.segments.closeAllSegments)); } @@ -309,6 +314,6 @@ export type ChatWrapperJinjaMatchConfiguration = A [ testConfig: FirstItemOfTupleOrFallback, object>, applyConfig: FirstItemOfTupleOrFallback, object>, - testJinjaChatWrapperOptions?: JinjaTemplateChatWrapperOptions + testJinjaChatWrapperOptions?: Omit ] >; diff --git a/src/chatWrappers/MuseChatWrapper.ts b/src/chatWrappers/MuseChatWrapper.ts new file mode 100644 index 00000000..06f70c54 --- /dev/null +++ b/src/chatWrappers/MuseChatWrapper.ts @@ -0,0 +1,566 @@ +import {ChatWrapper, ChatWrapperJinjaMatchConfiguration} from "../ChatWrapper.js"; +import { + ChatHistoryItem, + ChatModelFunctionCall, ChatModelFunctions, ChatModelResponse, ChatWrapperGenerateContextStateOptions, + ChatWrapperGeneratedContextState, ChatWrapperGenerateInitialHistoryOptions, ChatWrapperSettings +} from "../types.js"; +import {LlamaText, SpecialToken, SpecialTokensText} from "../utils/LlamaText.js"; +import {optionsMatrix} from "../utils/optionsMatrix.js"; +import {jsonDumps} from "./utils/jsonDumps.js"; + +// source: https://dev.meta.ai/docs/muse-glimmer/prompting +// https://huggingface.co/meta-models/Muse-Glimmer-30B/blob/main/chat_template.jinja +export class MuseChatWrapper extends ChatWrapper { + public readonly wrapperName: string = "Muse"; + + public readonly reasoningStrength: "xhigh" | "high" | "medium" | "low"; + public readonly keepOnlyLastThought: boolean; + public readonly knowledgeCutoff: Date | (() => Date) | null; + public readonly todayDate: Date | (() => Date) | null; + + /** @internal */ private readonly _systemMessageInSpecialTokensText: boolean; + + public override readonly settings: ChatWrapperSettings = { + supportsSystemMessages: true, + functions: { + call: { + optionalPrefixSpace: false, + prefix: LlamaText(new SpecialTokensText(" to=")), + paramsPrefix: LlamaText([ + new SpecialTokensText("<|message|>\n"), + new SpecialTokensText('\n'), + new SpecialTokensText('') + ]), + suffix: LlamaText(new SpecialTokensText("\n\n")), + emptyCallParamsPlaceholder: {} + }, + parallelism: { + call: { + sectionPrefix: "", + betweenCalls: LlamaText(new SpecialTokensText("<|eom|><|start|>assistant")), + sectionSuffix: LlamaText(new SpecialToken("EOT")) + }, + result: { + sectionPrefix: "" + } + }, + result: { + prefix: LlamaText([ + new SpecialTokensText("<|start|>tool "), "{{functionName}}", + new SpecialTokensText('<|message|>\n') + ]), + suffix: LlamaText([new SpecialTokensText("\n"), new SpecialToken("EOT")]) + } + }, + segments: { + thought: { + prefix: LlamaText(new SpecialTokensText(" to=self<|message|>")), + suffix: LlamaText(new SpecialTokensText("<|eom|>")) + } + } + }; + + public constructor(options: { + /** + * The amount of reasoning to instruct the model to use. + * + * Defaults to `"high"`. + */ + reasoningStrength?: "xhigh" | "high" | "medium" | "low", + + /** + * Whether to keep only the chain of thought from the last model response. + * + * Defaults to `false`, matching the original chat template. + */ + keepOnlyLastThought?: boolean, + + /** + * The knowledge cutoff used by the default system message. + * The default system message is applied only when you supply a chat history that don't have a system message at the beginning. + * + * Set to `null` to omit it. + * + * Defaults to `"2026-01-04"`. + */ + knowledgeCutoff?: Date | (() => Date) | string | null, + + /** + * The current date used by the default system message. + * The default system message is applied only when you supply a chat history that don't have a system message at the beginning. + * + * Set to `null` to omit it. + * + * Defaults to the current date. + */ + todayDate?: Date | (() => Date) | number | string | null, + + /** @internal */ + _systemMessageInSpecialTokensText?: boolean + } = {}) { + super(); + + const { + reasoningStrength = "high", + keepOnlyLastThought = false, + knowledgeCutoff = new Date("2026-01-04T00:00:00Z"), + todayDate = () => new Date(), + _systemMessageInSpecialTokensText = false + } = options; + + this.reasoningStrength = reasoningStrength; + this.keepOnlyLastThought = keepOnlyLastThought; + this.knowledgeCutoff = knowledgeCutoff == null + ? null + : knowledgeCutoff instanceof Function + ? knowledgeCutoff + : new Date(knowledgeCutoff); + this.todayDate = todayDate == null + ? null + : todayDate instanceof Function + ? todayDate + : new Date(todayDate); + + this._systemMessageInSpecialTokensText = _systemMessageInSpecialTokensText; + } + + public override generateContextState({ + chatHistory, availableFunctions, documentFunctionParams + }: ChatWrapperGenerateContextStateOptions): ChatWrapperGeneratedContextState { + const hasFunctions = Object.keys(availableFunctions ?? {}).length > 0; + const modifiedChatHistory = chatHistory.slice(); + + let systemMessage: LlamaText = LlamaText(); + if (modifiedChatHistory[0]?.type === "system") { + systemMessage = LlamaText([ + LlamaText.fromJSON(modifiedChatHistory[0].text), + this._systemMessageInSpecialTokensText + ? new SpecialTokensText("\n") + : "\n" + ]); + modifiedChatHistory.shift(); + } else + systemMessage = this._getDefaultSystemMessage(); + + const contextContent: LlamaText[] = [ + LlamaText(new SpecialToken("BOS")), + LlamaText(new SpecialTokensText("<|start|>system<|message|>")), + systemMessage, + this._getSystemMessage(availableFunctions, {documentParams: documentFunctionParams}) + ]; + + let needsTriggers = true; + for (let i = 0; i < modifiedChatHistory.length; i++) { + const isLastItem = i === modifiedChatHistory.length - 1; + const item = modifiedChatHistory[i]; + + if (item == null) + continue; + + if (item.type === "system") { + contextContent.push( + LlamaText([ + new SpecialTokensText("<|start|>system<|message|>"), + isLastItem + ? LlamaText([]) + : new SpecialTokensText("<|eot|>") + ]) + ); + + if (isLastItem) + needsTriggers = false; + } else if (item.type === "user") { + contextContent.push( + LlamaText([ + new SpecialTokensText("<|start|>user<|message|>"), + item.text, + isLastItem + ? LlamaText([]) + : new SpecialTokensText("<|eot|>") + ]) + ); + + if (isLastItem) + needsTriggers = false; + } else if (item.type === "model") { + const { + res, needsTriggers: modelNeedsTriggers + } = this._getModelResponse(item.response, true, isLastItem, this.keepOnlyLastThought); + + if (isLastItem) + needsTriggers = modelNeedsTriggers; + + contextContent.push(res); + } else + void (item satisfies never); + } + + const contextText = LlamaText(contextContent); + + if (!needsTriggers) + return { + contextText, + stopGenerationTriggers: [ + LlamaText(new SpecialToken("EOS")), + LlamaText(new SpecialTokensText("<|eot|>")), + LlamaText(new SpecialToken("EOT")) + ], + detectFunctionCalls: false, + rerender: { + triggers: [LlamaText(new SpecialTokensText("<|eom|>"))], + action: "closeResponseItem" + } + }; + + return { + contextText, + stopGenerationTriggers: [ + LlamaText(new SpecialToken("EOS")), + LlamaText(new SpecialTokensText("<|eot|>")), + LlamaText(new SpecialToken("EOT")) + ], + prefixTriggers: [{ + type: "segment", + segmentType: "thought", + triggers: [LlamaText(new SpecialTokensText(" to=self<|message|>"))] + }, + { + type: "response", + triggers: [LlamaText(new SpecialTokensText(" to=user<|message|>"))] + }], + noPrefixTrigger: hasFunctions + ? { + type: "functionCall", + inject: LlamaText(new SpecialTokensText("")) + } + : { + type: "response", + inject: LlamaText(new SpecialTokensText(" to=user<|message|>")) + }, + detectFunctionCalls: true, + rerender: { + triggers: [LlamaText(new SpecialTokensText("<|eom|>"))], + action: "closeResponseItem" + } + }; + } + + public override generateFunctionCall(name: string, params: any): LlamaText { + const emptyCallParamsPlaceholder = this.settings.functions.call.emptyCallParamsPlaceholder; + + return LlamaText([ + new SpecialTokensText(" to="), + name, + new SpecialTokensText("<|message|>\n"), + new SpecialTokensText('\n'), + new SpecialTokensText(''), + params === undefined + ? (emptyCallParamsPlaceholder === undefined || emptyCallParamsPlaceholder === "") + ? "" + : jsonDumps(emptyCallParamsPlaceholder) + : jsonDumps(params), + this.settings.functions.call.suffix + ]); + } + + public override generateFunctionCallResult(functionName: string, functionParams: any, result: any): LlamaText { + return LlamaText([ + new SpecialTokensText("<|start|>tool "), + functionName, + new SpecialTokensText('<|message|>\n'), + result === undefined + ? "" + : jsonDumps(result), + new SpecialTokensText("\n"), + new SpecialToken("EOT") + ]); + } + + public override generateModelResponseText(modelResponse: ChatModelResponse["response"], useRawValues: boolean = true): LlamaText { + const {res} = this._getModelResponse(modelResponse, useRawValues, false, false); + const [start, ...rest] = res.values; + let newStart = start; + let newEnd = rest.pop(); + + if (newStart instanceof SpecialTokensText && newStart.value.startsWith("<|start|>assistant")) + newStart = new SpecialTokensText(newStart.value.slice("<|start|>assistant".length)); + + if (newEnd instanceof SpecialTokensText && newEnd.value.endsWith("<|eot|>")) + newEnd = new SpecialTokensText(newEnd.value.slice(0, -"<|eot|>".length)); + + return LlamaText([ + newStart ?? [], + ...rest, + newEnd ?? [] + ]); + } + + public override generateAvailableFunctionsSystemText(availableFunctions: ChatModelFunctions, {documentParams = true}: { + documentParams?: boolean + }) { + if (Object.keys(availableFunctions).length === 0) + return LlamaText([]); + + const namespaceNames = new Set(); + for (const functionName of Object.keys(availableFunctions)) + namespaceNames.add(functionName.split(".")[0]!); + + const toolMetadata = [...namespaceNames] + .map((namespaceName) => jsonDumps({name: namespaceName, description: ""})) + .join("\n"); + const functionSchemas = Object.entries(availableFunctions) + .map(([name, definition]) => jsonDumps({ + name, + description: definition.description ?? "", + parameters: documentParams + ? (definition.params ?? {}) + : undefined + })) + .join("\n"); + + return LlamaText.joinValues("\n", [ + "In this environment you have access to a set of tools you can use to answer the user's question.", + "", + LlamaText([ + "You can invoke a function by writing a ", + new SpecialTokensText('""'), + " block like the following:" + ]), + new SpecialTokensText(""), + new SpecialTokensText(''), + new SpecialTokensText('$PARAMETER_VALUE'), + "...", + new SpecialTokensText(""), + new SpecialTokensText(""), + "", + ( + "String and scalar parameters should be specified as is, while lists and objects should use JSON format. " + + "Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with " + + "regular expressions." + ), + "Here are the functions available in JSONSchema format:", + "// Tool metadata", + toolMetadata, + "// Function schemas", + functionSchemas, + "Here's an example of how to call a function in the tool set:", + "(If the tool namespace is not specified, invoke the function directly as `example_function_name` rather than " + + "`example_tool_name.example_function_name`)", + "", + "to=example_tool_name.example_function_name", + "", + new SpecialTokensText(""), + new SpecialTokensText(''), + LlamaText([ + new SpecialTokensText(''), + "value_1", + new SpecialTokensText("") + ]), + LlamaText([ + new SpecialTokensText(''), + "This is the value for the second parameter" + ]), + "that can span", + '"multiple" lines', + new SpecialTokensText(""), + new SpecialTokensText(""), + new SpecialTokensText("") + ]); + } + + public override generateInitialChatHistory({ + systemPrompt + }: ChatWrapperGenerateInitialHistoryOptions = {}): ChatHistoryItem[] { + if (systemPrompt === "") + return [{ + type: "system", + text: this._getDefaultSystemMessage().toJSON() + }]; + + return super.generateInitialChatHistory({systemPrompt}); + } + + /** @internal */ + private _getModelResponse( + modelResponse: ChatModelResponse["response"], + useRawValues: boolean, + isLastItem: boolean, + keepOnlyLastThought: boolean + ) { + const res: LlamaText[] = []; + const pendingFunctionCalls: ChatModelFunctionCall[] = []; + let canEnableTriggers = true; + + const addPendingFunctions = () => { + if (pendingFunctionCalls.length === 0) + return; + + res.push(LlamaText(new SpecialTokensText("<|start|>assistant"))); + res.push(this.generateFunctionCallsAndResults(pendingFunctionCalls, useRawValues)); + + pendingFunctionCalls.length = 0; + }; + + for (let index = 0; index < modelResponse.length; index++) { + const isLastResponse = index === modelResponse.length - 1; + const response = modelResponse[index]; + + if (response == null) + continue; + else if (response === "" && (!isLastResponse || !isLastItem)) + continue; + + if (typeof response === "string") { + addPendingFunctions(); + res.push(LlamaText([ + new SpecialTokensText("<|start|>assistant to=user<|message|>"), + response, + (isLastResponse && isLastItem) + ? LlamaText([]) + : new SpecialTokensText("<|eot|>") + ])); + + if (isLastResponse && isLastItem) + canEnableTriggers = false; + } else if (response.type === "segment") { + addPendingFunctions(); + + if (response.ended && response.raw != null && useRawValues) + res.push(LlamaText([ + new SpecialTokensText("<|start|>assistant"), + LlamaText.fromJSON(response.raw) + ])); + else if (response.segmentType === "thought") { + if (keepOnlyLastThought && !isLastItem) + continue; + + res.push( + LlamaText([ + new SpecialTokensText("<|start|>assistant to=self<|message|>"), + response.text, + (isLastItem && !response.ended) + ? LlamaText([]) + : new SpecialTokensText("<|eom|>") + ]) + ); + + if (isLastItem && isLastResponse && !response.ended) + canEnableTriggers = false; + } else if (response.segmentType === "comment") + continue; // unsupported + else + void (response.segmentType satisfies never); + } else if (response.type === "functionCall") { + if (response.startsNewChunk) + addPendingFunctions(); + + pendingFunctionCalls.push(response); + } else + void (response satisfies never); + } + + addPendingFunctions(); + + const needsTriggers = canEnableTriggers && isLastItem; + if (needsTriggers) + res.push( + LlamaText([ + new SpecialTokensText("<|start|>assistant") + ]) + ); + + return { + res: LlamaText(res), + needsTriggers + }; + } + + /** @internal */ + private _getDefaultSystemMessage() { + const lines = ["You are a helpful AI assistant."]; + + if (this.knowledgeCutoff != null) { + const date = this.knowledgeCutoff instanceof Function + ? this.knowledgeCutoff() + : this.knowledgeCutoff; + + lines.push(`Knowledge cutoff: ${formatDate(date, "UTC")}.`); + } + + if (this.todayDate != null) { + const date = this.todayDate instanceof Function + ? this.todayDate() + : this.todayDate; + lines.push(`Current date: ${formatDate(date)}.`); + } + + return LlamaText(lines.join("\n")); + } + + /** @internal */ + private _getSystemMessage(availableFunctions?: ChatModelFunctions, {documentParams = true}: { + documentParams?: boolean + } = {}) { + const hasFunctions = Object.keys(availableFunctions ?? {}).length > 0; + const recipientNames = [ + '"self"', + ...(new Set(Object.keys(availableFunctions ?? {}).map((functionName) => ('"' + functionName.split(".")[0] + '.*"')))), + '"user"' + ]; + + let res = LlamaText([ + "\n", + `Reasoning strength: ${this.reasoningStrength}.`, + hasFunctions + ? LlamaText([ + "\n\n", + this.generateAvailableFunctionsSystemText(availableFunctions ?? {}, {documentParams}) + ]) + : [], + `\n\n# Valid recipients: ${recipientNames.join(", ")}.`, + new SpecialTokensText("<|eot|>") + ]); + + if (this._systemMessageInSpecialTokensText) + res = LlamaText(res.values.map((value) => { + if (typeof value === "string") + return new SpecialTokensText(value); + + return value; + })); + + return res; + } + + /** @internal */ + public static override _getOptionConfigurationsToTestIfCanSupersedeJinjaTemplate(): ChatWrapperJinjaMatchConfiguration { + return [...optionsMatrix({ + _systemMessageInSpecialTokensText: [false, true] + })].flatMap(({_systemMessageInSpecialTokensText}) => [ + [{_systemMessageInSpecialTokensText}, {}], + [ + { + todayDate: new Date("2026-08-11T00:00:00"), + cuttingKnowledgeDate: new Date("2026-01-04T00:00:00Z"), + _systemMessageInSpecialTokensText + }, + {}, + { + additionalRenderParameters: { + "current_date": "2026-08-11", + "knowledge_cutoff": "2026-01-04" + } + } + ] + ]); + } +} + +function formatDate(date: Date, timezone?: "UTC") { + const day = date.toLocaleDateString("en-US", {day: "numeric", timeZone: timezone}).padStart(2, "0"); + const month = date.toLocaleDateString("en-US", {month: "numeric", timeZone: timezone}).padStart(2, "0"); + const year = date.toLocaleDateString("en-US", {year: "numeric", timeZone: timezone}).padStart(4, "0"); + return `${day}-${month}-${year}`; +} diff --git a/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts b/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts index 39570a91..236831c4 100644 --- a/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts +++ b/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts @@ -13,6 +13,8 @@ import { import {removeUndefinedFields} from "../../utils/removeNullFields.js"; import {jsonDumps} from "../utils/jsonDumps.js"; import {tryMatrix} from "../../utils/optionsMatrix.js"; +import {getStandardizedChatWrapperSegmentDefinition} from "../../utils/getStandardizedChatWrapperSegmentDefinition.js"; +import {replaceRegularTextInLlamaText} from "../utils/replaceRegularTextInLlamaText.js"; import {ChatHistoryFunctionCallMessageTemplate, parseFunctionCallMessageTemplate} from "./utils/chatHistoryFunctionCallMessageTemplate.js"; import { templateSegmentOptionsToChatWrapperSettings, TemplateChatWrapperSegmentsOptions @@ -482,7 +484,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { : JSON.stringify(jsonDumps(emptyCallParamsPlaceholder)) : JSON.stringify(jsonDumps(params)) ), - this.settings.functions.call.suffix + replaceRegularTextInLlamaText(this.settings.functions.call.suffix, "{{functionName}}", name) ]); } @@ -747,6 +749,15 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { role: this.userRoleName, content: idsGenerator.generateId() } as OpenAiChatMessage); + } else if ( + lastJinjaItem?.role === this.modelRoleName && + typeof this.settings.segments?.thought?.prefix === "object" && + !LlamaText.isLlamaText(this.settings.segments?.thought?.prefix) && + this.settings.segments?.thought?.prefix.type === "openedOnStart" && + typeof lastJinjaItem.content === "string" + ) { + (lastJinjaItem as OpenAiChatAssistantMessage)["reasoning_content"] = lastJinjaItem.content; + lastJinjaItem.content = ""; } const renderJinjaText = () => { @@ -837,8 +848,9 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { const {splitJinjaParts, stopGenerationJinjaParts} = renderJinjaAndSplitIntoParts(); const messageIdsLeftToProcess = new Set(messageIds); - const thoughSegmentPrefix = getLlamaTextOnlyText(this.settings.segments?.thought?.prefix); - const thoughSegmentSuffix = getLlamaTextOnlyText(this.settings.segments?.thought?.suffix); + const standardizedSegmentDefinition = getStandardizedChatWrapperSegmentDefinition(this.settings, "thought"); + const thoughSegmentPrefix = getLlamaTextOnlyText(standardizedSegmentDefinition?.prefix); + const thoughSegmentSuffix = getLlamaTextOnlyText(standardizedSegmentDefinition?.suffix); let inLastModelResponseSection: boolean | null = ( thoughSegmentPrefix == null || thoughSegmentSuffix == null || diff --git a/src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts b/src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts index 3e564522..ca0b415e 100644 --- a/src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts +++ b/src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts @@ -177,22 +177,24 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ function extractControls() { const {responseOnly, withReasoning} = tryMatrix({ enableThinking: [true, null], - variation: ["simple", "separateReasoning", "nullContent", "reasoningFirst", "reasoningFirstNullMessage"] - }, ({enableThinking, variation}) => { + composition: ["simple", "separateReasoning", "nullContent", "reasoningFirst", "reasoningFirstNullMessage"] + }, ({enableThinking, composition}) => { const thinkingParam = enableThinking === true ? {"enable_thinking": true} : {}; - if (variation === "simple") + if (composition === "simple") return { + thinkingParam, responseOnly: renderTemplate(messagesWithModelResponseLongBase, thinkingParam), withReasoning: { long: renderTemplate(messagesWithModelReasoningLongBase, thinkingParam), short: renderTemplate(messagesWithModelReasoning, thinkingParam) } }; - else if (variation === "separateReasoning" || variation === "nullContent") + else if (composition === "separateReasoning" || composition === "nullContent") return { + thinkingParam, responseOnly: renderTemplate([...messagesWithModelResponseLongBase, { role: "assistant", content: "" @@ -200,27 +202,28 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ withReasoning: { long: renderTemplate([...messagesWithModelResponseLongBase, { role: "assistant", - ...(variation === "nullContent" ? {} : { + ...(composition === "nullContent" ? {} : { content: "" }), "reasoning_content": modelReasoning2 }], thinkingParam), short: renderTemplate([...messagesWithModelResponse, { role: "assistant", - ...(variation === "nullContent" ? {} : { + ...(composition === "nullContent" ? {} : { content: "" }), "reasoning_content": modelReasoning2 }], thinkingParam) } }; - else if (variation === "reasoningFirst" || variation === "reasoningFirstNullMessage") + else if (composition === "reasoningFirst" || composition === "reasoningFirstNullMessage") return { + thinkingParam, responseOnly: renderTemplate(messagesWithModelResponseLongBase, thinkingParam), withReasoning: { long: renderTemplate([...longBaseMessages, { role: "assistant", - ...(variation === "reasoningFirstNullMessage" ? {} : { + ...(composition === "reasoningFirstNullMessage" ? {} : { content: "" }), "reasoning_content": modelReasoning2 @@ -230,7 +233,7 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ }], thinkingParam), short: renderTemplate([...baseMessages, { role: "assistant", - ...(variation === "reasoningFirstNullMessage" ? {} : { + ...(composition === "reasoningFirstNullMessage" ? {} : { content: "" }), "reasoning_content": modelReasoning2 @@ -241,11 +244,11 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ } }; - void (variation satisfies never); - throw new Error(`Unsupported variation: ${variation}`); + void (composition satisfies never); + throw new Error(`Unsupported composition: ${composition}`); }); - let reasoningSectionStartPrefix: string | undefined = undefined; + let reasoningSectionStartPrefix: string | {type: "openedOnStart"} | undefined = undefined; let reasoningSectionEndPrefix: string | undefined = undefined; if (responseOnly === withReasoning.long) @@ -258,17 +261,27 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ const modelResponsePrefix = responseOnly.slice(0, modelResponseIndex); const withReasoningPrefixContent = withReasoning.short.slice(0, modelResponseIndex); - if (modelResponsePrefix !== withReasoningPrefixContent) - return undefined; - - const reasoningSectionStartIndex = modelResponseIndex; - const reasoningContentStartIndex = withReasoning.short.indexOf(modelReasoning2, reasoningSectionStartIndex); - if (reasoningContentStartIndex < 0) - return undefined; - - reasoningSectionStartPrefix = withReasoning.short.slice(modelResponseIndex, reasoningContentStartIndex); + let reasoningContentEndIndex: number; + if (modelResponsePrefix !== withReasoningPrefixContent) { + const reasoningContentStartIndex = withReasoning.short.indexOf(modelReasoning2); + if (reasoningContentStartIndex < 0) + return undefined; + + if (responseOnly.slice(0, reasoningContentStartIndex) !== withReasoning.short.slice(0, reasoningContentStartIndex)) + return undefined; + + reasoningSectionStartPrefix = {type: "openedOnStart"}; + reasoningContentEndIndex = reasoningContentStartIndex + modelReasoning2.length; + } else { + const reasoningSectionStartIndex = modelResponseIndex; + const reasoningContentStartIndex = withReasoning.short.indexOf(modelReasoning2, reasoningSectionStartIndex); + if (reasoningContentStartIndex < 0) + return undefined; + + reasoningSectionStartPrefix = withReasoning.short.slice(modelResponseIndex, reasoningContentStartIndex); + reasoningContentEndIndex = reasoningContentStartIndex + modelReasoning2.length; + } - const reasoningContentEndIndex = reasoningContentStartIndex + modelReasoning2.length; const modelResponseStartIndex = withReasoning.short.indexOf(modelResponse2, reasoningContentEndIndex); if (modelResponseStartIndex < 0) return undefined; @@ -284,13 +297,13 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ function shouldKeepPastThinking() { const renderedOutput = tryMatrix({ enableThinking: [true, null], - variation: ["simple", "reasoningFirst", "reasoningFirstNullMessage"] - }, ({enableThinking, variation}) => { + composition: ["simple", "reasoningFirst", "reasoningFirstNullMessage"] + }, ({enableThinking, composition}) => { const thinkingParam = enableThinking === true ? {"enable_thinking": true} : {}; - if (variation === "simple") + if (composition === "simple") return renderTemplate([...messagesWithModelReasoning, { role: "user", content: userMessage2 @@ -299,10 +312,10 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ content: modelResponse3, "reasoning_content": modelReasoning3 }], thinkingParam); - else if (variation === "reasoningFirst" || variation === "reasoningFirstNullMessage") + else if (composition === "reasoningFirst" || composition === "reasoningFirstNullMessage") return renderTemplate([...baseMessages, { role: "assistant", - ...(variation === "reasoningFirstNullMessage" ? {} : { + ...(composition === "reasoningFirstNullMessage" ? {} : { content: "" }), "reasoning_content": modelReasoning2 @@ -314,7 +327,7 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ content: userMessage2 }, { role: "assistant", - ...(variation === "reasoningFirstNullMessage" ? {} : { + ...(composition === "reasoningFirstNullMessage" ? {} : { content: "" }), "reasoning_content": modelReasoning3 @@ -323,14 +336,14 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ content: modelResponse3 }], thinkingParam); - void (variation satisfies never); - throw new Error(`Unsupported variation: ${variation}`); + void (composition satisfies never); + throw new Error(`Unsupported composition: ${composition}`); }); return renderedOutput.includes(modelReasoning2) && renderedOutput.includes(modelReasoning3); } - function shouldOpenThinkingSegmentOnModelResponseStart(reasoningSectionPrefix: string, reasoningSectionSuffix: string) { + function shouldOpenThinkingSegmentOnModelResponseStart(reasoningSectionPrefix: string | undefined, reasoningSectionSuffix: string) { if (!enableReasoning) return false; @@ -363,19 +376,25 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ if (withReasoningUserMessage1Index < 0) return false; - if (responseOnly.indexOf(reasoningSectionPrefix, userMessage1Index) >= 0) - return false; + let reasoningSectionPrefixIndex: number; + let reasoningSectionEndIndex: number; + if (reasoningSectionPrefix != null) { + if (responseOnly.indexOf(reasoningSectionPrefix, userMessage1Index) >= 0) + return false; - const reasoningSectionPrefixIndex = withGenerationPrompt.indexOf(reasoningSectionPrefix, withReasoningUserMessage1Index); - if (reasoningSectionPrefixIndex < 0) - return false; + reasoningSectionPrefixIndex = withGenerationPrompt.indexOf(reasoningSectionPrefix, withReasoningUserMessage1Index); + if (reasoningSectionPrefixIndex < 0) + return false; + + reasoningSectionEndIndex = reasoningSectionPrefixIndex + reasoningSectionPrefix.length; + } else { + reasoningSectionPrefixIndex = userMessage1Index; + reasoningSectionEndIndex = reasoningSectionPrefixIndex; + } const reasoningSectionSuffixIndex = withGenerationPrompt.indexOf(reasoningSectionSuffix, reasoningSectionPrefixIndex); if (reasoningSectionSuffixIndex >= 0) { - const reasoningSectionContent = withGenerationPrompt.slice( - reasoningSectionPrefixIndex + reasoningSectionPrefix.length, - reasoningSectionSuffixIndex - ); + const reasoningSectionContent = withGenerationPrompt.slice(reasoningSectionEndIndex, reasoningSectionSuffixIndex); if (reasoningSectionContent.trim() === "") return false; @@ -390,7 +409,7 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ try { controls = extractControls(); - if (controls == null || controls.prefix.trim() === "") + if (controls == null || (typeof controls.prefix === "string" && controls.prefix.trim() === "")) return { thoughtSegment: undefined, keepPastReasoning: undefined @@ -409,25 +428,41 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate({ } try { - openOnResponseStart = controls != null && shouldOpenThinkingSegmentOnModelResponseStart(controls.prefix, controls.suffix); + if (controls != null) { + if (typeof controls.prefix !== "string") { + if (controls.prefix.type === "openedOnStart") + openOnResponseStart = true; + else + void (controls.prefix.type satisfies never); + } else + openOnResponseStart = shouldOpenThinkingSegmentOnModelResponseStart(controls.prefix, controls.suffix); + } else + openOnResponseStart = false; } catch (err) { // do nothing } const thoughtSuffix = controls.suffix.trim() === "" - ? ( - knownThinkingSegmentControls.get(controls.prefix) ?? - knownThinkingSegmentControls.get(controls.prefix.trim()) - ) + ? typeof controls.prefix === "string" + ? ( + knownThinkingSegmentControls.get(controls.prefix) ?? + knownThinkingSegmentControls.get(controls.prefix.trim()) + ) + : undefined : controls.suffix; return { thoughtSegment: { - prefix: LlamaText(new SpecialTokensText(controls.prefix)), + prefix: typeof controls.prefix !== "string" + ? controls.prefix.type === "openedOnStart" + ? {type: "openedOnStart"} + : void (controls.prefix.type satisfies never) as never + : LlamaText(new SpecialTokensText(controls.prefix)), suffix: thoughtSuffix != null ? LlamaText(new SpecialTokensText(thoughtSuffix)) : undefined, - openOnResponseStart + openOnResponseStart, + reopenAfterFunctionCalls: typeof controls.prefix !== "string" && controls.prefix.type === "openedOnStart" }, keepPastReasoning }; diff --git a/src/chatWrappers/utils/isJinjaTemplateEquivalentToSpecializedChatWrapper.ts b/src/chatWrappers/utils/isJinjaTemplateEquivalentToSpecializedChatWrapper.ts index fd82af7a..83ff0ce2 100644 --- a/src/chatWrappers/utils/isJinjaTemplateEquivalentToSpecializedChatWrapper.ts +++ b/src/chatWrappers/utils/isJinjaTemplateEquivalentToSpecializedChatWrapper.ts @@ -136,8 +136,16 @@ function checkEquivalence( (specializedChatWrapper as Writable).settings = originalSpecializedSettings; } - if (jinjaChatWrapper.settings.segments?.thought?.openOnResponseStart === true && - specializedChatWrapper.settings.segments?.thought?.openOnResponseStart !== true + const jinjaThoughtSegment = jinjaChatWrapper.settings.segments?.thought; + const specializedThoughtSegment = specializedChatWrapper.settings.segments?.thought; + + const jinjaOpenOnStart = hasThoughtSegmentOpenedOnStart(jinjaChatWrapper.settings); + const specializedOpenOnStart = hasThoughtSegmentOpenedOnStart(specializedChatWrapper.settings); + if (jinjaOpenOnStart !== specializedOpenOnStart) + return false; + else if ( + !jinjaOpenOnStart && + (jinjaThoughtSegment?.openOnResponseStart ?? false) !== (specializedThoughtSegment?.openOnResponseStart ?? false) ) return false; @@ -318,8 +326,8 @@ function convertChatWrapperSettingsToUseSpecialTokensText(settings: ChatWrapperS call: { ...settings.functions.call, prefix: convertToSpecialTokensText(settings.functions.call.prefix), - suffix: convertToSpecialTokensText(settings.functions.call.suffix), - paramsPrefix: convertToSpecialTokensText(settings.functions.call.paramsPrefix) + suffix: convertToSpecialTokensText(settings.functions.call.suffix, ["{{functionName}}"]), + paramsPrefix: convertToSpecialTokensText(settings.functions.call.paramsPrefix, ["{{functionName}}"]) }, result: { ...settings.functions.result, @@ -512,3 +520,10 @@ function removeLeadingBos(llamaText: LlamaText) { return llamaText; } + +function hasThoughtSegmentOpenedOnStart(settings: ChatWrapperSettings): boolean { + const thoughtSegment = settings.segments?.thought; + + return typeof thoughtSegment?.prefix === "object" && !LlamaText.isLlamaText(thoughtSegment.prefix) && + thoughtSegment.prefix.type === "openedOnStart"; +} diff --git a/src/chatWrappers/utils/replaceRegularTextInLlamaText.ts b/src/chatWrappers/utils/replaceRegularTextInLlamaText.ts new file mode 100644 index 00000000..7a28c45c --- /dev/null +++ b/src/chatWrappers/utils/replaceRegularTextInLlamaText.ts @@ -0,0 +1,11 @@ +import {LlamaText} from "../../utils/LlamaText.js"; + +export function replaceRegularTextInLlamaText(text: string | LlamaText, find: string, replace: string) { + return LlamaText(text) + .mapValues((value) => { + if (typeof value !== "string") + return value; + + return value.replaceAll(find, replace); + }); +} diff --git a/src/chatWrappers/utils/resolveChatWrapper.ts b/src/chatWrappers/utils/resolveChatWrapper.ts index bbbca3b3..27ba453f 100644 --- a/src/chatWrappers/utils/resolveChatWrapper.ts +++ b/src/chatWrappers/utils/resolveChatWrapper.ts @@ -21,15 +21,16 @@ import {LlamaModel} from "../../evaluator/LlamaModel/LlamaModel.js"; import {QwenChatWrapper} from "../QwenChatWrapper.js"; import {HarmonyChatWrapper} from "../HarmonyChatWrapper.js"; import {SeedChatWrapper} from "../SeedChatWrapper.js"; +import {MuseChatWrapper} from "../MuseChatWrapper.js"; +import {GgufArchitectureType} from "../../gguf/types/GgufMetadataTypes.js"; import {isJinjaTemplateEquivalentToSpecializedChatWrapper} from "./isJinjaTemplateEquivalentToSpecializedChatWrapper.js"; import {getModelLinageNames} from "./getModelLinageNames.js"; import type {GgufFileInfo} from "../../gguf/types/GgufFileInfoTypes.js"; -import type {GgufArchitectureType} from "../../gguf/types/GgufMetadataTypes.js"; export const specializedChatWrapperTypeNames = Object.freeze([ "general", "deepSeek", "qwen", "llama3.2-lightweight", "llama3.1", "llama3", "llama2Chat", "mistral", "alpacaChat", "functionary", - "chatML", "falconChat", "gemma4", "gemma", "harmony", "seed" + "chatML", "falconChat", "gemma4", "gemma", "harmony", "muse", "seed" ] as const); export type SpecializedChatWrapperTypeName = (typeof specializedChatWrapperTypeNames)[number]; @@ -61,6 +62,7 @@ export const chatWrappers = Object.freeze({ "gemma4": Gemma4ChatWrapper, "gemma": GemmaChatWrapper, "harmony": HarmonyChatWrapper, + "muse": MuseChatWrapper, "seed": SeedChatWrapper, "template": TemplateChatWrapper, "jinjaTemplate": JinjaTemplateChatWrapper @@ -74,6 +76,7 @@ const chatWrapperToConfigType = new Map( const specializedChatWrapperRelatedTexts = { "harmony": ["gpt", "gpt-oss"], + "muse": ["muse", "muse glimmer", "muse-glimmer", "muse_glimmer"], "gemma4": ["gemma 4", "gemma-4"] } satisfies Partial>; @@ -380,6 +383,8 @@ export function resolveChatWrapper( return createSpecializedChatWrapper(GemmaChatWrapper); else if (includesText(modelNames, ["gpt-oss", "Gpt Oss", "Gpt-Oss", "openai_gpt-oss", "Openai_Gpt Oss", "openai.gpt-oss", "Openai.Gpt Oss"])) return createSpecializedChatWrapper(HarmonyChatWrapper); + else if (includesText(modelNames, ["Muse Glimmer", "Muse-Glimmer", "Muse_Glimmer", "muse-glimmer", "muse_glimmer"])) + return createSpecializedChatWrapper(MuseChatWrapper); else if (includesText(modelNames, ["seed-oss", "Seed Oss", "Seed OSS", "Seed-Oss", "Seed-OSS", "ByteDance-Seed_Seed-OSS", "ByteDance-Seed.Seed-OSS"])) return createSpecializedChatWrapper(SeedChatWrapper); } @@ -463,14 +468,16 @@ export function resolveChatWrapper( } if (architecture != null) { - if (architecture === "llama") + if (architecture === GgufArchitectureType.llama) return createSpecializedChatWrapper(GeneralChatWrapper); - else if (architecture === "falcon") + else if (architecture === GgufArchitectureType.falcon) return createSpecializedChatWrapper(FalconChatWrapper); - else if (architecture === "gemma" || architecture === "gemma2") + else if (architecture === GgufArchitectureType.gemma || architecture === GgufArchitectureType.gemma2) return createSpecializedChatWrapper(GemmaChatWrapper); - else if (architecture === "gemma4") + else if (architecture === GgufArchitectureType.gemma4) return createSpecializedChatWrapper(Gemma4ChatWrapper); + else if (architecture === GgufArchitectureType.museGlimmer) + return createSpecializedChatWrapper(MuseChatWrapper); } return null; diff --git a/src/cli/commands/inspect/commands/InspectGgufCommand.ts b/src/cli/commands/inspect/commands/InspectGgufCommand.ts index b73f42a8..3e004d3a 100644 --- a/src/cli/commands/inspect/commands/InspectGgufCommand.ts +++ b/src/cli/commands/inspect/commands/InspectGgufCommand.ts @@ -18,7 +18,8 @@ import {GgmlType, GgufTensorInfo} from "../../../../gguf/types/GgufTensorInfoTyp import {toBytes} from "../../../utils/toBytes.js"; import {printDidYouMeanUri} from "../../../utils/resolveCommandGgufPath.js"; import {isModelUri} from "../../../../utils/parseModelUri.js"; -import {getDominantTensorType} from "../../../../gguf/insights/GgufInsights.js"; +import {getDominantTensorType, getTotalModelParameters} from "../../../../gguf/insights/GgufInsights.js"; +import {formatModelParameterCount} from "../../../../utils/formatModelParameterCount.js"; const chatTemplateKey = ".chatTemplate"; @@ -226,6 +227,7 @@ export const InspectGgufCommand: CommandModule = { console.info(`${chalk.yellow("Spliced parts:")} ${parsedMetadata.splicedParts}`); const dominantTensorType = getDominantTensorType(parsedMetadata.fullTensorInfo ?? []); + const modelParameters = getTotalModelParameters(parsedMetadata.fullTensorInfo ?? []); console.info(`${chalk.yellow("GGUF version:")} ${parsedMetadata.version}`); console.info(`${chalk.yellow("Tensor count:")} ${parsedMetadata.totalTensorCount.toLocaleString("en-US", numberLocaleFormattingOptions)}`); @@ -236,6 +238,9 @@ export const InspectGgufCommand: CommandModule = { if (dominantTensorType != null) console.info(`${chalk.yellow("Dominant tensor type:")} ${dominantTensorType} (${GgmlType[dominantTensorType]})`); + if (!noSplice && modelParameters !== 0) + console.info(`${chalk.yellow("Total " + (parsedMetadata.splicedParts === 1 ? "file" : "files") + " parameters:")} ${formatModelParameterCount(modelParameters)}`); + console.info(`${chalk.yellow("Metadata:")} ${prettyPrintObject(parsedMetadata.metadata, undefined, metadataPrettyPrintOptions)}`); console.info(`${chalk.yellow("Tensor info:")} ${prettyPrintObject(parsedMetadata.fullTensorInfo, undefined, tensorInfoPrettyPrintOptions)}`); } diff --git a/src/cli/utils/resolveCommandGgufPath.ts b/src/cli/utils/resolveCommandGgufPath.ts index 4a6e25a9..55d48c67 100644 --- a/src/cli/utils/resolveCommandGgufPath.ts +++ b/src/cli/utils/resolveCommandGgufPath.ts @@ -5,7 +5,7 @@ import {cliModelsDirectory} from "../../config.js"; import {Llama} from "../../bindings/Llama.js"; import {createModelDownloader} from "../../utils/createModelDownloader.js"; import {resolveModelDestination} from "../../utils/resolveModelDestination.js"; -import {ggufQuantNames} from "../../gguf/utils/ggufQuantNames.js"; +import {ggufFileQuantNamesSet} from "../../gguf/utils/ggufQuantNames.js"; import {getConsoleLogPrefix} from "../../utils/getConsoleLogPrefix.js"; import {isModelUri} from "../../utils/parseModelUri.js"; import {GgmlType, resolveGgmlTypeOption} from "../../gguf/types/GgufTensorInfoTypes.js"; @@ -179,7 +179,7 @@ export function tryCoercingModelUri(ggufPath: string) { // / // /: foundSlashes === 1 && - (possibleQuant == null || possibleQuant === "latest" || ggufQuantNames.has(possibleQuant.toUpperCase())) + (possibleQuant == null || possibleQuant === "latest" || ggufFileQuantNamesSet.has(possibleQuant.toUpperCase())) ) ) { const possibleUri = "hf:" + ggufPath; diff --git a/src/evaluator/LlamaChat/LlamaChat.ts b/src/evaluator/LlamaChat/LlamaChat.ts index 8fdd0793..737c6c08 100644 --- a/src/evaluator/LlamaChat/LlamaChat.ts +++ b/src/evaluator/LlamaChat/LlamaChat.ts @@ -23,10 +23,11 @@ import {pushAll} from "../../utils/pushAll.js"; import {resolveLastTokens} from "../../utils/resolveLastTokens.js"; import {LlamaSampler} from "../LlamaContext/LlamaSampler.js"; import {LlamaModel} from "../LlamaModel/LlamaModel.js"; -import {getChatWrapperSegmentDefinition} from "../../utils/getChatWrapperSegmentDefinition.js"; +import {getStandardizedChatWrapperSegmentDefinition} from "../../utils/getStandardizedChatWrapperSegmentDefinition.js"; import {jsonDumps} from "../../chatWrappers/utils/jsonDumps.js"; import {defaultMaxPreloadTokens} from "../LlamaChatSession/utils/LlamaChatSessionPromptCompletionEngine.js"; import {LlamaLogLevel} from "../../bindings/types.js"; +import {replaceRegularTextInLlamaText} from "../../chatWrappers/utils/replaceRegularTextInLlamaText.js"; import { eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy } from "./utils/contextShiftStrategies/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.js"; @@ -712,8 +713,11 @@ export class LlamaChat { await loadContextWindow(); generateResponseState.isRerender = false; - if (generateResponseState.isFirstEvaluation && generateResponseState.onModelResponseStateShouldOpenThoughtSegment()) { - if (!tookInitialCheckpoint && this.sequence.needsCheckpoints) { + const shouldOpenThoughtSegment = generateResponseState.isFirstEvaluation + ? generateResponseState.onModelResponseStartShouldOpenThoughtSegment() + : false; + if (shouldOpenThoughtSegment !== false) { + if (!tookInitialCheckpoint && this.sequence.needsCheckpoints && shouldOpenThoughtSegment !== "openedOnStart") { await generateResponseState.alignCurrentSequenceStateWithCurrentTokens(false); await generateResponseState.evaluateWithoutGeneratingNewTokens(); @@ -721,7 +725,7 @@ export class LlamaChat { tookInitialCheckpoint = true; } - generateResponseState.openThoughtSegmentOnModelResponseStartIfNeeded(); + generateResponseState.openThoughtSegmentOnModelResponseStart(); generateResponseState.canAvoidReloadingHistory = false; generateResponseState.isRerender = shouldHandlePrefixTriggers; await loadContextWindow(); @@ -973,9 +977,12 @@ export class LlamaChat { )); allSegmentTypes - .map((segmentType) => getChatWrapperSegmentDefinition(this._chatWrapper.settings, segmentType)) - .filter((segmentDefinition) => segmentDefinition != null) - .flatMap((segmentDefinition) => [segmentDefinition?.prefix, segmentDefinition?.suffix]) + .map((segmentType) => getStandardizedChatWrapperSegmentDefinition(this._chatWrapper.settings, segmentType)) + .filter((standardizedSegmentDefinition) => standardizedSegmentDefinition != null) + .flatMap((standardizedSegmentDefinition) => [ + standardizedSegmentDefinition?.prefix, + standardizedSegmentDefinition?.suffix + ]) .filter((trigger) => trigger != null) .forEach((trigger) => ( generateResponseState.stopGenerationDetector.addStopTrigger( @@ -1946,23 +1953,28 @@ class GenerateResponseState[0]["segmentDefinitions"] = new Map(); for (const segmentType of allSegmentTypes) { - const segmentDefinition = getChatWrapperSegmentDefinition(this.chatWrapper.settings, segmentType); - if (segmentDefinition != null) - segmentDefinitions.set(segmentType, segmentDefinition); + const standardizedSegmentDefinition = getStandardizedChatWrapperSegmentDefinition(this.chatWrapper.settings, segmentType); + if (standardizedSegmentDefinition != null) + segmentDefinitions.set(segmentType, standardizedSegmentDefinition); } const lastModelMessageFullResponse = getLastModelMessageFullResponseFromChatHistory(this.resolvedHistory); @@ -1984,23 +1996,28 @@ class GenerateResponseState 0) return false; @@ -2055,17 +2078,11 @@ class GenerateResponseState 0) - return false; - const currentResponseSegmentsStack = SegmentHandler.getStackFromModelResponse(lastModelResponseItem.response); if (currentResponseSegmentsStack.includes("thought")) return false; @@ -2918,7 +2935,11 @@ class GenerateResponseState; private readonly _segmentsStack: S[] = []; @@ -3809,7 +3838,7 @@ class SegmentHandler; @@ -3823,7 +3852,7 @@ class SegmentHandler, closeAllSegments?: string | LlamaText, @@ -3853,8 +3882,10 @@ class SegmentHandler>; + private readonly _loadModelError?: Error; private _disposed = false; - public constructor(llama: Llama, lruCacheSize: number = 10) { + public constructor(llama: Llama, lruCacheSize: number = 10, loadModelError?: Error) { this._llama = llama; this._modelHandlePromises = new LruCache(lruCacheSize, { async onDelete(key, value) { @@ -1218,6 +1214,7 @@ export class GgufInsightsSimulatorSession { } } }); + this._loadModelError = loadModelError; } public async estimateModelResources({ @@ -1373,6 +1370,8 @@ export class GgufInsightsSimulatorSession { }) { if (this._disposed) throw new Error("simulator session is disposed"); + else if (this._loadModelError != null) + throw this._loadModelError; let preventDisposalHandle: DisposalPreventionHandle; try { @@ -2026,3 +2025,17 @@ export function getDominantTensorType(tensorInfo: GgufTensorInfo[]): GgmlType | return dominantType; } + +export function getTotalModelParameters(tensorInfo: GgufTensorInfo[]) { + let totalParameters = 0n; + + for (const tensor of tensorInfo) { + let tensorParameters = 1n; + for (const dim of tensor.dimensions ?? []) + tensorParameters *= BigInt(dim); + + totalParameters += tensorParameters; + } + + return Number(totalParameters); +} diff --git a/src/gguf/insights/GgufInsightsConfigurationResolver.ts b/src/gguf/insights/GgufInsightsConfigurationResolver.ts index d174e5b1..9688bb85 100644 --- a/src/gguf/insights/GgufInsightsConfigurationResolver.ts +++ b/src/gguf/insights/GgufInsightsConfigurationResolver.ts @@ -207,7 +207,7 @@ export class GgufInsightsConfigurationResolver { ? this._ggufInsights._getUseMmap() : useMmap; let gpuLayersFitMemory = false; - const simulatorSession = this._ggufInsights._createSimulatorSession(); + await using simulatorSession = this._ggufInsights._createSimulatorSession(); try { const layersResolution = await this.resolveModelGpuLayersV2( @@ -547,7 +547,7 @@ export class GgufInsightsConfigurationResolver { _simulatorSession } = options; - + return await resolveContextContextSizeOption({ contextSize, batchSize, diff --git a/src/gguf/types/GgufMetadataTypes.ts b/src/gguf/types/GgufMetadataTypes.ts index 1faae49c..18040122 100644 --- a/src/gguf/types/GgufMetadataTypes.ts +++ b/src/gguf/types/GgufMetadataTypes.ts @@ -48,6 +48,7 @@ export const enum GgufArchitectureType { gemma3 = "gemma3", gemma3n = "gemma3n", gemma4 = "gemma4", + gemma4assistant = "gemma4-assistant", gemmaEmbedding = "gemma-embedding", starcoder2 = "starcoder2", mamba = "mamba", @@ -57,15 +58,19 @@ export const enum GgufArchitectureType { xverse = "xverse", commandR = "command-r", cohere2 = "cohere2", + cohere2moe = "cohere2moe", dbrx = "dbrx", olmo = "olmo", olmo2 = "olmo2", olmoe = "olmoe", + museGlimmer = "muse-glimmer", openelm = "openelm", arctic = "arctic", deepseek = "deepseek", deepseek2 = "deepseek2", deepseek2ocr = "deepseek2-ocr", + deepseek32 = "deepseek32", + deepseek4 = "deepseek4", chatglm = "chatglm", glm4 = "glm4", glm4moe = "glm4moe", @@ -88,6 +93,7 @@ export const enum GgufArchitectureType { granite = "granite", granitemoe = "granitemoe", granitehybrid = "granitehybrid", + graniteswitch = "graniteswitch", chameleon = "chameleon", wavtokenizerDec = "wavtokenizer-dec", plm = "plm", @@ -96,10 +102,13 @@ export const enum GgufArchitectureType { dots1 = "dots1", arcee = "arcee", afmoe = "afmoe", + laguna = "laguna", ernie4_5 = "ernie4_5", ernie4_5Moe = "ernie4_5-moe", hunyuanMoe = "hunyuan-moe", hunyuanDense = "hunyuan-dense", + hunyuanVl = "hunyuan_vl", + hyV3 = "hy_v3", smollm3 = "smollm3", gptOss = "gpt-oss", lfm2 = "lfm2", @@ -112,10 +121,13 @@ export const enum GgufArchitectureType { grovemoe = "grovemoe", apertus = "apertus", minimaxM2 = "minimax-m2", + minimaxM3 = "minimax-m3", cogvlm = "cogvlm", rnd1 = "rnd1", panguEmbedded = "pangu-embedded", mistral3 = "mistral3", + eagle3 = "eagle3", + dflash = "dflash", mistral4 = "mistral4", paddleocr = "paddleocr", mimo2 = "mimo2", @@ -123,6 +135,10 @@ export const enum GgufArchitectureType { llamaEmbed = "llama-embed", maincoder = "maincoder", kimiLinear = "kimi-linear", + talkie = "talkie", + mellum = "mellum", + nanbeige = "nanbeige", + qwen3tts = "qwen3tts", clip = "clip", unknown = "(unknown)" } diff --git a/src/gguf/utils/ggufQuantNames.ts b/src/gguf/utils/ggufQuantNames.ts index 0b9492cf..e3c9f453 100644 --- a/src/gguf/utils/ggufQuantNames.ts +++ b/src/gguf/utils/ggufQuantNames.ts @@ -42,8 +42,7 @@ export const ggufQuantNames = new Map([ ["F32", GgufFileType.ALL_F32], ["COPY", GgufFileType.ALL_F32] ]); - -export const ggufFileQuantNames = Object.freeze([ +export const ggufFileQuantNamesSet = new Set([ ...ggufQuantNames.keys(), "Q2_K_XL", "Q3_K_XL", @@ -52,4 +51,7 @@ export const ggufFileQuantNames = Object.freeze([ "Q6_K_XL", "Q7_K_XL", "Q8_K_XL" -].filter((name) => name !== "COPY")); +]); +ggufFileQuantNamesSet.delete("COPY"); + +export const ggufFileQuantNames = Object.freeze([...ggufFileQuantNamesSet]); diff --git a/src/index.ts b/src/index.ts index 7dad0945..541d15a3 100644 --- a/src/index.ts +++ b/src/index.ts @@ -65,6 +65,7 @@ import {FunctionaryChatWrapper} from "./chatWrappers/FunctionaryChatWrapper.js"; import {GemmaChatWrapper} from "./chatWrappers/GemmaChatWrapper.js"; import {Gemma4ChatWrapper} from "./chatWrappers/Gemma4ChatWrapper.js"; import {HarmonyChatWrapper} from "./chatWrappers/HarmonyChatWrapper.js"; +import {MuseChatWrapper} from "./chatWrappers/MuseChatWrapper.js"; import {SeedChatWrapper} from "./chatWrappers/SeedChatWrapper.js"; import {TemplateChatWrapper, type TemplateChatWrapperOptions} from "./chatWrappers/generic/TemplateChatWrapper.js"; import { @@ -240,6 +241,7 @@ export { GemmaChatWrapper, Gemma4ChatWrapper, HarmonyChatWrapper, + MuseChatWrapper, SeedChatWrapper, TemplateChatWrapper, type TemplateChatWrapperOptions, diff --git a/src/types.ts b/src/types.ts index 7034a5b6..6d0f3e3d 100644 --- a/src/types.ts +++ b/src/types.ts @@ -26,7 +26,39 @@ export type ChatWrapperSettings = { readonly call: { readonly optionalPrefixSpace: boolean, readonly prefix: string | LlamaText, + + /** + * Supported template parameters: + * - `{{functionName}}` + * + * Template parameters can only appear in a string or a string in a `LlamaText`. + * + * Template parameters inside a `SpecialTokensText` inside a `LlamaText` won't be replaced. + * + * Example of supported values: + * - `"text{{functionName}}text"` + * - `LlamaText(["text{{functionName}}text"])` + * + * Example of unsupported values: + * - `LlamaText([new SpecialTokensText("text{{functionName}}text")])` + */ readonly paramsPrefix: string | LlamaText, + + /** + * Supported template parameters: + * - `{{functionName}}` + * + * Template parameters can only appear in a string or a string in a `LlamaText`. + * + * Template parameters inside a `SpecialTokensText` inside a `LlamaText` won't be replaced. + * + * Example of supported values: + * - `"text{{functionName}}text"` + * - `LlamaText(["text{{functionName}}text"])` + * + * Example of unsupported values: + * - `LlamaText([new SpecialTokensText("text{{functionName}}text")])` + */ readonly suffix: string | LlamaText, /** @@ -36,7 +68,13 @@ export type ChatWrapperSettings = { * * Defaults to `""`. */ - readonly emptyCallParamsPlaceholder?: object | string | number | boolean | null + readonly emptyCallParamsPlaceholder?: object | string | number | boolean | null, + + /** + * Alternate prefixes that can be used to detect a function call, + * but won't be used to construct the context when building it from scratch. + */ + readonly prefixAlternateMatches?: Array }, readonly result: { @@ -111,7 +149,9 @@ export type ChatWrapperSettings = { readonly reiterateStackAfterFunctionCalls?: boolean, /** Chain of Thought text segment */ - readonly thought?: ChatWrapperSettingsSegment & { + readonly thought?: { + readonly prefix: string | LlamaText | {type: "openedOnStart"}, + readonly suffix?: string | LlamaText, openOnResponseStart?: boolean, reopenAfterFunctionCalls?: boolean }, diff --git a/src/utils/OpenAIFormat.ts b/src/utils/OpenAIFormat.ts index c3d7426c..eeed28c6 100644 --- a/src/utils/OpenAIFormat.ts +++ b/src/utils/OpenAIFormat.ts @@ -8,7 +8,7 @@ import {LlamaGrammar} from "../evaluator/LlamaGrammar.js"; import {Llama} from "../bindings/Llama.js"; import {LlamaModel} from "../evaluator/LlamaModel/LlamaModel.js"; import {GbnfJsonSchema} from "./gbnfJson/types.js"; -import {getChatWrapperSegmentDefinition} from "./getChatWrapperSegmentDefinition.js"; +import {getStandardizedChatWrapperSegmentDefinition} from "./getStandardizedChatWrapperSegmentDefinition.js"; import {LlamaText} from "./LlamaText.js"; import {removeUndefinedFields} from "./removeNullFields.js"; @@ -221,7 +221,10 @@ export function fromChatHistoryToIntermediateOpenAiMessages 0 && segmentStack.at(-1) === response.segmentType) ? "" - : segmentDefinition?.prefix ?? "", + : (standardizedSegmentDefinition?.prefix ?? ""), response.text, response.ended - ? (segmentDefinition?.suffix ?? "") + ? (standardizedSegmentDefinition?.suffix ?? "") : "" ]) ); @@ -243,7 +246,7 @@ export function fromChatHistoryToIntermediateOpenAiMessages(); for (const segmentType of allSegmentTypes) { - const segmentDefinition = getChatWrapperSegmentDefinition(chatWrapper.settings, segmentType); - if (segmentDefinition != null) + const standardizedSegmentDefinition = getStandardizedChatWrapperSegmentDefinition(chatWrapper.settings, segmentType); + if (standardizedSegmentDefinition != null) segmentDefinitions.set(segmentType, { - prefix: LlamaText(segmentDefinition.prefix).toString(), - suffix: segmentDefinition.suffix != null - ? LlamaText(segmentDefinition.suffix).toString() + prefix: standardizedSegmentDefinition.prefix != null + ? LlamaText(standardizedSegmentDefinition.prefix).toString() + : undefined, + suffix: standardizedSegmentDefinition.suffix != null + ? LlamaText(standardizedSegmentDefinition.suffix).toString() : undefined }); } @@ -653,7 +658,7 @@ function segmentModelResponseText, closeAllSegments?: string @@ -665,7 +670,8 @@ function segmentModelResponseText(); for (const [segmentType, {prefix, suffix}] of segmentDefinitions) { - separatorActions.set(prefix, {type: "prefix", segmentType}); + if (prefix != null) + separatorActions.set(prefix, {type: "prefix", segmentType}); if (suffix != null) separatorActions.set(suffix, {type: "suffix", segmentType}); diff --git a/src/utils/formatModelParameterCount.ts b/src/utils/formatModelParameterCount.ts new file mode 100644 index 00000000..4826f217 --- /dev/null +++ b/src/utils/formatModelParameterCount.ts @@ -0,0 +1,18 @@ +/** + * Format a parameter count as compact text, such as `560M`, `2B`, or `2T`. + * Promotes `950K` to `1M`, `950M` to `1B`, `980B` to `1T`, and `980T` to `1Q`. + */ +export function formatModelParameterCount(parameters: number): string { + if (parameters < 1000) + return String(parameters); + else if (parameters < 950 * 1000) + return Math.round(parameters / 1000) + "K"; + else if (parameters < 950 * (1000 ** 2)) + return Math.round(parameters / (1000 ** 2)) + "M"; + else if (parameters < 980 * (1000 ** 3)) + return Math.round(parameters / (1000 ** 3)) + "B"; + else if (parameters < 980 * (1000 ** 4)) + return Math.round(parameters / (1000 ** 4)) + "T"; + + return Math.round(parameters / (1000 ** 5)) + "Q"; +} diff --git a/src/utils/getChatWrapperSegmentDefinition.ts b/src/utils/getChatWrapperSegmentDefinition.ts deleted file mode 100644 index 4d699835..00000000 --- a/src/utils/getChatWrapperSegmentDefinition.ts +++ /dev/null @@ -1,14 +0,0 @@ -import {ChatModelSegmentType, ChatWrapperSettings, ChatWrapperSettingsSegment} from "../types.js"; - -export function getChatWrapperSegmentDefinition( - chatWrapperSetting: ChatWrapperSettings, - segmentType: ChatModelSegmentType -): ChatWrapperSettingsSegment | undefined { - if (segmentType === "thought") - return chatWrapperSetting.segments?.thought; - else if (segmentType === "comment") - return chatWrapperSetting.segments?.comment; - - void (segmentType satisfies never); - return undefined; -} diff --git a/src/utils/getStandardizedChatWrapperSegmentDefinition.ts b/src/utils/getStandardizedChatWrapperSegmentDefinition.ts new file mode 100644 index 00000000..66a3a65e --- /dev/null +++ b/src/utils/getStandardizedChatWrapperSegmentDefinition.ts @@ -0,0 +1,27 @@ +import {ChatModelSegmentType, ChatWrapperSettings} from "../types.js"; +import {LlamaText} from "./LlamaText.js"; + +export function getStandardizedChatWrapperSegmentDefinition( + chatWrapperSetting: ChatWrapperSettings, + segmentType: ChatModelSegmentType +): StandardizedChatWrapperSettingsSegment | undefined { + if (segmentType === "thought") { + const thoughtSegment = chatWrapperSetting.segments?.thought; + if (typeof thoughtSegment?.prefix === "object" && !LlamaText.isLlamaText(thoughtSegment.prefix) && thoughtSegment.prefix.type != null) + return { + ...thoughtSegment, + prefix: undefined + }; + + return thoughtSegment as StandardizedChatWrapperSettingsSegment; + } else if (segmentType === "comment") + return chatWrapperSetting.segments?.comment; + + void (segmentType satisfies never); + return undefined; +} + +export type StandardizedChatWrapperSettingsSegment = { + readonly prefix?: string | LlamaText, + readonly suffix?: string | LlamaText +}; diff --git a/src/utils/optionsMatrix.ts b/src/utils/optionsMatrix.ts index 281e38ed..bbd2c32c 100644 --- a/src/utils/optionsMatrix.ts +++ b/src/utils/optionsMatrix.ts @@ -14,7 +14,9 @@ * 2 y * ``` */ -export function* optionsMatrix>(options: {[K in keyof T]: T[K][]}): Generator<{[K in keyof T]: T[K]}> { +export function* optionsMatrix>( + options: {[K in keyof T]: readonly T[K][]} +): Generator<{[K in keyof T]: T[K]}> { const keys: Array = Object.keys(options); const indexes = keys.map(() => 0); @@ -83,7 +85,7 @@ export function* optionsMatrix>(options: {[K * ``` */ export function tryMatrix, R>( - options: {[K in keyof T]: T[K][]}, + options: {[K in keyof T]: readonly T[K][]}, callback: (options: {[K in keyof T]: T[K]}) => R ): R { let nextOption: {[K in keyof T]: T[K]} | undefined = undefined; diff --git a/src/utils/parseModelFileName.ts b/src/utils/parseModelFileName.ts index 9de3efe9..b46fa3f5 100644 --- a/src/utils/parseModelFileName.ts +++ b/src/utils/parseModelFileName.ts @@ -1,4 +1,4 @@ -import {ggufQuantNames} from "../gguf/utils/ggufQuantNames.js"; +import {ggufFileQuantNamesSet} from "../gguf/utils/ggufQuantNames.js"; export function parseModelFileName(filename: string) { const parts = filename.split("-"); @@ -36,13 +36,13 @@ export function parseModelFileName(filename: string) { if (parts.length > 0 && (quantization == null || quantization === "")) { const lastPart = parts.at(-1)!.toUpperCase(); - if (ggufQuantNames.has(lastPart)) { + if (ggufFileQuantNamesSet.has(lastPart)) { quantization = lastPart; parts.pop(); } } - if (quantization != null && !ggufQuantNames.has(quantization)) + if (quantization != null && !ggufFileQuantNamesSet.has(quantization)) quantization = undefined; if (quantization == null) { @@ -54,7 +54,7 @@ export function parseModelFileName(filename: string) { for (const part of potentialParts) { const upperPart = part.toUpperCase(); - if (ggufQuantNames.has(upperPart)) { + if (ggufFileQuantNamesSet.has(upperPart)) { quantization = upperPart; break; } diff --git a/src/utils/parseModelUri.ts b/src/utils/parseModelUri.ts index 2a46050f..35f0f3bc 100644 --- a/src/utils/parseModelUri.ts +++ b/src/utils/parseModelUri.ts @@ -3,7 +3,7 @@ import prettyMilliseconds from "pretty-ms"; import {normalizeGgufDownloadUrl} from "../gguf/utils/normalizeGgufDownloadUrl.js"; import {getFilenameForBinarySplitGgufPartUrls, resolveBinarySplitGgufPartUrls} from "../gguf/utils/resolveBinarySplitGgufPartUrls.js"; import {createSplitPartFilename, getGgufSplitPartsInfo} from "../gguf/utils/resolveSplitGgufParts.js"; -import {ggufQuantNames} from "../gguf/utils/ggufQuantNames.js"; +import {ggufFileQuantNamesSet} from "../gguf/utils/ggufQuantNames.js"; import {isUrl} from "./isUrl.js"; import {ModelFileAccessTokens, resolveModelFileAccessTokensTryHeaders} from "./modelFileAccessTokens.js"; import {isHuggingFaceUrl, ModelDownloadEndpoints, resolveHuggingFaceEndpoint} from "./modelDownloadEndpoints.js"; @@ -164,7 +164,7 @@ export async function resolveParsedModelUri( return defaultHuggingFaceFileQuantization; const quantizationText = parseModelFileName(filename).quantization; - if (quantizationText != null && ggufQuantNames.has(quantizationText)) + if (quantizationText != null && ggufFileQuantNamesSet.has(quantizationText)) return quantizationText; return ""; @@ -220,12 +220,12 @@ async function fetchHuggingFaceModelManifest({ ...await resolveModelFileAccessTokensTryHeaders(manifestUrl, tokens, endpoints) ]; let rateLimitPendingRetries = 0; - + for (let i = 0; i < headersToTry.length * (1 + rateLimitPendingRetries); i++) { const headers = headersToTry[i % headersToTry.length]; if (headers == null) continue; - + let response: Awaited> | undefined; try { response = await fetch(manifestUrl, { @@ -326,7 +326,7 @@ function parseHuggingFaceUriContent(uri: string, fullUri: string, endpoints: Mod const actualTag = tagParts.length > 0 ? [tag, ...tagParts].join(":").trimEnd() : (tag ?? "").trimEnd(); - const assumedQuant = ggufQuantNames.has(actualTag.toUpperCase()) + const assumedQuant = ggufFileQuantNamesSet.has(actualTag.toUpperCase()) ? actualTag.toUpperCase() : undefined; const resolvedTag = assumedQuant != null diff --git a/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts b/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts index 69c8934f..72c6f31d 100644 --- a/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts +++ b/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts @@ -1593,8 +1593,8 @@ describe("functionary", () => { totalRam: s1GB * 8, freeRam: s1GB * 8 }); - expect(res.gpuLayers).to.toMatchInlineSnapshot("9"); - expect(res.contextSize).to.toMatchInlineSnapshot("7424"); + expect(res.gpuLayers).to.toMatchInlineSnapshot("0"); + expect(res.contextSize).to.toMatchInlineSnapshot("8192"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); expect(res.contextSize).to.be.gte(contextSize); } @@ -1690,7 +1690,7 @@ describe("functionary", () => { }); expect(res.gpuLayers).to.toMatchInlineSnapshot("9"); expect(res.contextSize).to.toMatchInlineSnapshot("7424"); - expect(res.useMmap).to.toMatchInlineSnapshot("false"); + expect(res.useMmap).to.toMatchInlineSnapshot("true"); expect(res.contextSize).to.be.gte(contextSize); } { diff --git a/test/modelDependent/stableCode/stableCodeModelGpuLayersOptions.test.ts b/test/modelDependent/stableCode/stableCodeModelGpuLayersOptions.test.ts index f11e0647..9a6ee495 100644 --- a/test/modelDependent/stableCode/stableCodeModelGpuLayersOptions.test.ts +++ b/test/modelDependent/stableCode/stableCodeModelGpuLayersOptions.test.ts @@ -619,9 +619,9 @@ describe("stableCode", () => { totalVram: s1GB * 2, freeVram: s1GB * 1 }); - expect(res.gpuLayers).to.toMatchInlineSnapshot("0"); - expect(res.contextSize).to.toMatchInlineSnapshot("16384"); - expect(res.useMmap).to.toMatchInlineSnapshot("true"); + expect(res.gpuLayers).to.toMatchInlineSnapshot("8"); + expect(res.contextSize).to.toMatchInlineSnapshot("6144"); + expect(res.useMmap).to.toMatchInlineSnapshot("false"); expect(res.contextSize).to.be.gte(contextSize); } { diff --git a/test/standalone/chatWrappers/MuseChatWrapper.test.ts b/test/standalone/chatWrappers/MuseChatWrapper.test.ts new file mode 100644 index 00000000..900b1b3c --- /dev/null +++ b/test/standalone/chatWrappers/MuseChatWrapper.test.ts @@ -0,0 +1,200 @@ +import {describe, expect, test} from "vitest"; +import {ChatHistoryItem, ChatModelFunctions, MuseChatWrapper} from "../../../src/index.js"; + + +describe("MuseChatWrapper", () => { + test("should generate valid context text with a configurable reasoning strength", () => { + const chatWrapper = new MuseChatWrapper({reasoningStrength: "medium"}); + const chatHistory: ChatHistoryItem[] = [{ + type: "system", + text: "Be concise." + }, { + type: "user", + text: "Hello" + }, { + type: "model", + response: ["Hi!"] + }]; + + const {contextText} = chatWrapper.generateContextState({chatHistory}); + + expect(contextText).toMatchInlineSnapshot(` + LlamaText([ + new SpecialToken("BOS"), + new SpecialTokensText("<|start|>system<|message|>"), + "Be concise. + + Reasoning strength: medium. + + # Valid recipients: "self", "user".", + new SpecialTokensText("<|eot|><|start|>user<|message|>"), + "Hello", + new SpecialTokensText("<|eot|><|start|>assistant to=user<|message|>"), + "Hi!", + ]) + `); + }); + + test("should generate the default system message", () => { + const chatWrapper = new MuseChatWrapper({ + reasoningStrength: "low", + todayDate: new Date("2026-08-11T00:00:00Z") + }); + + const {contextText} = chatWrapper.generateContextState({ + chatHistory: [{type: "user", text: "Hello"}] + }); + + expect(contextText).toMatchInlineSnapshot(` + LlamaText([ + new SpecialToken("BOS"), + new SpecialTokensText("<|start|>system<|message|>"), + "You are a helpful AI assistant. + Knowledge cutoff: 04-01-2026. + Current date: 11-08-2026. + Reasoning strength: low. + + # Valid recipients: "self", "user".", + new SpecialTokensText("<|eot|><|start|>user<|message|>"), + "Hello", + ]) + `); + }); + + test("should replay reasoning and ATEM function calls", () => { + const chatWrapper = new MuseChatWrapper(); + const availableFunctions: ChatModelFunctions = { + "get_weather": { + description: "Get the current weather for a city.", + params: { + type: "object", + properties: { + city: {type: "string"}, + includeForecast: {type: "boolean"} + }, + required: ["city"] + } + } + }; + const chatHistory: ChatHistoryItem[] = [{ + type: "system", + text: "Answer weather questions." + }, { + type: "user", + text: "What is the weather?" + }, { + type: "model", + response: [{ + type: "segment", + segmentType: "thought", + text: "I should check the weather.", + ended: true + }, { + type: "functionCall", + name: "get_weather", + params: {city: "Tokyo", includeForecast: false}, + result: {temperature: 24} + }, "It is 24 degrees in Tokyo."] + }]; + + const {contextText} = chatWrapper.generateContextState({chatHistory, availableFunctions}); + expect(contextText).toMatchInlineSnapshot(` + LlamaText([ + new SpecialToken("BOS"), + new SpecialTokensText("<|start|>system<|message|>"), + "Answer weather questions. + + Reasoning strength: high. + + In this environment you have access to a set of tools you can use to answer the user's question. + + You can invoke a function by writing a ", + new SpecialTokensText(""""), + " block like the following: + ", + new SpecialTokensText(""), + " + ", + new SpecialTokensText(""), + " + ", + new SpecialTokensText("$PARAMETER_VALUE"), + " + ... + ", + new SpecialTokensText(""), + " + ", + new SpecialTokensText(""), + " + + String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions. + Here are the functions available in JSONSchema format: + // Tool metadata + {"name": "get_weather", "description": ""} + // Function schemas + {"name": "get_weather", "description": "Get the current weather for a city.", "parameters": {"type": "object", "properties": {"city": {"type": "string"}, "includeForecast": {"type": "boolean"}}, "required": ["city"]}} + Here's an example of how to call a function in the tool set: + (If the tool namespace is not specified, invoke the function directly as \`example_function_name\` rather than \`example_tool_name.example_function_name\`) + + to=example_tool_name.example_function_name + + ", + new SpecialTokensText(""), + " + ", + new SpecialTokensText(""), + " + ", + new SpecialTokensText(""), + "value_1", + new SpecialTokensText(""), + " + ", + new SpecialTokensText(""), + "This is the value for the second parameter + that can span + "multiple" lines + ", + new SpecialTokensText(""), + " + ", + new SpecialTokensText(""), + " + ", + new SpecialTokensText(""), + " + + # Valid recipients: "self", "get_weather.*", "user".", + new SpecialTokensText("<|eot|><|start|>user<|message|>"), + "What is the weather?", + new SpecialTokensText("<|eot|><|start|>assistant to=self<|message|>"), + "I should check the weather.", + new SpecialTokensText("<|eom|><|start|>assistant to="), + "get_weather", + new SpecialTokensText("<|message|> + + "), + "{"city": "Tokyo", "includeForecast": false}", + new SpecialTokensText(" + + "), + new SpecialToken("EOT"), + new SpecialTokensText("<|start|>tool "), + "get_weather", + new SpecialTokensText("<|message|> + "), + "{"temperature": 24}", + new SpecialTokensText(" + "), + new SpecialToken("EOT"), + new SpecialTokensText("<|start|>assistant to=user<|message|>"), + "It is 24 degrees in Tokyo.", + ]) + `); + }); +}); diff --git a/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts b/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts index 83d5b88a..a294f010 100644 --- a/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts +++ b/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts @@ -1703,6 +1703,7 @@ describe("JinjaTemplateChatWrapper", () => { expect(chatWrapper.keepOnlyLastThought).to.be.eql(true); expect(chatWrapper.settings.segments?.thought).to.eql({ openOnResponseStart: true, + reopenAfterFunctionCalls: false, prefix: LlamaText(new SpecialTokensText("\n")), suffix: LlamaText(new SpecialTokensText("\n\n\n")) }); @@ -1724,6 +1725,7 @@ describe("JinjaTemplateChatWrapper", () => { const chatWrapper = new JinjaTemplateChatWrapper({template}); expect(chatWrapper.settings.segments?.thought).to.eql({ openOnResponseStart: true, + reopenAfterFunctionCalls: false, prefix: LlamaText(new SpecialTokensText("\n")), suffix: LlamaText(new SpecialTokensText("\n\n")) }); diff --git a/test/standalone/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.test.ts b/test/standalone/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.test.ts index 8c49cc17..8b4bba46 100644 --- a/test/standalone/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.test.ts +++ b/test/standalone/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.test.ts @@ -210,11 +210,14 @@ describe("JinjaTemplateChatWrapper", () => { }, "segments": { "thought": { - "prefix": LlamaText([ - new SpecialTokensText(""), - ]), + "openOnResponseStart": true, + "prefix": { + "type": "openedOnStart", + }, + "reopenAfterFunctionCalls": true, "suffix": LlamaText([ - new SpecialTokensText(""), + new SpecialTokensText(" + "), ]), }, }, diff --git a/test/standalone/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.test.ts b/test/standalone/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.test.ts new file mode 100644 index 00000000..619b010c --- /dev/null +++ b/test/standalone/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.test.ts @@ -0,0 +1,94 @@ +import {describe, expect, test} from "vitest"; +import {JinjaTemplateChatWrapper} from "../../../../../src/index.js"; +import {functionGemma270mJinjaTemplate, glm4_5airJinjaTemplate, glm4_7flashJinjaTemplate, LagunaXS2_1JinjaTemplate, lfm2_5JinjaTemplate} from "../../utils/jinjaTemplates.js"; + + +describe("JinjaTemplateChatWrapper", () => { + describe("extractSegmentSettingsFromTokenizerAndChatTemplate", () => { + test("lfm2_5JinjaTemplate", () => { + const chatWrapper = new JinjaTemplateChatWrapper({ + template: lfm2_5JinjaTemplate + }); + expect(chatWrapper.settings.segments).toMatchInlineSnapshot(` + { + "thought": { + "prefix": LlamaText([ + new SpecialTokensText(""), + ]), + "suffix": LlamaText([ + new SpecialTokensText(""), + ]), + }, + } + `); + }); + + test("functionGemma270mJinjaTemplate", () => { + const chatWrapper = new JinjaTemplateChatWrapper({ + template: functionGemma270mJinjaTemplate + }); + expect(chatWrapper.settings.segments).toMatchInlineSnapshot("{}"); + }); + + test("glm4_7flashJinjaTemplate", () => { + const chatWrapper = new JinjaTemplateChatWrapper({ + template: glm4_7flashJinjaTemplate + }); + expect(chatWrapper.settings.segments).toMatchInlineSnapshot(` + { + "thought": { + "prefix": LlamaText([ + new SpecialTokensText(""), + ]), + "suffix": LlamaText([ + new SpecialTokensText(""), + ]), + }, + } + `); + }); + + test("glm4_5airJinjaTemplate", () => { + const chatWrapper = new JinjaTemplateChatWrapper({ + template: glm4_5airJinjaTemplate + }); + expect(chatWrapper.settings.segments).toMatchInlineSnapshot(` + { + "thought": { + "openOnResponseStart": true, + "prefix": { + "type": "openedOnStart", + }, + "reopenAfterFunctionCalls": true, + "suffix": LlamaText([ + new SpecialTokensText(" + "), + ]), + }, + } + `); + }); + + + + test("LagunaXS2_1JinjaTemplate", () => { + const chatWrapper = new JinjaTemplateChatWrapper({ + template: LagunaXS2_1JinjaTemplate + }); + expect(chatWrapper.settings.segments).toMatchInlineSnapshot(` + { + "thought": { + "openOnResponseStart": true, + "prefix": { + "type": "openedOnStart", + }, + "reopenAfterFunctionCalls": true, + "suffix": LlamaText([ + new SpecialTokensText(""), + ]), + }, + } + `); + }); + }); +}); diff --git a/test/standalone/chatWrappers/utils/jinjaTemplates.ts b/test/standalone/chatWrappers/utils/jinjaTemplates.ts index 654702e1..1facd494 100644 --- a/test/standalone/chatWrappers/utils/jinjaTemplates.ts +++ b/test/standalone/chatWrappers/utils/jinjaTemplates.ts @@ -3211,3 +3211,293 @@ For each function call, output the function name and arguments within the follow {%- endif -%} {# Copyright 2025-present Unsloth. Apache 2.0 License. #} `.slice(1, -1); + +export const LagunaXS2_1JinjaTemplate = ` +{#- Iteration on laguna_glm_thinking_v8/chat_template.jinja -#} +{#- No formatting instructions -#} +{{- "〈|EOS|〉" -}} +{%- set enable_thinking = enable_thinking | default(false) -%} +{%- set add_generation_prompt = add_generation_prompt | default(false) -%} + +{#- ───── header (system message) ───── -#} +{#- A caller-supplied system message with empty content opts out of the default below, producing no block — used to train without a system message. -#} +{%- set system_message = "You are a helpful, conversationally-fluent assistant made by Poolside. You are here to be helpful to users through natural language conversations." -%} +{%- if messages and messages[0].role == "system" -%} + {%- set system_message = messages[0].content -%} + {%- set messages = messages[1:] -%} +{%- endif -%} + +{%- set has_sys = system_message and system_message.strip() -%} +{%- if has_sys or tools or enable_thinking -%} + {{- "" -}} + + {%- if has_sys -%} + {{- system_message.rstrip() -}} + {%- if tools -%}{{- "\n\n" -}}{%- endif -%} + {%- endif -%} + + {%- if tools -%} + {{- "### Tools\n\n" -}} + {{- "You may call functions to assist with the user query.\n" -}} + {{- "All available function signatures are listed below:\n" -}} + {{- "\n" -}} + {%- for tool in tools -%} + {{- (tool | tojson) ~ "\n" -}} + {%- endfor -%} + {{- "" -}} + {%- endif -%} + + {{- "\n" -}} +{%- endif -%} + +{#- ───── main loop ───── -#} +{%- for message in messages -%} + {%- set content = message.content if message.content is string else "" -%} + {%- if message.role == "user" -%} + {{- "" + content + "\n" -}} + {%- elif message.role == "assistant" -%} + {%- generation -%} + {{- "" -}} + {#- Extract reasoning content from message.reasoning (vLLM field name) or message.reasoning_content -#} + {%- set reasoning_content = '' -%} + {%- if message.reasoning is string -%} + {%- set reasoning_content = message.reasoning -%} + {%- elif message.reasoning_content is string -%} + {%- set reasoning_content = message.reasoning_content -%} + {%- endif -%} + {#- Display reasoning content for all messages if enable_thinking -#} + {%- if enable_thinking -%} + {{- '' + reasoning_content + '' -}} + {%- else -%} + {{- '' -}} + {%- endif -%} + {#- Display main content (trailing newline only when no tool_calls follow) -#} + {%- if content -%} + {{- content -}} + {%- endif -%} + {%- if message.tool_calls -%} + {%- for tool_call in message.tool_calls -%} + {%- set function_data = tool_call.function -%} + {{- '' + function_data.name -}} + {%- set _args = function_data.arguments -%} + {%- for k, v in _args.items() -%} + {{- "" ~ k ~ "" -}} + {{- "" -}}{{- v | tojson(ensure_ascii=False) if v is not string else v -}}{{- "" -}} + {%- endfor -%} + {{- "" -}} + {%- endfor -%} + {%- endif -%} + {{- "\n" -}} + {%- endgeneration -%} + {%- elif message.role == "tool" -%} + {{- "" + content + "\n" -}} + {%- elif message.role == "system" -%} + {#- Render additional system messages (the first one, if any, is handled separately in the header and was sliced off above) -#} + {{- "" + content + "\n" -}} + {%- endif -%} +{%- endfor -%} +{#- ───── generation prompt ───── -#} +{%- if add_generation_prompt -%} + {{- "" -}} + {#- ───── Include reasoning mode directive ───── -#} + {%- if enable_thinking -%} + {{- '' -}} + {%- else -%} + {{- '' -}} + {%- endif -%} +{%- endif -%} +`.slice(1, -1); + +// source: https://huggingface.co/meta-models/Muse-Glimmer-30B/blob/main/chat_template.jinja +export const museGlimmerJinjaTemplate = String.raw` +{%- macro render_content(content) -%} + {%- if content is string -%} + {{- content -}} + {%- elif content is not none -%} + {%- for part in content -%} + {%- if part["type"] == "image" -%} + {{- "<|patch|>" -}} + {%- elif part["type"] == "video" -%} + {{- "<|video|>" -}} + {%- elif part["type"] == "text" -%} + {{- part["text"] -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} +{%- endmacro -%} +{%- macro render_atem(tc) -%} + {%- set args = tc.function.arguments -%} + {%- if args is not mapping -%} + {{- raise_exception("Onyx ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.") -}} + {%- endif -%} + {{- "\n\n" -}} + {%- for (k, v) in args.items() -%} + {{- "" -}} + {%- if v is boolean -%} + {%- if v -%} + {{- "true" -}} + {%- else -%} + {{- "false" -}} + {%- endif -%} + {%- elif v is none -%} + {{- "null" -}} + {%- elif v is mapping or v is iterable and v is not string -%} + {{- v | tojson -}} + {%- else -%} + {{- v -}} + {%- endif -%} + {{- "\n" -}} + {%- endfor -%} + {{- "\n" -}} +{%- endmacro -%} +{%- macro render_tool_defs(tools) -%} + {{- "In this environment you have access to a set of tools you can use to answer the user's question.\n\n" -}} + {{- "You can invoke a function by writing a \"\" block like the following:\n" -}} + {{- "\n\n$PARAMETER_VALUE\n...\n\n\n\n" -}} + {{- "String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n" -}} + {{- "Here are the functions available in JSONSchema format:\n" -}} + {{- "// Tool metadata\n" -}} + {%- set nsns = namespace(seen=[]) -%} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set tns = fn.name.split(".")[0] -%} + {%- if tns not in nsns.seen -%} + {%- set nsns.seen = nsns.seen + [tns] -%} + {%- endif -%} + {%- endfor -%} + {%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%} + {%- for tns in nsns.seen -%} + {{- "{\"name\": " + tns | tojson + ", \"description\": " + (nd[tns] if tns in nd else "") | tojson + "}\n" -}} + {%- endfor -%} + {{- "// Function schemas" -}} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {{- "\n{\"name\": " + fn.name | tojson + ", \"description\": " + fn.description | tojson + ", \"parameters\": " + fn.parameters | tojson + "}" -}} + {%- endfor -%} + {{- "\n\nHere's an example of how to call a function in the tool set:\n" -}} + {{- "(If the tool namespace is not specified, invoke the function directly as \`example_function_name\` rather than \`example_tool_name.example_function_name\`)\n\n" -}} + {{- "to=example_tool_name.example_function_name\n\n" -}} + {{- "\n\n" -}} + {{- "value_1\n" -}} + {{- "This is the value for the second parameter\nthat can span\n\"multiple\" lines\n\n" -}} + {{- "\n" -}} +{%- endmacro -%} +{%- macro render_reasoning() -%} + {%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else "high" -%} + {{- "Reasoning strength: " + rs + "." -}} +{%- endmacro -%} +{%- macro render_system_meta(tools) -%} + {%- set rns = namespace(recipients=["\"self\""], nslist=[]) -%} + {%- if tools -%} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set tns = fn.name.split(".")[0] -%} + {%- if tns not in rns.nslist -%} + {%- set rns.nslist = rns.nslist + [tns] -%} + {%- endif -%} + {%- endfor -%} + {%- for tns in rns.nslist -%} + {%- set rns.recipients = rns.recipients + ["\"" + tns + ".*\""] -%} + {%- endfor -%} + {%- endif -%} + {%- set rns.recipients = rns.recipients + ["\"user\""] -%} + {{- "# Valid recipients: " + rns.recipients | join(", ") + "." -}} +{%- endmacro -%} +{{- bos_token -}} +{%- set ns = namespace(has_system=false) -%} +{%- for m in messages -%} + {%- if m["role"] == "system" -%} + {%- set ns.has_system = true -%} + {%- endif -%} +{%- endfor -%} +{%- if not ns.has_system -%} + {{- "<|start|>system<|message|>You are a helpful AI assistant." -}} + {%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else "2026-01-04" -%} + {{- "\nKnowledge cutoff: " + kc + "." -}} + {%- if current_date is defined and current_date -%} + {{- "\nCurrent date: " + current_date + "." -}} + {%- elif strftime_now is defined -%} + {{- "\nCurrent date: " + strftime_now("%Y-%m-%d") + "." -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_reasoning() -}} + {%- if tools -%} + {{- "\n\n" -}} + {{- render_tool_defs(tools) -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_system_meta(tools) -}} + {{- "<|eot|>" -}} +{%- endif -%} +{%- for message in messages -%} + {%- set role = message["role"] -%} + {%- set end_token = "<|eom|>" if not loop.last and messages[loop.index0 + 1]["role"] == role else "<|eot|>" -%} + {%- if role == "system" -%} + {{- "<|start|>system<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "\n\n" -}} + {{- render_reasoning() -}} + {%- if tools -%} + {{- "\n\n" -}} + {{- render_tool_defs(tools) -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_system_meta(tools) -}} + {{- "<|eot|>" -}} + {%- elif role == "user" -%} + {{- "<|start|>user<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "<|eot|>" -}} + {%- elif role == "tool" -%} + {%- set tname = message.get("name") -%} + {%- if not tname -%} + {%- set tcid = message.get("tool_call_id") -%} + {%- set rns = namespace(name=tcid if tcid else "") -%} + {%- for m in messages -%} + {%- if m.get("tool_calls") -%} + {%- for tc in m["tool_calls"] -%} + {%- if tcid is not none and tc.id is defined and tc.id == tcid -%} + {%- set rns.name = tc.function.name -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endfor -%} + {%- set tname = rns.name -%} + {%- endif -%} + {{- "<|start|>tool " + tname + "<|message|>\n" -}} + {{- render_content(message["content"]) -}} + {{- "\n<|eot|>" -}} + {%- elif role == "assistant" -%} + {%- if message.get("reasoning_content") -%} + {{- "<|start|>assistant to=self<|message|>" + message["reasoning_content"] + "<|eom|>" -}} + {%- endif -%} + {%- if message.get("tool_calls") -%} + {%- for tc in message["tool_calls"] -%} + {{- "<|start|>assistant to=" + tc.function.name + "<|message|>" -}} + {{- render_atem(tc) -}} + {%- if loop.last -%} + {{- end_token -}} + {%- else -%} + {{- "<|eom|>" -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {%- set recipient = message.get("recipient") or "user" -%} + {%- set end_turn = message.get("end_turn") -%} + {%- if end_turn is none -%} + {%- set end_turn = not (recipient and recipient != "user") -%} + {%- endif -%} + {{- "<|start|>assistant" -}} + {%- if recipient -%} + {{- " to=" + recipient -}} + {%- endif -%} + {{- "<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "<|eot|>" if end_turn else "<|eom|>" -}} + {%- endif -%} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- "<|start|>assistant" -}} +{%- endif -%} +`.slice(1, -1).replaceAll("\\`", "`"); diff --git a/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts b/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts index 39da3d5d..4704d89e 100644 --- a/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts +++ b/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts @@ -2,11 +2,11 @@ import {describe, expect, test} from "vitest"; import { AlpacaChatWrapper, ChatMLChatWrapper, DeepSeekChatWrapper, FalconChatWrapper, FunctionaryChatWrapper, GemmaChatWrapper, Gemma4ChatWrapper, GeneralChatWrapper, Llama2ChatWrapper, Llama3_1ChatWrapper, MistralChatWrapper, QwenChatWrapper, - resolveChatWrapper, HarmonyChatWrapper + resolveChatWrapper, HarmonyChatWrapper, MuseChatWrapper } from "../../../../src/index.js"; import { harmonyJinjaTemplate, harmonyJinjaTemplate2, harmonyJinjaTemplate3, harmonyJinjaTemplate4, harmonyJinjaTemplate5, - gemma4JinjaTemplate1, gemma4JinjaTemplate2, gemma4JinjaTemplate3 + gemma4JinjaTemplate1, gemma4JinjaTemplate2, gemma4JinjaTemplate3, museGlimmerJinjaTemplate } from "./jinjaTemplates.js"; @@ -992,4 +992,17 @@ describe("resolveChatWrapper", () => { }); expect(chatWrapper).to.be.instanceof(HarmonyChatWrapper); }); + + test("should resolve to specialized MuseChatWrapper ", {timeout: 1000 * 60 * 60 * 2}, async () => { + const chatWrapper = resolveChatWrapper({ + customWrapperSettings: { + jinjaTemplate: { + template: museGlimmerJinjaTemplate + } + }, + fallbackToOtherWrappersOnJinjaError: false + }); + + expect(chatWrapper).to.be.instanceof(MuseChatWrapper); + }); });