diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 2c75123a..2e3e7bc3 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -136,11 +136,11 @@ jobs: run: | sudo apt-get update sudo apt-get install ninja-build libtbb-dev - + wget -c https://github.com/Kitware/CMake/releases/download/v4.2.1/cmake-4.2.1-linux-x86_64.tar.gz sudo tar --strip-components=1 -C /usr/local -xzf cmake-4.2.1-linux-x86_64.tar.gz rm -f ./cmake-4.2.1-linux-x86_64.tar.gz - + cmake --version - name: Install dependencies on Ubuntu (2) @@ -148,13 +148,13 @@ jobs: run: | sudo apt-get update sudo apt-get install ninja-build libtbb-dev - + wget -c https://github.com/Kitware/CMake/releases/download/v4.2.1/cmake-4.2.1-linux-x86_64.tar.gz sudo tar --strip-components=1 -C /usr/local -xzf cmake-4.2.1-linux-x86_64.tar.gz rm -f ./cmake-4.2.1-linux-x86_64.tar.gz - + cmake --version - + - name: Install dependencies on Ubuntu (arm64) if: matrix.config.name == 'Ubuntu (arm64)' env: @@ -163,7 +163,7 @@ jobs: wget https://apt.llvm.org/llvm.sh chmod +x llvm.sh sudo ./llvm.sh "${LLVM_VERSION}" - + sudo apt-get update sudo apt-get install -y ninja-build libtbb-dev "clang-${LLVM_VERSION}" "lld-${LLVM_VERSION}" "llvm-${LLVM_VERSION}" g++-arm-linux-gnueabihf gcc-arm-linux-gnueabihf @@ -171,11 +171,11 @@ jobs: sudo update-alternatives --install /usr/bin/c++ c++ "/usr/bin/clang++-${LLVM_VERSION}" 100 sudo update-alternatives --install /usr/bin/gcc gcc "/usr/bin/clang-${LLVM_VERSION}" 100 sudo update-alternatives --install /usr/bin/g++ g++ "/usr/bin/clang++-${LLVM_VERSION}" 100 - + wget -c https://github.com/Kitware/CMake/releases/download/v4.2.1/cmake-4.2.1-linux-aarch64.tar.gz sudo tar --strip-components=1 -C /usr/local -xzf cmake-4.2.1-linux-aarch64.tar.gz rm -f ./cmake-4.2.1-linux-aarch64.tar.gz - + cc --version c++ --version gcc --version @@ -183,7 +183,7 @@ jobs: which arm-linux-gnueabihf-gcc which arm-linux-gnueabihf-g++ - + cmake --version # llama.cpp's ggml-cpu uses the RVV _Float16 (zvfh) vector intrinsics, @@ -238,7 +238,7 @@ jobs: run: | # $env:VULKAN_VERSION = & curl.exe -fsSL https://vulkan.lunarg.com/sdk/latest/linux.txt # Write-Host "Downloading Vulkan SDK version $env:VULKAN_VERSION" - + curl.exe -o $env:RUNNER_TEMP/VulkanSDK-Installer.exe -L "https://sdk.lunarg.com/sdk/download/${env:VULKAN_VERSION}/windows/vulkansdk-windows-X64-${env:VULKAN_VERSION}.exe" & "$env:RUNNER_TEMP\VulkanSDK-Installer.exe" --accept-licenses --default-answer --confirm-command install Add-Content $env:GITHUB_ENV "VULKAN_SDK=C:\VulkanSDK\${env:VULKAN_VERSION}" @@ -251,7 +251,7 @@ jobs: run: | # export VULKAN_VERSION="$(curl -fsSL https://vulkan.lunarg.com/sdk/latest/linux.txt)" # echo "Downloading Vulkan SDK version ${VULKAN_VERSION}" - + sudo apt-get update sudo apt-get install xz-utils curl --no-progress-meter "https://sdk.lunarg.com/sdk/download/${VULKAN_VERSION}/linux/vulkan_sdk.tar.xz" -o "/opt/vulkan-sdk.tar.xz" @@ -286,50 +286,50 @@ jobs: fi npx zx -y <<'EOF' - + async function getLatestNodeVersions(maxDate) { const res = await fetch("https://nodejs.org/dist/index.json"); const data = await res.json(); const versions = new Map(); let latestVersion = null; - + for (const version of data) { const majorVersion = Number(version.version.split(".")[0].slice("v".length)); const versionDate = new Date(version.date); - + if (maxDate != null && versionDate.getTime() > maxDate) continue; - + if (!versions.has(majorVersion)) { versions.set(majorVersion, version.version); } - + if (latestVersion === null || majorVersion > latestVersion) { latestVersion = majorVersion; } } - + return {versions, latestVersion}; } - + const {versions: latestNodeVersions} = await getLatestNodeVersions(Date.now() - 1000 * 60 * 60 * 24 * 14); - + const nodeVersion = latestNodeVersions.get(20); const windowsOnArmNodeVersion = latestNodeVersions.get(20); - + if (nodeVersion == null || windowsOnArmNodeVersion == null) { throw new Error("Could not find node versions"); } - + $.verbose = true; await $`mkdir -p bins`; - + async function buildBinary(arch, flags = [], nodeTarget = nodeVersion) { console.log(`Building ${arch} for node ${nodeTarget} with flags`, flags); - + await $`node ./dist/cli/cli.js source build --ciMode --noUsageExample --arch ${arch} --nodeTarget ${nodeVersion} ${flags}`; } - + // build binaries if (process.env.ARTIFACT_NAME === "win-1") { await buildBinary("x64", ["--gpu", "false"]); @@ -354,7 +354,7 @@ jobs: } else if (process.env.ARTIFACT_NAME === "linux-riscv64") { await buildBinary("riscv64", ["--gpu", "false"]); } - + // move binaries to bins const localBuildsDirectoryPath = path.join(process.cwd(), "llama", "localBuilds"); const llamaBinsDirectoryPath = path.join(process.cwd(), "bins"); @@ -364,13 +364,13 @@ jobs: path.join(llamaBinsDirectoryPath, folderName) ); } - + if (process.env.ARTIFACT_NAME === "win-2") { await fs.move( path.join(llamaBinsDirectoryPath, "win-x64-cuda"), path.join(llamaBinsDirectoryPath, "win-x64-cuda-2") ); - + if (!(await fs.pathExists(path.join(llamaBinsDirectoryPath, "win-x64-cuda-2", "ggml-cuda.dll")))) { throw new Error("ggml-cuda.dll not found in win-x64-cuda-2"); } @@ -379,15 +379,15 @@ jobs: path.join(llamaBinsDirectoryPath, "linux-x64-cuda"), path.join(llamaBinsDirectoryPath, "linux-x64-cuda-2") ); - + if (!(await fs.pathExists(path.join(llamaBinsDirectoryPath, "linux-x64-cuda-2", "libggml-cuda.so")))) { throw new Error("libggml-cuda.so not found in linux-x64-cuda-2"); } } - + await $`echo "Built binaries:"`; await $`ls bins`; - + EOF # - name: Cache UPX @@ -530,7 +530,7 @@ jobs: - name: Install dependencies on ubuntu run: | sudo apt-get update - sudo apt-get install ninja-build cmake + sudo apt-get install ninja-build cmake - name: Install modules run: npm ci @@ -650,30 +650,30 @@ jobs: mkdir -p bins mv artifacts/bins-*/* bins/ mv artifacts/build dist/ - + mkdir -p ./bins/win-x64-cuda/fallback mv ./bins/win-x64-cuda-2/ggml-cuda.dll ./bins/win-x64-cuda/fallback/ggml-cuda.dll rm -rf ./bins/win-x64-cuda-2 - + mkdir -p ./bins/linux-x64-cuda/fallback mv ./bins/linux-x64-cuda-2/libggml-cuda.so ./bins/linux-x64-cuda/fallback/libggml-cuda.so rm -rf ./bins/linux-x64-cuda-2 - + cp -r artifacts/llama.cpp/llama.cpp/grammars llama/grammars - + rm -f ./llama/binariesGithubRelease.json mv artifacts/llama.cpp/binariesGithubRelease.json ./llama/binariesGithubRelease.json - + rm -f ./llama/llama.cpp.info.json mv artifacts/llama.cpp/llama.cpp.info.json ./llama/llama.cpp.info.json - + rm -f ./llama/gitRelease.bundle mv artifacts/llama.cpp/gitRelease.bundle ./llama/gitRelease.bundle - + mv artifacts/build-templates templates/packed/ rm -f ./templates/package.json rm -f ./templates/package-lock.json - + echo "Built binaries:" ls bins - name: Move binaries to standalone prebuilt binary modules @@ -690,7 +690,7 @@ jobs: GH_RELEASE_REF: ${{ github.ref }} run: | export DRY_RUN_RESULT_FILE_PATH="$(pwd)/semanticReleaseDryRunReleaseResult.json" - + git apply --ignore-whitespace ./scripts/patches/@semantic-release+github+11.0.0.patch npx semantic-release - name: Set npm package url to GITHUB_OUTPUT @@ -709,13 +709,13 @@ jobs: if: steps.set-npm-url.outputs.npm-url != '' run: | export DEPLOYED_PACKAGE_VERSION=$(cat .semanticRelease.npmPackage.deployedVersion.txt) - + pushd packages/create-node-llama-cpp npm ci --ignore-scripts popd - + npx --no vite-node ./scripts/prepareCreateNodeLlamaCppModuleForPublish.ts --packageVersion "$DEPLOYED_PACKAGE_VERSION" - + pushd packages/create-node-llama-cpp npm run build - name: Release `create-node-llama-cpp` module @@ -725,7 +725,7 @@ jobs: GH_RELEASE_REF: ${{ github.ref }} run: | cd packages/create-node-llama-cpp - + if [ "$GH_RELEASE_REF" == "refs/heads/beta" ]; then npm publish --tag beta else @@ -817,12 +817,12 @@ jobs: RELEASE_TAG: ${{ needs.release.outputs.package-version }} run: | shopt -s nullglob - + for file in ./electron-app-example/release/*.{dmg,zip,exe,appx,AppImage,snap,assert,deb,tar.gz}; do echo "Adding $file to release $RELEASE_TAG" gh release upload "v$RELEASE_TAG" "$file" done - + shopt -u nullglob update-documentation-website: @@ -875,7 +875,7 @@ jobs: - name: Move artifacts run: | mv artifacts/build dist/ - + cp -r artifacts/llama.cpp/llama.cpp llama/llama.cpp rm -f ./llama/binariesGithubRelease.json @@ -898,7 +898,7 @@ jobs: run: | export DOCS_PACKAGE_VERSION="$(cat ./docsVersion.txt)" echo "Package version: $DOCS_PACKAGE_VERSION" - + npm run docs:build - name: Upload docs to GitHub Pages uses: actions/upload-pages-artifact@v5 diff --git a/.gitignore b/.gitignore index 781d098b..76d2b307 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,7 @@ node_modules /.env /.eslintcache /.vitepress/.cache +/.vitepress/.temp /test/.models /test/temp /test/.temp diff --git a/.vitepress/config.ts b/.vitepress/config.ts index ade9ae4c..381b3c7e 100644 --- a/.vitepress/config.ts +++ b/.vitepress/config.ts @@ -4,7 +4,7 @@ import process from "process"; import {fileURLToPath} from "url"; import fs from "fs-extra"; import {createContentLoader, defineConfig, HeadConfig, Plugin as VitepressPlugin} from "vitepress"; -import {transformerTwoslash} from "@shikijs/vitepress-twoslash"; +import {rendererFloatingVue, transformerTwoslash} from "@shikijs/vitepress-twoslash"; import ts from "typescript"; import envVar from "env-var"; import {Feed} from "feed"; @@ -19,6 +19,8 @@ import {ensureLocalImage} from "./utils/ensureLocalImage.js"; import {getExcerptFromMarkdownFile} from "./utils/getExcerptFromMarkdownFile.js"; import {getVitepressSidebar, getVitepressSidebarWithBlog} from "./config/sidebar.js"; import {getBlogPosts} from "./config/getBlogPosts.js"; +import {apiTypeLinksTransformer} from "./config/apiLinksTransformer.js"; +import {createCustomTwoSlashRenderer, preserveTwoslashLanguage} from "./config/createCustomTwoSlashRenderer.js"; import type {Element as HastElement, Parent} from "hast"; import type {Node as UnistNode} from "unist"; @@ -386,40 +388,53 @@ export default defineConfig({ "js-highlight": "javascript" }, codeTransformers: [ - transformerTwoslash({ - floatingVue: { - classFloatingPanel: "twoslash-floating vp-code" - }, - explicitTrigger: false, - filter(lang, code, options) { - return options.lang?.toLowerCase() === "typescript"; - }, - twoslashOptions: { - compilerOptions: { - ...(await fs.readJSON(path.join(__dirname, "..", "tsconfig.json"))).compilerOptions, - moduleResolution: undefined, - paths: { - "node-llama-cpp": [ - path.resolve(__dirname, "..", "dist", "index.d.ts"), - path.resolve(__dirname, "..", "src", "index.ts") + preserveTwoslashLanguage( + transformerTwoslash({ + renderer: createCustomTwoSlashRenderer({ + baseRenderer: rendererFloatingVue({ + lang: "ts", + floatingVue: { + classFloatingPanel: "twoslash-floating vp-code" + } + }), + transformers: [apiTypeLinksTransformer({resolveHref})] + }), + floatingVue: { + classFloatingPanel: "twoslash-floating vp-code" + }, + explicitTrigger: false, + filter(lang, code, options) { + return options.lang?.toLowerCase() === "typescript"; + }, + twoslashOptions: { + compilerOptions: { + ...(await fs.readJSON(path.join(__dirname, "..", "tsconfig.json"))).compilerOptions, + moduleResolution: undefined, + paths: { + "node-llama-cpp": [ + path.resolve(__dirname, "..", "dist", "index.d.ts"), + path.resolve(__dirname, "..", "src", "index.ts") + ], + "node-llama-cpp/commands": [ + path.resolve(__dirname, "..", "dist", "commands.d.ts"), + path.resolve(__dirname, "..", "src", "commands.ts") + ] + }, + typeRoots: [ + path.resolve(__dirname, "..", "node_modules"), + path.resolve(__dirname, "..", "node_modules", "@types") ], - "node-llama-cpp/commands": [ - path.resolve(__dirname, "..", "dist", "commands.d.ts"), - path.resolve(__dirname, "..", "src", "commands.ts") - ] + module: ts.ModuleKind.ES2022, + target: ts.ScriptTarget.ES2022, + moduleDetection: ts.ModuleDetectionKind.Force, + rootDir: undefined }, - typeRoots: [ - path.resolve(__dirname, "..", "node_modules"), - path.resolve(__dirname, "..", "node_modules", "@types") - ], - module: ts.ModuleKind.ES2022, - target: ts.ScriptTarget.ES2022, - moduleDetection: ts.ModuleDetectionKind.Force, - rootDir: undefined - }, - tsModule: ts - } - }) as ShikiTransformer + tsModule: ts + } + }) as ShikiTransformer + ), + + apiTypeLinksTransformer({resolveHref}) ] }, themeConfig: { diff --git a/.vitepress/config/apiLinksTransformer.ts b/.vitepress/config/apiLinksTransformer.ts new file mode 100644 index 00000000..48374646 --- /dev/null +++ b/.vitepress/config/apiLinksTransformer.ts @@ -0,0 +1,413 @@ +/* eslint import/no-unresolved: "off" */ +import typedocSidebar from "../../docs/api/typedoc-sidebar.json"; +import type {ShikiTransformer} from "shiki"; +import type {Element} from "hast"; + + +const apiSymbolLinks = typedocSidebarToSymbolMap(); + +const typeDocMemberModifiers = [ + "public", "protected", "private", "static", "abstract", "readonly", "override", "declare", "optional" +] as const; +const typeDocMemberModifiersRegex = typeDocMemberModifiers.join("|"); +const tsMemberModifiersRegex = ["public", "protected", "private", "static", "abstract", "readonly", "override", "declare"].join("|"); +const optionalModifierRegex = new RegExp(`^([\\t ]*(?:(?:${tsMemberModifiersRegex})[\\t ]+)*)optional(?=[\\t ])`, "gm"); + +export function apiTypeLinksTransformer({ + resolveHref +}: { + resolveHref(href: string, withDomain?: boolean): string +}): ShikiTransformer { + const enabled = new WeakSet(); + const originalCode = new WeakMap(); + const qualifiedSymbolLinks = new WeakMap>(); + const signaturePrefixes = new WeakMap>(); + + return { + name: "api-type-links", + enforce: "pre", + + preprocess(code, options) { + if (options.lang.toLowerCase() !== "ts") + return; + + enabled.add(this.meta); + + options.includeExplanation = "scopeName"; + options.mergeWhitespaces = "never"; + + const links = getQualifiedSymbolLinks(code); + if (links.size > 0) + qualifiedSymbolLinks.set(this.meta, links); + + const normalizedSignature = normalizeTypeDocSignature(code); + if (normalizedSignature != null) { + options.grammarContextCode = "declare interface __TypeDoc {\n"; + + originalCode.set(this.meta, code); + signaturePrefixes.set(this.meta, normalizedSignature.prefixes); + + return normalizedSignature.code.replace( + optionalModifierRegex, + "$1abstract" + ); + } + + if (looksLikeTypeDocMember(code)) { + options.grammarContextCode = "declare class __TypeDoc {\n"; + + originalCode.set(this.meta, code); + + return code.replace( + optionalModifierRegex, + "$1abstract" + ); + } + + return; + }, + + tokens(tokens) { + if (!enabled.has(this.meta)) + return; + + const links = qualifiedSymbolLinks.get(this.meta); + const prefixes = signaturePrefixes.get(this.meta); + + if (links == null && prefixes == null) + return; + + return tokens.map((line, lineIndex) => { + let column = 0; + + return line.flatMap((token) => { + const tokenColumn = column; + column += token.content.length; + + const signaturePrefix = prefixes?.get(lineIndex + 1)?.find((prefix) => ( + tokenColumn <= prefix.column && + tokenColumn + token.content.length >= prefix.column + prefix.content.length + )); + + if (signaturePrefix != null) { + const prefixOffset = signaturePrefix.column - tokenColumn; + const prefixEnd = prefixOffset + signaturePrefix.content.length; + const result: typeof line = []; + + if (prefixOffset > 0) + result.push({ + ...token, + content: token.content.slice(0, prefixOffset) + }); + + if (signaturePrefix.typeName != null) { + const typeStyleToken = line.find((otherToken) => ( + otherToken.content === signaturePrefix.typeName && + otherToken !== token && + isTypeToken(otherToken) + )) ?? line.find(isTypeToken); + + result.push({ + ...(typeStyleToken ?? token), + content: signaturePrefix.typeName, + offset: token.offset + prefixOffset + }); + + const remainingPrefix = signaturePrefix.content.slice(signaturePrefix.typeName.length); + if (remainingPrefix !== "") + result.push({ + ...token, + content: remainingPrefix, + offset: token.offset + prefixOffset + signaturePrefix.typeName.length + }); + } else + result.push({ + ...token, + content: signaturePrefix.content, + offset: token.offset + prefixOffset + }); + + if (prefixEnd < token.content.length) + result.push({ + ...token, + content: token.content.slice(prefixEnd), + offset: token.offset + prefixEnd + }); + + return result; + } + + if (!links?.has(`${lineIndex + 1}:${tokenColumn}`) || !token.content.endsWith(".")) + return token; + + return [{ + ...token, + content: token.content.slice(0, -".".length) + }, + { + ...token, + content: ".", + offset: token.offset + token.content.length - 1 + }]; + }); + }); + }, + + span(node, line, column, lineElement, token) { + if (!enabled.has(this.meta)) + return; + + const signaturePrefix = signaturePrefixes.get(this.meta)?.get(line) + ?.find((prefix) => ( + prefix.column === column && + prefix.typeName === token.content && + prefix.href != null + )); + + if (signaturePrefix != null) { + node.tagName = "a"; + node.properties.href = resolveHref(signaturePrefix.href!); + node.properties["class"] = "nlc-api-link" + ( + node.properties["class"] + ? (" " + node.properties["class"]) + : "" + ); + return; + } + + const original = originalCode.get(this.meta); + if (original != null && token.content === "abstract" && restoreOptionalModifier(original, line, column, node)) + return; + + const qualifiedLinks = qualifiedSymbolLinks.get(this.meta); + let href = qualifiedLinks?.get(`${line}:${column}`); + + const isType = token.explanation?.some((explanation) => ( + explanation.scopes.some(({scopeName}) => ( + scopeName.startsWith("entity.name.type.") || + scopeName === "entity.name.type.ts" || + scopeName === "entity.other.inherited-class.ts" + )) + )) ?? false; + + if (href == null && (isType || qualifiedLinks != null)) + href = apiSymbolLinks.get(token.content)?.link; + if (href == null) + return; + + const isTypeDeclaration = token.explanation?.some((explanation) => ( + explanation.scopes.some(({scopeName}) => ( + scopeName === "entity.name.type.alias.ts" || + scopeName === "entity.name.type.class.ts" || + scopeName === "entity.name.type.interface.ts" || + scopeName === "entity.name.type.enum.ts" + )) + )) ?? false; + + if (qualifiedLinks == null && isTypeDeclaration) + return; + + if (href.toLowerCase().endsWith(".md")) + href = href.slice(0, -".md".length); + + node.tagName = "a"; + node.properties.href = resolveHref(href); + node.properties["class"] = "nlc-api-link" + ( + node.properties["class"] + ? (" " + node.properties["class"]) + : "" + ); + } + }; +} + +function looksLikeTypeDocMember(code: string): boolean { + const trimmed = code.trim(); + + const modifiers = `(?:(?:${typeDocMemberModifiersRegex})[\\t ]+)*`; + const identifier = "[A-Za-z_$][\\w$]*"; + + return ( + new RegExp(`^${modifiers}(?:get|set)[\\t ]+${identifier}[\\t ]*\\(`).test(trimmed) || + new RegExp(`^${modifiers}${identifier}(?:<[^>]+>)?[\\t ]*\\(`).test(trimmed) || + new RegExp(`^${modifiers}${identifier}\\??[\\t ]*:`).test(trimmed) + ); +} + +function normalizeTypeDocSignature(code: string) { + const lines = code.split("\n"); + const prefixes = new Map(); + + let changed = false; + + for (const [lineIndex, line] of lines.entries()) { + const qualifiedMatch = /^([\t ]*)([A-Za-z_$][\w$]*)\.([A-Za-z_$][\w$]*)(?=\s*(?:<[^>\n]*>)?\s*\()/.exec(line); + if (qualifiedMatch != null) { + const indentation = qualifiedMatch[1] ?? ""; + const typeName = qualifiedMatch[2]; + + if (typeName != null) { + let href = apiSymbolLinks.get(typeName)?.link; + + if (href != null) { + if (href.toLowerCase().endsWith(".md")) + href = href.slice(0, -".md".length); + + const prefix = `${typeName}.`; + const column = indentation.length; + + prefixes.set(lineIndex + 1, [{ + column, + content: prefix, + typeName, + href + }]); + + lines[lineIndex] = indentation + " ".repeat(prefix.length) + line.slice(column + prefix.length); + + changed = true; + continue; + } + } + } + + const constructorMatch = /^([\t ]*)new([\t ]+)([A-Za-z_$][\w$]*)(?=\s*(?:<[^>\n]*>)?\s*\()/.exec(line); + if (constructorMatch != null) { + const indentation = constructorMatch[1] ?? ""; + const spacing = constructorMatch[2] ?? " "; + const typeName = constructorMatch[3]; + + if (typeName != null) { + let href = apiSymbolLinks.get(typeName)?.link; + + if (href != null) { + if (href.toLowerCase().endsWith(".md")) + href = href.slice(0, -".md".length); + + const column = indentation.length + "new".length + spacing.length; + + prefixes.set(lineIndex + 1, [{ + column, + content: typeName, + typeName, + href + }]); + + lines[lineIndex] = line.slice(0, column) + " ".repeat(typeName.length) + line.slice(column + typeName.length); + changed = true; + } + } + } + } + + if (!changed) + return undefined; + + return { + code: lines.join("\n"), + prefixes + }; +} + +function restoreOptionalModifier(original: string, line: number, column: number, node: Element) { + const originalLine = original.split("\n")[line - 1]; + if (originalLine == null) + return false; + + if (originalLine.slice(column, column + "optional".length) !== "optional") + return false; + + node.children = [{ + type: "text", + value: "optional" + }]; + return true; +} + +function getQualifiedSymbolLinks(code: string) { + const links = new Map(); + + for (const [lineIndex, line] of code.split("\n").entries()) { + for (const match of line.matchAll(/\b([A-Za-z_$][\w$]*)\.([A-Za-z_$][\w$]*)\b/g)) { + const typeName = match[1]; + const memberName = match[2]; + + if (typeName == null || memberName == null) + continue; + + let href = apiSymbolLinks.get(typeName)?.link; + if (href == null || match.index == null) + continue; + + if (href.toLowerCase().endsWith(".md")) + href = href.slice(0, -".md".length); + + links.set(`${lineIndex + 1}:${match.index}`, href); + links.set(`${lineIndex + 1}:${match.index + typeName.length + 1}`, `${href}#${memberName.toLowerCase()}`); + } + + for (const match of line.matchAll(/\bnew\s+([A-Za-z_$][\w$]*)\b/g)) { + const typeName = match[1]; + + if (typeName == null || match.index == null) + continue; + + let href = apiSymbolLinks.get(typeName)?.link; + if (href == null) + continue; + + if (href.toLowerCase().endsWith(".md")) + href = href.slice(0, -".md".length); + + const typeIndex = match.index + match[0].lastIndexOf(typeName); + + links.set(`${lineIndex + 1}:${typeIndex}`, href); + } + } + + return links; +} + +function typedocSidebarToSymbolMap() { + const map = new Map(); + + for (const category of typedocSidebar) { + for (const child of category.items) { + map.set(child.text, { + link: child.link, + section: category.text + }); + } + } + + return map; +} + +function isTypeToken(token: { + explanation?: { + scopes: { + scopeName: string + }[] + }[] +}) { + return token.explanation?.some((explanation) => ( + explanation.scopes.some(({scopeName}) => ( + scopeName.startsWith("entity.name.type.") || + scopeName === "entity.name.type.ts" || + scopeName === "entity.other.inherited-class.ts" + )) + )) ?? false; +} diff --git a/.vitepress/config/createCustomTwoSlashRenderer.ts b/.vitepress/config/createCustomTwoSlashRenderer.ts new file mode 100644 index 00000000..c48a08cf --- /dev/null +++ b/.vitepress/config/createCustomTwoSlashRenderer.ts @@ -0,0 +1,65 @@ +import type {ShikiTransformer, ShikiTransformerContext} from "shiki"; +import type {TwoslashRenderer} from "@shikijs/twoslash/core"; + +export function createCustomTwoSlashRenderer({ + baseRenderer, transformers +}: { + baseRenderer: TwoslashRenderer, + transformers: ShikiTransformer[] +}): TwoslashRenderer { + function createContext(context: ShikiTransformerContext): ShikiTransformerContext { + return { + ...context, + codeToHast( + code: Parameters[0], + options: Parameters[1] + ) { + return context.codeToHast(code, { + ...options, + lang: options.lang.toLowerCase() === "typescript" + ? "ts" + : options.lang, + transformers: [ + ...(options.transformers ?? []), + ...transformers + ] + }); + } + }; + } + + function wrap(fn: ((this: ShikiTransformerContext, ...args: Args) => Return)): ( + ((this: ShikiTransformerContext, ...args: Args) => Return) + ); + function wrap(fn: ((this: ShikiTransformerContext, ...args: Args) => Return) | undefined): ( + ((this: ShikiTransformerContext, ...args: Args) => Return) | undefined + ); + function wrap(method: ((this: ShikiTransformerContext, ...args: Args) => Return) | undefined) { + if (method == null) + return undefined; + + return function (this: ShikiTransformerContext, ...args: Args): Return { + return method.apply(createContext(this), args); + }; + } + + return Object.fromEntries( + Object.entries(baseRenderer) + .map(([key, value]) => [key, wrap(value)] as [typeof key, typeof value]) + ) as TwoslashRenderer; +} + +export function preserveTwoslashLanguage(transformer: ShikiTransformer): ShikiTransformer { + const preprocess = transformer.preprocess; + + return { + ...transformer, + preprocess(code, options) { + const originalLang = options.lang; + const result = preprocess?.call(this, code, options); + options.lang = originalLang; + + return result; + } + }; +} diff --git a/.vitepress/theme/index.ts b/.vitepress/theme/index.ts index 52f4288c..e2581f53 100644 --- a/.vitepress/theme/index.ts +++ b/.vitepress/theme/index.ts @@ -37,7 +37,19 @@ export default { }, enhanceApp({app, router, siteData}: EnhanceAppContext) { app.component("YouTubePlayer", YouTubePlayer); - app.use(TwoslashFloatingVue); + app.use(TwoslashFloatingVue, { + disposeTimeout: 300, + themes: { + twoslash: { + instantMove: false, + distance: 2, + delay: { + show: 50, + hide: 150 + } + } + } + }); app.use(NolebaseGitChangelogPlugin, { displayAuthorsInsideCommitLine: true, hideChangelogHeader: true, diff --git a/.vitepress/theme/style.css b/.vitepress/theme/style.css index 339f2fe4..b0cb5762 100644 --- a/.vitepress/theme/style.css +++ b/.vitepress/theme/style.css @@ -288,6 +288,36 @@ img.blog-coverImage { text-wrap: wrap; } +.twoslash-floating>.v-popper__wrapper { + transition: opacity 0.15s ease-out, transform 0.15s ease-out; + transform-origin: 50% 0%; + transform: scale(var(--hide-scale)) translateY(var(--hide-distance)); + --hide-scale: 0.96; + --hide-distance: -2px; + opacity: 0; +} +.twoslash-floating.v-popper__popper--shown>.v-popper__wrapper { + transition: opacity 0.15s ease-out, transform 0.15s ease-out; + transform-origin: 50% 0%; + transform: scale(1) translateY(0); + opacity: 1; +} +.twoslash-floating[data-popper-placement^="bottom-"] > .v-popper__wrapper { + transform-origin: 50% 0%; +} +.twoslash-floating[data-popper-placement^="top-"] > .v-popper__wrapper { + transform-origin: 50% 100%; +} +.twoslash-floating.v-popper__popper--hidden[data-popper-placement^="bottom-"]>.v-popper__wrapper { + transform: scale(var(--hide-scale)) translateY(var(--hide-distance)); +} +.twoslash-floating.v-popper__popper--hidden[data-popper-placement^="top-"]>.v-popper__wrapper { + transform: scale(var(--hide-scale)) translateY(calc(var(--hide-distance) * -1)); +} +.twoslash-floating .v-popper__arrow-container { + display: none; +} + span.twoslash-popup-docs-tag-value>code:has(>pre>code) { background-color: transparent; } @@ -297,6 +327,68 @@ span.twoslash-popup-docs-tag-value>code:has(>pre>code) { border-radius: 6px; } +.dark .vp-code a.nlc-api-link { + color: var(--shiki-dark, inherit); + --color: var(--shiki-dark, inherit); +} + +html:not(.dark) .vp-code a.nlc-api-link { + color: var(--shiki-light, inherit); + --color: var(--shiki-light, inherit); +} + +.dark .vp-code a.nlc-api-link:hover { + color: color-mix(in srgb, var(--color) 64%, white); +} +html:not(.dark) .vp-code a.nlc-api-link:hover { + color: color-mix(in srgb, var(--color) 64%, black); +} + +.vp-code a.nlc-api-link { + border-bottom: none; + text-decoration: none; + position: relative; + transition: color 0.2s ease-in-out, opacity 0.2s ease-in-out; +} +.vp-code a.nlc-api-link:after { + content: ""; + display: block; + position: absolute; + height: 2px; + inset-inline-end: 0px; + width: 100%; + z-index: -1; + pointer-events: none; + background-color: var(--color); + border-radius: 12px; + overflow: hidden; + bottom: -3px; + opacity: 0; + transition: opacity 0.3s ease-in-out, transform 0s 0.3s ease-in-out; + transform: scaleX(0.72); + transform-origin: 50% 50%; +} +.vp-code .twoslash-popup-code a.nlc-api-link:after { + transform: scaleX(0.86); + opacity: 0.16; + transition: opacity 0.3s ease-in-out, transform 0.3s ease-in-out; +} +.vp-code:hover a.nlc-api-link:after, +.vp-code:focus-visible a.nlc-api-link:after, +.vp-code:has(:focus-visible) a.nlc-api-link:after { + transform: scaleX(1); + opacity: 0.32; + transition: opacity 0.3s ease-in-out, transform 0.3s ease-in-out; +} +.vp-code a.nlc-api-link:focus-visible { + outline: color-mix(in srgb, var(--color) 48%, transparent) solid 2px; + outline-offset: 2px; + border-radius: 4px; +} +.vp-code a.nlc-api-link:focus-visible:after { + display: none; +} + .VPFeature { border-radius: 16px; } diff --git a/llama/addon/AddonModel.cpp b/llama/addon/AddonModel.cpp index bd447e11..356a30e3 100644 --- a/llama/addon/AddonModel.cpp +++ b/llama/addon/AddonModel.cpp @@ -331,16 +331,14 @@ AddonModel::AddonModel(const Napi::CallbackInfo& info) : model_params.vocab_only = options.Get("vocabOnly").As().Value(); } - if (options.Has("useMmap")) { - model_params.use_mmap = options.Get("useMmap").As().Value(); - } - - if (options.Has("useDirectIo")) { - model_params.use_direct_io = options.Get("useDirectIo").As().Value(); - } - - if (options.Has("useMlock")) { - model_params.use_mlock = options.Get("useMlock").As().Value(); + if (options.Has("useMlock") && options.Get("useMlock").As().Value()) { + model_params.load_mode = LLAMA_LOAD_MODE_MLOCK; + } else if (options.Has("useDirectIo") && options.Get("useDirectIo").As().Value()) { + model_params.load_mode = LLAMA_LOAD_MODE_DIRECT_IO; + } else if (options.Has("useMmap") && options.Get("useMmap").As().Value()) { + model_params.load_mode = LLAMA_LOAD_MODE_MMAP; + } else { + model_params.load_mode = LLAMA_LOAD_MODE_NONE; } if (options.Has("checkTensors")) { @@ -440,8 +438,7 @@ AddonModel::AddonModel(const Napi::CallbackInfo& info) : } if (model_params.no_alloc) { - model_params.use_mlock = false; - model_params.use_mmap = false; + model_params.load_mode = LLAMA_LOAD_MODE_NONE; } } diff --git a/llama/addon/AddonSampler.cpp b/llama/addon/AddonSampler.cpp index cfea8fa4..35828e23 100644 --- a/llama/addon/AddonSampler.cpp +++ b/llama/addon/AddonSampler.cpp @@ -492,6 +492,7 @@ Napi::Value AddonSampler::ApplyConfig(const Napi::CallbackInfo& info) { if (shouldCreateSampler) { repeatPenaltySampler = llama_sampler_init_penalties( + llama_vocab_n_tokens(model->vocab), repeatPenaltyMaxTokens, repeatPenalty, repeatPenaltyFrequencyPenalty, @@ -541,7 +542,7 @@ Napi::Value AddonSampler::ApplyConfig(const Napi::CallbackInfo& info) { sequenceBreakers.reserve(sequenceBreakersArray.Length()); for (size_t i = 0; i < sequenceBreakersArray.Length(); i++) { std::string breaker = sequenceBreakersArray.Get(i).As().Utf8Value(); - + if (sequenceBreaksIsTheSame && dryRepeatPenalty_sequenceBreakers[i] != breaker) { sequenceBreaksIsTheSame = false; } @@ -596,7 +597,6 @@ Napi::Value AddonSampler::ApplyConfig(const Napi::CallbackInfo& info) { dryRepeatPenaltySampler = llama_sampler_init_dry( model->vocab, - llama_model_n_ctx_train(model->model), strength, base, allowedLength, diff --git a/llama/addon/addon.cpp b/llama/addon/addon.cpp index 51347210..f8ce97af 100644 --- a/llama/addon/addon.cpp +++ b/llama/addon/addon.cpp @@ -1,6 +1,8 @@ #include +#include #include #include +#include #include "AddonContext.h" #include "AddonGgufMetadata.h" @@ -22,6 +24,20 @@ std::mutex backendMutex; bool backendInitialized = false; bool backendDisposed = false; +static bool compareWithUpperString(std::string_view source, std::string_view target) { + if (source.size() != target.size()) { + return false; + } + + for (std::size_t i = 0; i < source.size(); i++) { + if (static_cast(source[i]) != std::toupper(static_cast(target[i]))) { + return false; + } + } + + return true; +} + Napi::Value systemInfo(const Napi::CallbackInfo& info) { return Napi::String::From(info.Env(), llama_print_system_info()); } @@ -94,6 +110,45 @@ Napi::Value addonGetGgmlGraphOverheadCustom(const Napi::CallbackInfo& info) { return Napi::Number::New(info.Env(), graphOverhead); } +Napi::Value addonGetGgmlType(const Napi::CallbackInfo& info) { + if (info.Length() < 1) { + return info.Env().Undefined(); + } + + const auto typeParam = info[0]; + if (typeParam.IsNumber()) { + const auto typeParamValue = typeParam.As().Int32Value(); + if (typeParamValue < 0 || typeParamValue >= GGML_TYPE_COUNT) { + return info.Env().Undefined(); + } + + if (ggml_type_size(static_cast(typeParamValue)) == 0) { + return info.Env().Undefined(); + } + + return Napi::Number::New(info.Env(), typeParamValue); + } else if (typeParam.IsString()) { + const auto typeParamValue = typeParam.As().Utf8Value(); + + for (int i = 0; i < GGML_TYPE_COUNT; i++) { + if (ggml_type_size(static_cast(i)) == 0) { + continue; + } + + const auto typeName = ggml_type_name(static_cast(i)); + if (typeName == nullptr) { + continue; + } + + if (compareWithUpperString(typeParamValue, typeName)) { + return Napi::Number::New(info.Env(), i); + } + } + } + + return info.Env().Undefined(); +} + Napi::Value addonGetConsts(const Napi::CallbackInfo& info) { Napi::Object consts = Napi::Object::New(info.Env()); consts.Set("ggmlMaxDims", Napi::Number::New(info.Env(), GGML_MAX_DIMS)); @@ -301,6 +356,7 @@ Napi::Object registerCallback(Napi::Env env, Napi::Object exports) { Napi::PropertyDescriptor::Function("getBlockSizeForGgmlType", addonGetBlockSizeForGgmlType), Napi::PropertyDescriptor::Function("getTypeSizeForGgmlType", addonGetTypeSizeForGgmlType), Napi::PropertyDescriptor::Function("getGgmlGraphOverheadCustom", addonGetGgmlGraphOverheadCustom), + Napi::PropertyDescriptor::Function("getGgmlType", addonGetGgmlType), Napi::PropertyDescriptor::Function("getConsts", addonGetConsts), Napi::PropertyDescriptor::Function("setLogger", setLogger), Napi::PropertyDescriptor::Function("setLoggerLogLevel", setLoggerLogLevel), diff --git a/llama/gpuInfo/vulkan-gpu-info.cpp b/llama/gpuInfo/vulkan-gpu-info.cpp index 3c30b040..6ee9a4ab 100644 --- a/llama/gpuInfo/vulkan-gpu-info.cpp +++ b/llama/gpuInfo/vulkan-gpu-info.cpp @@ -1,6 +1,11 @@ #include +#include #include +#include +#include #include +#include +#include #include #include @@ -13,13 +18,83 @@ constexpr std::uint32_t VK_VENDOR_ID_QUALCOMM = 0x5143; typedef void (*gpuInfoVulkanWarningLogCallback_t)(const char* message); +static bool addWithoutOverflow(uint64_t& target, uint64_t value) { + if (value > std::numeric_limits::max() - target) { + return false; + } + + target += value; + return true; +} + static vk::Instance vulkanInstance() { - vk::ApplicationInfo appInfo("node-llama-cpp GPU info", 1, "llama.cpp", 1, VK_API_VERSION_1_2); - vk::InstanceCreateInfo createInfo(vk::InstanceCreateFlags(), &appInfo, {}, {}); - return vk::createInstance(createInfo); + static vk::Instance instance = []() { + const uint32_t apiVersion = vk::enumerateInstanceVersion(); + if (apiVersion < VK_API_VERSION_1_2) { + throw std::runtime_error("Vulkan 1.2 is not supported by the current system. Please update your Vulkan driver"); + } + + vk::ApplicationInfo appInfo("node-llama-cpp GPU info", 1, "llama.cpp", 1, VK_API_VERSION_1_2); + vk::InstanceCreateInfo createInfo(vk::InstanceCreateFlags(), &appInfo, {}, {}); + return vk::createInstance(createInfo); + }(); + + return instance; +} + +static bool deviceSupportsMemoryBudget(const vk::PhysicalDevice& physicalDevice) { + std::vector extensionProperties = physicalDevice.enumerateDeviceExtensionProperties(); + + return std::any_of( + extensionProperties.begin(), + extensionProperties.end(), + [](const vk::ExtensionProperties& ext) { + return std::string(ext.extensionName.data()) == VK_EXT_MEMORY_BUDGET_EXTENSION_NAME; + } + ); +} + +static bool isVulkanDeviceSupported(const vk::PhysicalDevice& physicalDevice, std::string* unsupportedReason = nullptr) { + if (unsupportedReason != nullptr) { + unsupportedReason->clear(); + } + + VkPhysicalDeviceFeatures2 features2 = {}; + features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; + + VkPhysicalDeviceVulkan11Features vk11Features = {}; + vk11Features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; + features2.pNext = &vk11Features; + + vkGetPhysicalDeviceFeatures2(physicalDevice, &features2); + + if (!vk11Features.storageBuffer16BitAccess) { + if (unsupportedReason != nullptr) { + vk::PhysicalDeviceProperties deviceProps = physicalDevice.getProperties(); + *unsupportedReason = + "Vulkan storageBuffer16BitAccess not supported for device \"" + + std::string(deviceProps.deviceName.data()) + "\""; + } + + return false; + } + + if (!deviceSupportsMemoryBudget(physicalDevice)) { + // VK_EXT_memory_budget extension is not supported, so we cannot determine used memory + + if (unsupportedReason != nullptr) { + vk::PhysicalDeviceProperties deviceProps = physicalDevice.getProperties(); + *unsupportedReason = "Vulkan VK_EXT_memory_budget extension not supported for device \"" + + std::string(deviceProps.deviceName.data()) + "\", so VRAM info cannot be determined for it"; + } + + return false; + } + + return true; } -static std::vector dedupedDevices() { +static std::vector dedupedDevices(gpuInfoVulkanWarningLogCallback_t warningLogCallback = nullptr) { vk::Instance instance = vulkanInstance(); auto physicalDevices = instance.enumeratePhysicalDevices(); std::vector dedupedDevices; @@ -27,6 +102,22 @@ static std::vector dedupedDevices() { // adapted from `ggml_vk_instance_init` in `ggml-vulkan.cpp` for (const auto& device : physicalDevices) { + vk::PhysicalDeviceProperties deviceProps = device.getProperties(); + + // ignore CPU devices, as we don't want to count RAM from the CPU as VRAM + if (deviceProps.deviceType == vk::PhysicalDeviceType::eCpu) { + continue; + } + + std::string unsupportedReason; + if (!isVulkanDeviceSupported(device, warningLogCallback != nullptr ? &unsupportedReason : nullptr)) { + if (warningLogCallback != nullptr) { + warningLogCallback(unsupportedReason.c_str()); + } + + continue; + } + vk::PhysicalDeviceProperties2 newProps; vk::PhysicalDeviceDriverProperties newDriver; vk::PhysicalDeviceIDProperties newId; @@ -109,78 +200,58 @@ static std::vector dedupedDevices() { return dedupedDevices; } -static bool enumerateVulkanDevices(size_t* total, size_t* used, size_t* unifiedMemorySize, bool addDeviceNames, std::vector * deviceNames, gpuInfoVulkanWarningLogCallback_t warningLogCallback, bool * checkSupported) { - auto physicalDevices = dedupedDevices(); +static bool enumerateVulkanDevices(uint64_t* total, uint64_t* used, uint64_t* unifiedMemorySize, bool addDeviceNames, std::vector * deviceNames, gpuInfoVulkanWarningLogCallback_t warningLogCallback) { + auto physicalDevices = dedupedDevices(warningLogCallback); - size_t usedMem = 0; - size_t totalMem = 0; - size_t totalUnifiedMemorySize = 0; + uint64_t usedMem = 0; + uint64_t totalMem = 0; + uint64_t totalUnifiedMemorySize = 0; for (size_t i = 0; i < physicalDevices.size(); i++) { vk::PhysicalDevice physicalDevice = physicalDevices[i]; - vk::PhysicalDeviceMemoryProperties memProps = physicalDevice.getMemoryProperties(); vk::PhysicalDeviceProperties deviceProps = physicalDevice.getProperties(); - if (deviceProps.deviceType == vk::PhysicalDeviceType::eCpu) { - // ignore CPU devices, as we don't want to count RAM from the CPU as VRAM - continue; - } - - std::vector extensionProperties = physicalDevice.enumerateDeviceExtensionProperties(); - bool memoryBudgetExtensionSupported = - std::any_of( - extensionProperties.begin(), - extensionProperties.end(), - [](const vk::ExtensionProperties& ext) { return std::string(ext.extensionName.data()) == VK_EXT_MEMORY_BUDGET_EXTENSION_NAME;} - ); + vk::PhysicalDeviceMemoryBudgetPropertiesEXT memoryBudgetProperties; + vk::PhysicalDeviceMemoryProperties2 memProps2 = {}; + memProps2.pNext = &memoryBudgetProperties; - if (memoryBudgetExtensionSupported) { - vk::PhysicalDeviceMemoryBudgetPropertiesEXT memoryBudgetProperties; - vk::PhysicalDeviceMemoryProperties2 memProps2 = {}; - memProps2.pNext = &memoryBudgetProperties; + physicalDevice.getMemoryProperties2(&memProps2); - physicalDevice.getMemoryProperties2(&memProps2); + bool hasDeviceLocalHeap = false; - for (uint32_t i = 0; i < memProps.memoryHeapCount; ++i) { - const auto heap = memProps2.memoryProperties.memoryHeaps[i]; + for (uint32_t i = 0; i < memProps2.memoryProperties.memoryHeapCount; ++i) { + const auto heap = memProps2.memoryProperties.memoryHeaps[i]; - if (heap.flags & vk::MemoryHeapFlagBits::eDeviceLocal) { - totalMem += heap.size; - usedMem += memoryBudgetProperties.heapUsage[i] + (heap.size - memoryBudgetProperties.heapBudget[i]); + if (heap.flags & vk::MemoryHeapFlagBits::eDeviceLocal) { + const uint64_t heapSize = heap.size; + const uint64_t heapBudget = std::min(memoryBudgetProperties.heapBudget[i], heapSize); + const uint64_t heapUsage = std::min(memoryBudgetProperties.heapUsage[i], heapBudget); + const uint64_t heapUsed = heapSize - (heapBudget - heapUsage); - if (heap.flags & vk::MemoryHeapFlagBits::eMultiInstance) { - totalUnifiedMemorySize += heap.size; - } + hasDeviceLocalHeap = heapSize != 0; - if (heap.size > 0 && addDeviceNames) { - (*deviceNames).push_back(std::string(deviceProps.deviceName.data())); + if (!addWithoutOverflow(totalMem, heapSize) || !addWithoutOverflow(usedMem, heapUsed)) { + if (warningLogCallback != nullptr) { + warningLogCallback("Vulkan VRAM size overflow"); } - if (checkSupported != nullptr && checkSupported) { - VkPhysicalDeviceFeatures2 features2 = {}; - features2.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2; - - VkPhysicalDeviceVulkan11Features vk11Features = {}; - vk11Features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES; - features2.pNext = &vk11Features; - - vkGetPhysicalDeviceFeatures2(physicalDevice, &features2); + return false; + } - if (!vk11Features.storageBuffer16BitAccess) { - *checkSupported = false; + if (deviceProps.deviceType == vk::PhysicalDeviceType::eIntegratedGpu) { + if (!addWithoutOverflow(totalUnifiedMemorySize, heapSize)) { + if (warningLogCallback != nullptr) { + warningLogCallback("Vulkan unified VRAM size overflow"); } + + return false; } } } - } else { - // VK_EXT_memory_budget extension is not supported, so we cannot determine used memory - warningLogCallback( - ( - "Vulkan VK_EXT_memory_budget extension not supported for device \"" + - std::string(deviceProps.deviceName.data()) + "\", so VRAM info cannot be determined for it" - ).c_str() - ); - return false; + } + + if (hasDeviceLocalHeap && addDeviceNames) { + (*deviceNames).push_back(std::string(deviceProps.deviceName.data())); } } @@ -191,17 +262,29 @@ static bool enumerateVulkanDevices(size_t* total, size_t* used, size_t* unifiedM return true; } -bool gpuInfoGetTotalVulkanDevicesInfo(size_t* total, size_t* used, size_t* unifiedMemorySize, gpuInfoVulkanWarningLogCallback_t warningLogCallback) { - return enumerateVulkanDevices(total, used, unifiedMemorySize, false, nullptr, warningLogCallback, nullptr); +bool gpuInfoGetTotalVulkanDevicesInfo(uint64_t* total, uint64_t* used, uint64_t* unifiedMemorySize, gpuInfoVulkanWarningLogCallback_t warningLogCallback) { + try { + return enumerateVulkanDevices(total, used, unifiedMemorySize, false, nullptr, warningLogCallback); + } catch (const std::exception& err) { + if (warningLogCallback != nullptr) { + std::string message = "Failed to get Vulkan GPU info: " + std::string(err.what()); + warningLogCallback(message.c_str()); + } + + return false; + } } bool checkIsVulkanEnvSupported(gpuInfoVulkanWarningLogCallback_t warningLogCallback) { - size_t total = 0; - size_t used = 0; - size_t unifiedMemorySize = 0; - - bool isSupported = true; - enumerateVulkanDevices(&total, &used, &unifiedMemorySize, false, nullptr, warningLogCallback, &isSupported); + try { + static_cast(vulkanInstance().enumeratePhysicalDevices()); + return true; + } catch (const std::exception& err) { + if (warningLogCallback != nullptr) { + std::string message = "Failed to check Vulkan support: " + std::string(err.what()); + warningLogCallback(message.c_str()); + } - return isSupported; -} + return false; + } +} \ No newline at end of file diff --git a/llama/gpuInfo/vulkan-gpu-info.h b/llama/gpuInfo/vulkan-gpu-info.h index 09f63406..e4be4943 100644 --- a/llama/gpuInfo/vulkan-gpu-info.h +++ b/llama/gpuInfo/vulkan-gpu-info.h @@ -5,5 +5,5 @@ typedef void (*gpuInfoVulkanWarningLogCallback_t)(const char* message); -bool gpuInfoGetTotalVulkanDevicesInfo(size_t* total, size_t* used, size_t* unifiedMemorySize, gpuInfoVulkanWarningLogCallback_t warningLogCallback); +bool gpuInfoGetTotalVulkanDevicesInfo(uint64_t* total, uint64_t* used, uint64_t* unifiedMemorySize, gpuInfoVulkanWarningLogCallback_t warningLogCallback); bool checkIsVulkanEnvSupported(gpuInfoVulkanWarningLogCallback_t warningLogCallback); diff --git a/llama/patches/PR-22566.diff b/llama/patches/PR-22566.diff index 37e9ba27..126a4532 100644 --- a/llama/patches/PR-22566.diff +++ b/llama/patches/PR-22566.diff @@ -1,11 +1,11 @@ diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp -index 474cabdfc095..76692620408d 100644 +index 71bc9f7ef0aa..9e91dd492a11 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp -@@ -698,8 +698,13 @@ llama_model_loader::llama_model_loader( +@@ -694,8 +694,13 @@ llama_model_loader::llama_model_loader( llm_kv = LLM_KV(llm_arch_from_string(arch_name)); } - + - n_kv = gguf_get_n_kv(metadata); - n_tensors = weights_map.size(); + n_kv = gguf_get_n_kv(metadata); @@ -15,12 +15,12 @@ index 474cabdfc095..76692620408d 100644 + n_tensors = weights_map.size(); + GGML_ASSERT(files.size() != 1 || static_cast(n_tensors) == gguf_get_n_tensors(metadata)); + } - + fver = (enum llama_fver) gguf_get_version(metadata); - -@@ -1213,13 +1218,20 @@ struct ggml_tensor * llama_model_loader::create_tensor( + +@@ -1225,13 +1230,20 @@ struct ggml_tensor * llama_model_loader::create_tensor( }; - + if (files.empty()) { - if (flags & TENSOR_SKIP_IF_VIRTUAL) { - return nullptr; @@ -40,24 +40,11 @@ index 474cabdfc095..76692620408d 100644 + throw std::runtime_error(format("missing tensor '%s'", tn.str().c_str())); + } } - + // for tensors that are not required some of the dimensions can be invalid: -@@ -1237,14 +1249,32 @@ struct ggml_tensor * llama_model_loader::create_tensor( - for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) { - t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1; - GGML_ASSERT(t_meta.ne[dim] >= 1); -- t_meta.nb[dim] = dim == 0 ? ggml_type_size(type) : t_meta.ne[dim-1]*t_meta.nb[dim-1]; -+ } -+ t_meta.nb[0] = ggml_type_size(type); -+ t_meta.nb[1] = t_meta.nb[0] * (t_meta.ne[0] / ggml_blck_size(type)); -+ GGML_ASSERT(t_meta.nb[0] >= 1); -+ GGML_ASSERT(t_meta.nb[1] >= 1); -+ for (size_t dim = 2; dim < GGML_MAX_DIMS; ++dim) { -+ t_meta.nb[dim] = t_meta.nb[dim - 1] * t_meta.ne[dim - 1]; - GGML_ASSERT(t_meta.nb[dim] >= 1); - } +@@ -1261,8 +1273,20 @@ struct ggml_tensor * llama_model_loader::create_tensor( ggml_set_name(&t_meta, tn.str().c_str()); - + ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta); - GGML_ASSERT(buft != nullptr); + if (buft == nullptr) { diff --git a/package-lock.json b/package-lock.json index a1c02768..9e7338e3 100644 --- a/package-lock.json +++ b/package-lock.json @@ -23,7 +23,7 @@ "ignore": "^7.0.4", "ipull": "^3.9.5", "is-unicode-supported": "^2.1.0", - "lifecycle-utils": "^3.1.1", + "lifecycle-utils": "^4.0.1", "log-symbols": "^7.0.1", "nanoid": "^5.1.6", "node-addon-api": "^8.6.0", @@ -100,22 +100,6 @@ "type": "github", "url": "https://github.com/sponsors/giladgd" }, - "optionalDependencies": { - "@node-llama-cpp/linux-arm64": "0.1.0", - "@node-llama-cpp/linux-armv7l": "0.1.0", - "@node-llama-cpp/linux-riscv64": "0.1.0", - "@node-llama-cpp/linux-x64": "0.1.0", - "@node-llama-cpp/linux-x64-cuda": "0.1.0", - "@node-llama-cpp/linux-x64-cuda-ext": "0.1.0", - "@node-llama-cpp/linux-x64-vulkan": "0.1.0", - "@node-llama-cpp/mac-arm64-metal": "0.1.0", - "@node-llama-cpp/mac-x64": "0.1.0", - "@node-llama-cpp/win-arm64": "0.1.0", - "@node-llama-cpp/win-x64": "0.1.0", - "@node-llama-cpp/win-x64-cuda": "0.1.0", - "@node-llama-cpp/win-x64-cuda-ext": "0.1.0", - "@node-llama-cpp/win-x64-vulkan": "0.1.0" - }, "peerDependencies": { "typescript": ">=5.0.0" }, @@ -987,422 +971,6 @@ "node": ">=10" } }, - "node_modules/@esbuild/aix-ppc64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/aix-ppc64/-/aix-ppc64-0.27.4.tgz", - "integrity": "sha512-cQPwL2mp2nSmHHJlCyoXgHGhbEPMrEEU5xhkcy3Hs/O7nGZqEpZ2sUtLaL9MORLtDfRvVl2/3PAuEkYZH0Ty8Q==", - "cpu": [ - "ppc64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "aix" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/android-arm": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/android-arm/-/android-arm-0.27.4.tgz", - "integrity": "sha512-X9bUgvxiC8CHAGKYufLIHGXPJWnr0OCdR0anD2e21vdvgCI8lIfqFbnoeOz7lBjdrAGUhqLZLcQo6MLhTO2DKQ==", - "cpu": [ - "arm" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "android" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/android-arm64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/android-arm64/-/android-arm64-0.27.4.tgz", - "integrity": "sha512-gdLscB7v75wRfu7QSm/zg6Rx29VLdy9eTr2t44sfTW7CxwAtQghZ4ZnqHk3/ogz7xao0QAgrkradbBzcqFPasw==", - "cpu": [ - "arm64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "android" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/android-x64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/android-x64/-/android-x64-0.27.4.tgz", - "integrity": "sha512-PzPFnBNVF292sfpfhiyiXCGSn9HZg5BcAz+ivBuSsl6Rk4ga1oEXAamhOXRFyMcjwr2DVtm40G65N3GLeH1Lvw==", - "cpu": [ - "x64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "android" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/darwin-arm64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/darwin-arm64/-/darwin-arm64-0.27.4.tgz", - "integrity": "sha512-b7xaGIwdJlht8ZFCvMkpDN6uiSmnxxK56N2GDTMYPr2/gzvfdQN8rTfBsvVKmIVY/X7EM+/hJKEIbbHs9oA4tQ==", - "cpu": [ - "arm64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "darwin" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/darwin-x64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/darwin-x64/-/darwin-x64-0.27.4.tgz", - "integrity": "sha512-sR+OiKLwd15nmCdqpXMnuJ9W2kpy0KigzqScqHI3Hqwr7IXxBp3Yva+yJwoqh7rE8V77tdoheRYataNKL4QrPw==", - "cpu": [ - "x64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "darwin" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/freebsd-arm64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/freebsd-arm64/-/freebsd-arm64-0.27.4.tgz", - "integrity": "sha512-jnfpKe+p79tCnm4GVav68A7tUFeKQwQyLgESwEAUzyxk/TJr4QdGog9sqWNcUbr/bZt/O/HXouspuQDd9JxFSw==", - "cpu": [ - "arm64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "freebsd" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/freebsd-x64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/freebsd-x64/-/freebsd-x64-0.27.4.tgz", - "integrity": "sha512-2kb4ceA/CpfUrIcTUl1wrP/9ad9Atrp5J94Lq69w7UwOMolPIGrfLSvAKJp0RTvkPPyn6CIWrNy13kyLikZRZQ==", - "cpu": [ - "x64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "freebsd" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-arm": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-arm/-/linux-arm-0.27.4.tgz", - "integrity": "sha512-aBYgcIxX/wd5n2ys0yESGeYMGF+pv6g0DhZr3G1ZG4jMfruU9Tl1i2Z+Wnj9/KjGz1lTLCcorqE2viePZqj4Eg==", - "cpu": [ - "arm" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-arm64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-arm64/-/linux-arm64-0.27.4.tgz", - "integrity": "sha512-7nQOttdzVGth1iz57kxg9uCz57dxQLHWxopL6mYuYthohPKEK0vU0C3O21CcBK6KDlkYVcnDXY099HcCDXd9dA==", - "cpu": [ - "arm64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-ia32": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-ia32/-/linux-ia32-0.27.4.tgz", - "integrity": "sha512-oPtixtAIzgvzYcKBQM/qZ3R+9TEUd1aNJQu0HhGyqtx6oS7qTpvjheIWBbes4+qu1bNlo2V4cbkISr8q6gRBFA==", - "cpu": [ - "ia32" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-loong64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-loong64/-/linux-loong64-0.27.4.tgz", - "integrity": "sha512-8mL/vh8qeCoRcFH2nM8wm5uJP+ZcVYGGayMavi8GmRJjuI3g1v6Z7Ni0JJKAJW+m0EtUuARb6Lmp4hMjzCBWzA==", - "cpu": [ - "loong64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-mips64el": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-mips64el/-/linux-mips64el-0.27.4.tgz", - "integrity": "sha512-1RdrWFFiiLIW7LQq9Q2NES+HiD4NyT8Itj9AUeCl0IVCA459WnPhREKgwrpaIfTOe+/2rdntisegiPWn/r/aAw==", - "cpu": [ - "mips64el" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-ppc64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-ppc64/-/linux-ppc64-0.27.4.tgz", - "integrity": "sha512-tLCwNG47l3sd9lpfyx9LAGEGItCUeRCWeAx6x2Jmbav65nAwoPXfewtAdtbtit/pJFLUWOhpv0FpS6GQAmPrHA==", - "cpu": [ - "ppc64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-riscv64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-riscv64/-/linux-riscv64-0.27.4.tgz", - "integrity": "sha512-BnASypppbUWyqjd1KIpU4AUBiIhVr6YlHx/cnPgqEkNoVOhHg+YiSVxM1RLfiy4t9cAulbRGTNCKOcqHrEQLIw==", - "cpu": [ - "riscv64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-s390x": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-s390x/-/linux-s390x-0.27.4.tgz", - "integrity": "sha512-+eUqgb/Z7vxVLezG8bVB9SfBie89gMueS+I0xYh2tJdw3vqA/0ImZJ2ROeWwVJN59ihBeZ7Tu92dF/5dy5FttA==", - "cpu": [ - "s390x" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/linux-x64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/linux-x64/-/linux-x64-0.27.4.tgz", - "integrity": "sha512-S5qOXrKV8BQEzJPVxAwnryi2+Iq5pB40gTEIT69BQONqR7JH1EPIcQ/Uiv9mCnn05jff9umq/5nqzxlqTOg9NA==", - "cpu": [ - "x64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "linux" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/netbsd-arm64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/netbsd-arm64/-/netbsd-arm64-0.27.4.tgz", - "integrity": "sha512-xHT8X4sb0GS8qTqiwzHqpY00C95DPAq7nAwX35Ie/s+LO9830hrMd3oX0ZMKLvy7vsonee73x0lmcdOVXFzd6Q==", - "cpu": [ - "arm64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "netbsd" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/netbsd-x64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/netbsd-x64/-/netbsd-x64-0.27.4.tgz", - "integrity": "sha512-RugOvOdXfdyi5Tyv40kgQnI0byv66BFgAqjdgtAKqHoZTbTF2QqfQrFwa7cHEORJf6X2ht+l9ABLMP0dnKYsgg==", - "cpu": [ - "x64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "netbsd" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/openbsd-arm64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/openbsd-arm64/-/openbsd-arm64-0.27.4.tgz", - "integrity": "sha512-2MyL3IAaTX+1/qP0O1SwskwcwCoOI4kV2IBX1xYnDDqthmq5ArrW94qSIKCAuRraMgPOmG0RDTA74mzYNQA9ow==", - "cpu": [ - "arm64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "openbsd" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/openbsd-x64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/openbsd-x64/-/openbsd-x64-0.27.4.tgz", - "integrity": "sha512-u8fg/jQ5aQDfsnIV6+KwLOf1CmJnfu1ShpwqdwC0uA7ZPwFws55Ngc12vBdeUdnuWoQYx/SOQLGDcdlfXhYmXQ==", - "cpu": [ - "x64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "openbsd" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/openharmony-arm64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/openharmony-arm64/-/openharmony-arm64-0.27.4.tgz", - "integrity": "sha512-JkTZrl6VbyO8lDQO3yv26nNr2RM2yZzNrNHEsj9bm6dOwwu9OYN28CjzZkH57bh4w0I2F7IodpQvUAEd1mbWXg==", - "cpu": [ - "arm64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "openharmony" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/sunos-x64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/sunos-x64/-/sunos-x64-0.27.4.tgz", - "integrity": "sha512-/gOzgaewZJfeJTlsWhvUEmUG4tWEY2Spp5M20INYRg2ZKl9QPO3QEEgPeRtLjEWSW8FilRNacPOg8R1uaYkA6g==", - "cpu": [ - "x64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "sunos" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/win32-arm64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/win32-arm64/-/win32-arm64-0.27.4.tgz", - "integrity": "sha512-Z9SExBg2y32smoDQdf1HRwHRt6vAHLXcxD2uGgO/v2jK7Y718Ix4ndsbNMU/+1Qiem9OiOdaqitioZwxivhXYg==", - "cpu": [ - "arm64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "win32" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/win32-ia32": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/win32-ia32/-/win32-ia32-0.27.4.tgz", - "integrity": "sha512-DAyGLS0Jz5G5iixEbMHi5KdiApqHBWMGzTtMiJ72ZOLhbu/bzxgAe8Ue8CTS3n3HbIUHQz/L51yMdGMeoxXNJw==", - "cpu": [ - "ia32" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "win32" - ], - "engines": { - "node": ">=18" - } - }, - "node_modules/@esbuild/win32-x64": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/@esbuild/win32-x64/-/win32-x64-0.27.4.tgz", - "integrity": "sha512-+knoa0BDoeXgkNvvV1vvbZX4+hizelrkwmGJBdT17t8FNPwG2lKemmuMZlmaNQ3ws3DKKCxpb4zRZEIp3UxFCg==", - "cpu": [ - "x64" - ], - "extraneous": true, - "license": "MIT", - "os": [ - "win32" - ], - "engines": { - "node": ">=18" - } - }, "node_modules/@eslint-community/eslint-utils": { "version": "4.9.1", "resolved": "https://registry.npmjs.org/@eslint-community/eslint-utils/-/eslint-utils-4.9.1.tgz", @@ -2373,48 +1941,6 @@ "@tybys/wasm-util": "^0.10.0" } }, - "node_modules/@node-llama-cpp/linux-arm64": { - "optional": true - }, - "node_modules/@node-llama-cpp/linux-armv7l": { - "optional": true - }, - "node_modules/@node-llama-cpp/linux-riscv64": { - "optional": true - }, - "node_modules/@node-llama-cpp/linux-x64": { - "optional": true - }, - "node_modules/@node-llama-cpp/linux-x64-cuda": { - "optional": true - }, - "node_modules/@node-llama-cpp/linux-x64-cuda-ext": { - "optional": true - }, - "node_modules/@node-llama-cpp/linux-x64-vulkan": { - "optional": true - }, - "node_modules/@node-llama-cpp/mac-arm64-metal": { - "optional": true - }, - "node_modules/@node-llama-cpp/mac-x64": { - "optional": true - }, - "node_modules/@node-llama-cpp/win-arm64": { - "optional": true - }, - "node_modules/@node-llama-cpp/win-x64": { - "optional": true - }, - "node_modules/@node-llama-cpp/win-x64-cuda": { - "optional": true - }, - "node_modules/@node-llama-cpp/win-x64-cuda-ext": { - "optional": true - }, - "node_modules/@node-llama-cpp/win-x64-vulkan": { - "optional": true - }, "node_modules/@nodelib/fs.scandir": { "version": "2.1.5", "resolved": "https://registry.npmjs.org/@nodelib/fs.scandir/-/fs.scandir-2.1.5.tgz", @@ -8164,48 +7690,6 @@ "license": "MIT", "optional": true }, - "node_modules/esbuild": { - "version": "0.27.4", - "resolved": "https://registry.npmjs.org/esbuild/-/esbuild-0.27.4.tgz", - "integrity": "sha512-Rq4vbHnYkK5fws5NF7MYTU68FPRE1ajX7heQ/8QXXWqNgqqJ/GkmmyxIzUnf2Sr/bakf8l54716CcMGHYhMrrQ==", - "extraneous": true, - "hasInstallScript": true, - "license": "MIT", - "bin": { - "esbuild": "bin/esbuild" - }, - "engines": { - "node": ">=18" - }, - "optionalDependencies": { - "@esbuild/aix-ppc64": "0.27.4", - "@esbuild/android-arm": "0.27.4", - "@esbuild/android-arm64": "0.27.4", - "@esbuild/android-x64": "0.27.4", - "@esbuild/darwin-arm64": "0.27.4", - "@esbuild/darwin-x64": "0.27.4", - "@esbuild/freebsd-arm64": "0.27.4", - "@esbuild/freebsd-x64": "0.27.4", - "@esbuild/linux-arm": "0.27.4", - "@esbuild/linux-arm64": "0.27.4", - "@esbuild/linux-ia32": "0.27.4", - "@esbuild/linux-loong64": "0.27.4", - "@esbuild/linux-mips64el": "0.27.4", - "@esbuild/linux-ppc64": "0.27.4", - "@esbuild/linux-riscv64": "0.27.4", - "@esbuild/linux-s390x": "0.27.4", - "@esbuild/linux-x64": "0.27.4", - "@esbuild/netbsd-arm64": "0.27.4", - "@esbuild/netbsd-x64": "0.27.4", - "@esbuild/openbsd-arm64": "0.27.4", - "@esbuild/openbsd-x64": "0.27.4", - "@esbuild/openharmony-arm64": "0.27.4", - "@esbuild/sunos-x64": "0.27.4", - "@esbuild/win32-arm64": "0.27.4", - "@esbuild/win32-ia32": "0.27.4", - "@esbuild/win32-x64": "0.27.4" - } - }, "node_modules/escalade": { "version": "3.2.0", "resolved": "https://registry.npmjs.org/escalade/-/escalade-3.2.0.tgz", @@ -11301,10 +10785,14 @@ } }, "node_modules/lifecycle-utils": { - "version": "3.1.1", - "resolved": "https://registry.npmjs.org/lifecycle-utils/-/lifecycle-utils-3.1.1.tgz", - "integrity": "sha512-gNd3OvhFNjHykJE3uGntz7UuPzWlK9phrIdXxU9Adis0+ExkwnZibfxCJWiWWZ+a6VbKiZrb+9D9hCQWd4vjTg==", - "license": "MIT" + "version": "4.0.1", + "resolved": "https://registry.npmjs.org/lifecycle-utils/-/lifecycle-utils-4.0.1.tgz", + "integrity": "sha512-u4FZ8oW6M1dWj9sJ6e1qNya3p0p6oMGAohjk8FZBU/blBhvRbl7bUdZL5m9cYRhopx/5OSDzLVzxKcoMiWJ34w==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/giladgd" + } }, "node_modules/lightningcss": { "version": "1.32.0", diff --git a/package.json b/package.json index 777b4720..c148af0a 100644 --- a/package.json +++ b/package.json @@ -68,7 +68,7 @@ "clean": "rm -rf ./node_modules ./dist ./tsconfig.tsbuildinfo ./test/.models ./docs/api ./docs/api-overrides ./templates/packed", "docs:generateTypedoc": "typedoc && rimraf ./docs/api/index.md ./docs/api/globals.md ./docs/api/functions/LlamaText.md", "docs:dev": "npm run docs:generateTypedoc && vitepress dev --cors", - "docs:build": "npm run docs:generateTypedoc && vitepress build", + "docs:build": "npm run docs:generateTypedoc && NODE_OPTIONS=\"--max-old-space-size=8192\" vitepress build", "docs:preview": "npm run docs:generateTypedoc && vitepress preview" }, "repository": { @@ -199,7 +199,7 @@ "ignore": "^7.0.4", "ipull": "^3.9.5", "is-unicode-supported": "^2.1.0", - "lifecycle-utils": "^3.1.1", + "lifecycle-utils": "^4.0.1", "log-symbols": "^7.0.1", "nanoid": "^5.1.6", "node-addon-api": "^8.6.0", @@ -223,7 +223,7 @@ "optional": true } }, - "optionalDependencies": { + "_optionalDependencies": { "@node-llama-cpp/linux-arm64": "0.1.0", "@node-llama-cpp/linux-armv7l": "0.1.0", "@node-llama-cpp/linux-riscv64": "0.1.0", diff --git a/scripts/postVersion.ts b/scripts/postVersion.ts index a78f5d25..1769f812 100644 --- a/scripts/postVersion.ts +++ b/scripts/postVersion.ts @@ -8,6 +8,14 @@ const packageJsonPath = path.join(__dirname, "..", "package.json"); const packageJson = await fs.readJson(packageJsonPath); const currentVersion = packageJson.version; +if (packageJson._optionalDependencies != null) { + packageJson.optionalDependencies = { + ...(packageJson.optionalDependencies ?? {}), + ...packageJson._optionalDependencies + }; + delete packageJson._optionalDependencies; +} + if (packageJson.optionalDependencies != null) { for (const packageName of Object.keys(packageJson.optionalDependencies)) { if (!packageName.startsWith("@node-llama-cpp/")) diff --git a/src/bindings/AddonTypes.ts b/src/bindings/AddonTypes.ts index 63a879ed..9c4911be 100644 --- a/src/bindings/AddonTypes.ts +++ b/src/bindings/AddonTypes.ts @@ -72,6 +72,7 @@ export type BindingModule = { getBlockSizeForGgmlType(ggmlType: number): number | undefined, getTypeSizeForGgmlType(ggmlType: number): number | undefined, getGgmlGraphOverheadCustom(size: number, grads: boolean): number, + getGgmlType(ggmlType: string | number): number | undefined, getConsts(): { ggmlMaxDims: number, ggmlTypeF16Size: number, diff --git a/src/bindings/utils/MemoryOrchestrator.ts b/src/bindings/utils/MemoryOrchestrator.ts index f2b870dd..3a7106b7 100644 --- a/src/bindings/utils/MemoryOrchestrator.ts +++ b/src/bindings/utils/MemoryOrchestrator.ts @@ -7,6 +7,7 @@ export class MemoryOrchestrator { /** @internal */ public _markedMemory: number = 0; /** @internal */ private _memoryCap: number | null = null; /** @internal */ private _padding: number = 0; + /** @internal */ public _markingsFinalizationRegistry: FinalizationRegistry; public readonly onMemoryReservationRelease = new EventRelay(); public readonly onMemoryMarkingRelease = new EventRelay(); @@ -15,6 +16,8 @@ export class MemoryOrchestrator { this._getMemoryState = getMemoryState; this._onMarkFinalized = this._onMarkFinalized.bind(this); + + this._markingsFinalizationRegistry = new FinalizationRegistry(this._onMarkFinalized); } public reserveMemory(bytes: number) { @@ -77,7 +80,7 @@ export class MemoryOrchestrator { public _onMarkFinalized(bytes: number) { this._markedMemory -= bytes; this.onMemoryMarkingRelease.dispatchEvent(); - } + } } export class MemoryReservation { @@ -116,13 +119,11 @@ export class MemoryReservation { export class MemoryMarking { /** @internal */ private readonly _size: number; /** @internal */ private _orchestrator?: MemoryOrchestrator; - /** @internal */ private _finalizationRegistry: FinalizationRegistry; private constructor(size: number, orchestrator: MemoryOrchestrator) { this._size = size; this._orchestrator = orchestrator; - this._finalizationRegistry = new FinalizationRegistry(orchestrator._onMarkFinalized); - this._finalizationRegistry.register(this, size, this); + this._orchestrator._markingsFinalizationRegistry.register(this, size, this); } public get size(): number { @@ -140,7 +141,7 @@ export class MemoryMarking { public dispose(): void { if (this._orchestrator != null) { this._orchestrator._onMarkFinalized(this._size); - this._finalizationRegistry.unregister(this); + this._orchestrator._markingsFinalizationRegistry.unregister(this); } this._orchestrator = undefined; diff --git a/src/chatWrappers/Gemma4ChatWrapper.ts b/src/chatWrappers/Gemma4ChatWrapper.ts index ed1c9a34..b5f1a1bd 100644 --- a/src/chatWrappers/Gemma4ChatWrapper.ts +++ b/src/chatWrappers/Gemma4ChatWrapper.ts @@ -13,6 +13,8 @@ export class Gemma4ChatWrapper extends ChatWrapper { public readonly reasoning: boolean; public readonly keepOnlyLastThought: boolean; + /** @internal */ private readonly _lineBreakAfterThinkIndicator: boolean; + public override readonly settings: ChatWrapperSettings = { supportsSystemMessages: true, functions: { @@ -52,17 +54,22 @@ export class Gemma4ChatWrapper extends ChatWrapper { * * Defaults to `true`. */ - keepOnlyLastThought?: boolean + keepOnlyLastThought?: boolean, + + /** @internal */ + _lineBreakAfterThinkIndicator?: boolean } = {}) { super(); const { reasoning = true, - keepOnlyLastThought = true + keepOnlyLastThought = true, + _lineBreakAfterThinkIndicator = true } = options; this.reasoning = reasoning; this.keepOnlyLastThought = keepOnlyLastThought; + this._lineBreakAfterThinkIndicator = _lineBreakAfterThinkIndicator; } public override generateContextState({ @@ -86,6 +93,9 @@ export class Gemma4ChatWrapper extends ChatWrapper { if (this.reasoning) systemMessage = LlamaText([ new SpecialTokensText("<|think|>"), + this._lineBreakAfterThinkIndicator + ? new SpecialTokensText("\n") + : "", systemMessage ]); @@ -250,6 +260,12 @@ export class Gemma4ChatWrapper extends ChatWrapper { {}, {}, {additionalRenderParameters: {"enable_thinking": true}} + ], + [{_lineBreakAfterThinkIndicator: false}, {}], + [ + {_lineBreakAfterThinkIndicator: false}, + {}, + {additionalRenderParameters: {"enable_thinking": true}} ] ]; } diff --git a/src/chatWrappers/QwenChatWrapper.ts b/src/chatWrappers/QwenChatWrapper.ts index cd34057c..62f52b66 100644 --- a/src/chatWrappers/QwenChatWrapper.ts +++ b/src/chatWrappers/QwenChatWrapper.ts @@ -14,7 +14,7 @@ export class QwenChatWrapper extends ChatWrapper { public readonly variation: "3" | "3.5"; public readonly keepOnlyLastThought: boolean; - public readonly thoughts: "auto" | "discourage"; + public readonly thoughts: "auto" | "discourage" | "modelInitiated"; /** @internal */ private readonly _flatFunctionResultString: boolean; /** @internal */ private readonly _ensureModelThoughtBeforeTextOnLastResponse: boolean; @@ -33,9 +33,12 @@ export class QwenChatWrapper extends ChatWrapper { /** * Control the usage of thoughts in the model responses. * + * When set to `"modelInitiated"`, thought segments won't be force-opened at the start of the model response, + * even if that's what the model would normally expect. + * * Defaults to `"auto"`. */ - thoughts?: "auto" | "discourage", + thoughts?: "auto" | "discourage" | "modelInitiated", /** * Chat template variation to use. @@ -145,7 +148,8 @@ export class QwenChatWrapper extends ChatWrapper { reiterateStackAfterFunctionCalls: true, thought: { prefix: LlamaText(new SpecialTokensText("\n")), - suffix: LlamaText(new SpecialTokensText("\n")) + suffix: LlamaText(new SpecialTokensText("\n")), + openOnResponseStart: thoughts === "auto" } } }; @@ -354,7 +358,7 @@ export class QwenChatWrapper extends ChatWrapper { /** @internal */ public static override _checkModelCompatibility(options: ChatWrapperCheckModelCompatibilityParams): boolean { - const architecture = options.fileInfo?.metadata.general.architecture; + const architecture = options.architecture; return ( architecture == null || architecture === GgufArchitectureType.qwen2 || @@ -390,26 +394,33 @@ export class QwenChatWrapper extends ChatWrapper { {_requireFunctionCallSettingsExtraction: true} ], - [ - {variation: "3.5"}, - {variation: "3.5"}, - {_requireFunctionCallSettingsExtraction: true, _functionCallExtractionExamineNonFirst: true} - ], - [ - {variation: "3.5", _lineBreakBeforeFunctionCallPrefix: true}, - {variation: "3.5"}, - {_requireFunctionCallSettingsExtraction: true, _functionCallExtractionExamineNonFirst: true} - ], - [ - {variation: "3.5", _ensureModelThoughtBeforeTextOnLastResponse: true, _lineBreakBeforeFunctionCallPrefix: true}, - {variation: "3.5"}, - {_requireFunctionCallSettingsExtraction: true, _functionCallExtractionExamineNonFirst: true} - ], - [ - {variation: "3.5", _ensureModelThoughtBeforeTextOnLastResponse: true}, - {variation: "3.5"}, - {_requireFunctionCallSettingsExtraction: true, _functionCallExtractionExamineNonFirst: true} - ] + ...([undefined, "modelInitiated"] satisfies (undefined | typeof this.prototype.thoughts)[]).flatMap((thoughts) => [ + [ + {variation: "3.5", thoughts}, + {variation: "3.5", thoughts}, + {_requireFunctionCallSettingsExtraction: true, _functionCallExtractionExamineNonFirst: true} + ], + [ + {variation: "3.5", thoughts, _lineBreakBeforeFunctionCallPrefix: true}, + {variation: "3.5", thoughts}, + {_requireFunctionCallSettingsExtraction: true, _functionCallExtractionExamineNonFirst: true} + ], + [ + { + variation: "3.5", + thoughts, + _ensureModelThoughtBeforeTextOnLastResponse: true, + _lineBreakBeforeFunctionCallPrefix: true + }, + {variation: "3.5", thoughts}, + {_requireFunctionCallSettingsExtraction: true, _functionCallExtractionExamineNonFirst: true} + ], + [ + {variation: "3.5", thoughts, _ensureModelThoughtBeforeTextOnLastResponse: true}, + {variation: "3.5", thoughts}, + {_requireFunctionCallSettingsExtraction: true, _functionCallExtractionExamineNonFirst: true} + ] + ]) ]; } } diff --git a/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts b/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts index 7ad8d990..39570a91 100644 --- a/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts +++ b/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts @@ -2,7 +2,7 @@ import {Template} from "@huggingface/jinja"; import {splitText} from "lifecycle-utils"; import { ChatHistoryItem, ChatModelFunctions, ChatUserMessage, ChatWrapperGenerateContextStateOptions, ChatWrapperGeneratedContextState, - ChatWrapperSettings, Tokenizer + ChatWrapperSettings, isChatModelResponseSegment, Tokenizer } from "../../types.js"; import {SpecialToken, LlamaText, SpecialTokensText} from "../../utils/LlamaText.js"; import {ChatWrapper} from "../../ChatWrapper.js"; @@ -27,6 +27,15 @@ import {extractSegmentSettingsFromTokenizerAndChatTemplate} from "./utils/extrac export type JinjaTemplateChatWrapperOptions = { template: string, + /** + * Whether to enable reasoning in the Jinja template. + * + * When set to `null`, the thinking setting will be omitted from the Jinja template, which would cause its default setting to be used. + * + * Defaults to `true`. + */ + reasoning?: boolean | null, + /** * Defaults to `"assistant"`. */ @@ -97,6 +106,15 @@ export type JinjaTemplateChatWrapperOptions = { */ segments?: TemplateChatWrapperSegmentsOptions, + /** + * Whether to keep only the chain of thought from the last model response. + * + * When `false`, all the chain of thoughts from the model responses will be kept in the context state. + * + * The default setting is extracted from the Jinja template, and the extraction fails, defaults to `false`. + */ + keepOnlyLastThought?: boolean, + /** * Pass a model's tokenizer to attempt to detect common tokens used for chat formatting from it. * @@ -154,12 +172,14 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { public override readonly settings: ChatWrapperSettings; public readonly template: string; + public readonly reasoning: boolean | null; public readonly modelRoleName: string; public readonly userRoleName: string; public readonly systemRoleName: string; public readonly convertUnsupportedSystemMessagesToUserMessages?: JinjaTemplateChatWrapperOptionsConvertMessageFormat; public readonly joinAdjacentMessagesOfTheSameType: boolean; public readonly trimLeadingWhitespaceInResponses: boolean; + public readonly keepOnlyLastThought: boolean; public readonly additionalRenderParameters?: Record; /** @internal */ private readonly _jinjaTemplate: Template; @@ -178,6 +198,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { const { template, + reasoning = true, modelRoleName = "assistant", userRoleName = "user", systemRoleName = "system", @@ -185,6 +206,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { functionCallMessageTemplate = "auto", joinAdjacentMessagesOfTheSameType = true, trimLeadingWhitespaceInResponses = true, + keepOnlyLastThought, additionalRenderParameters, segments, tokenizer, @@ -196,6 +218,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { throw new Error("template cannot be null"); this.template = template; + this.reasoning = reasoning; this.modelRoleName = modelRoleName; this.userRoleName = userRoleName; this.systemRoleName = systemRoleName; @@ -218,11 +241,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { const {supportsSystemMessages, needsToEndJinjaMessagesWithUserMessage} = this._runSanityTest(); this.settings = { ...this.settings, - supportsSystemMessages, - segments: { - ...this.settings.segments, - ...extractSegmentSettingsFromTokenizerAndChatTemplate(this.template, tokenizer) - } + supportsSystemMessages }; if (needsToEndJinjaMessagesWithUserMessage) @@ -233,138 +252,137 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { ? undefined : functionCallMessageTemplate ); - if (functionCallSettings == null && functionCallMessageTemplate !== "noJinja") { - try { - const renderTemplate: ExtractFunctionCallSettingsRenderTemplate = ({ - chatHistory, functions, additionalParams, stringifyFunctionParams, stringifyFunctionResults, - combineModelMessageAndToolCalls, squashModelTextResponses = true, setFunctionNameInResponse - }) => { - const render = ( - convertSystemMessagesToUserMessagesFormat: - JinjaTemplateChatWrapperOptionsConvertMessageFormat["format"] | undefined, - wipeFunctionCallIds: boolean | "align", - setFunctionNameInResponse?: string | boolean - ) => { - let inputChatHistory = chatHistory; - if (this._wrapFunctionParamsInsideMapKey != null) - inputChatHistory = inputChatHistory.map((item) => { - if (item.type !== "model") - return item; + const renderTemplate: ExtractFunctionCallSettingsRenderTemplate = ({ + chatHistory, functions, additionalParams, stringifyFunctionParams, stringifyFunctionResults, + combineModelMessageAndToolCalls, squashModelTextResponses = true, setFunctionNameInResponse + }) => { + const render = ( + convertSystemMessagesToUserMessagesFormat: JinjaTemplateChatWrapperOptionsConvertMessageFormat["format"] | undefined, + wipeFunctionCallIds: boolean | "align", + setFunctionNameInResponse?: string | boolean + ) => { + let inputChatHistory = chatHistory; + if (this._wrapFunctionParamsInsideMapKey != null) + inputChatHistory = inputChatHistory.map((item) => { + if (item.type !== "model") + return item; + + return { + ...item, + response: item.response.map((response) => { + if (typeof response === "string" || response.type !== "functionCall") + return response; return { - ...item, - response: item.response.map((response) => { - if (typeof response === "string" || response.type !== "functionCall") - return response; - - return { - ...response, - params: {[this._wrapFunctionParamsInsideMapKey!]: response.params} - }; - }) + ...response, + params: {[this._wrapFunctionParamsInsideMapKey!]: response.params} }; - }); - - const {messages: intermediateMessages, tools} = fromChatHistoryToIntermediateOpenAiMessages({ - chatHistory: this._transformChatHistory(inputChatHistory, { - convertSystemMessagesToUserMessagesFormat, - joinAdjacentMessagesOfTheSameType: !squashModelTextResponses - ? false - : undefined - }).transformedHistory, - chatWrapperSettings: this.settings, - useRawValues: false, - functions, - stringifyFunctionParams, - stringifyFunctionResults, - combineModelMessageAndToolCalls, - squashModelTextResponses, - setFunctionNameInResponse - }); - - const messages = fromIntermediateToCompleteOpenAiMessages(intermediateMessages) - .map((item) => { - if (!wipeFunctionCallIds) - return item; - - if (item.role === "assistant" && item["tool_calls"] != null && item["tool_calls"].length > 0) { - for (const toolCall of item["tool_calls"]) { - if (wipeFunctionCallIds === "align") - toolCall.id = "fc_1_0001"; - else - delete (toolCall as {id?: string}).id; - } - } else if (item.role === "tool") { - if (wipeFunctionCallIds === "align") - item["tool_call_id"] = "fc_1_0001"; - else - delete (item as {"tool_call_id"?: string})["tool_call_id"]; - } - - return item; - }); - - const lastJinjaItem = messages.at(-1); - let eraseRenderedJinjaFromId: string | undefined; - if (this._endJinjaMessagesWithUserMessage && lastJinjaItem?.role === this.modelRoleName && - typeof lastJinjaItem.content === "string" && - lastJinjaItem.content.length > 0 && - ( - (lastJinjaItem as OpenAiChatAssistantMessage)["tool_calls"] == null || - (lastJinjaItem as OpenAiChatAssistantMessage)["tool_calls"]?.length === 0 - ) - ) { - eraseRenderedJinjaFromId = lastJinjaItem.content; - messages.push({ - role: this.userRoleName, - content: idsGenerator.generateId() - } as OpenAiChatMessage); - } + }) + }; + }); - let res = this._jinjaTemplate.render({ - ...( - this.additionalRenderParameters == null - ? {} - : structuredClone(this.additionalRenderParameters) - ), - ...additionalParams, - messages, - ...removeUndefinedFields({tools}) - }); - - if (eraseRenderedJinjaFromId != null) { - const eraseIndex = res.lastIndexOf(eraseRenderedJinjaFromId); - if (eraseIndex >= 0) - res = res.slice(0, eraseIndex + eraseRenderedJinjaFromId.length); + const {messages: intermediateMessages, tools} = fromChatHistoryToIntermediateOpenAiMessages({ + chatHistory: this._transformChatHistory(inputChatHistory, { + convertSystemMessagesToUserMessagesFormat, + joinAdjacentMessagesOfTheSameType: !squashModelTextResponses + ? false + : undefined + }).transformedHistory, + chatWrapperSettings: this.settings, + useRawValues: false, + functions, + stringifyFunctionParams, + stringifyFunctionResults, + combineModelMessageAndToolCalls, + squashModelTextResponses, + setFunctionNameInResponse + }); + + const messages = fromIntermediateToCompleteOpenAiMessages(intermediateMessages) + .map((item) => { + if (!wipeFunctionCallIds) + return item; + + if (item.role === "assistant" && item["tool_calls"] != null && item["tool_calls"].length > 0) { + for (const toolCall of item["tool_calls"]) { + if (wipeFunctionCallIds === "align") + toolCall.id = "fc_1_0001"; + else + delete (toolCall as {id?: string}).id; + } + } else if (item.role === "tool") { + if (wipeFunctionCallIds === "align") + item["tool_call_id"] = "fc_1_0001"; + else + delete (item as {"tool_call_id"?: string})["tool_call_id"]; } - // attempt to remove the ID pattern from the output - if (wipeFunctionCallIds === "align") - res = res - .replaceAll(/,\s*"(tool_call_id|call_id|id)":\s*"fc_1_0001"/g, "") - .replaceAll(/"(tool_call_id|call_id|id)":\s*"fc_1_0001"\s*,/g, ""); + return item; + }); - return res; - }; + const lastJinjaItem = messages.at(-1); + let eraseRenderedJinjaFromId: string | undefined; + if (this._endJinjaMessagesWithUserMessage && lastJinjaItem?.role === this.modelRoleName && + typeof lastJinjaItem.content === "string" && + lastJinjaItem.content.length > 0 && + ( + (lastJinjaItem as OpenAiChatAssistantMessage)["tool_calls"] == null || + (lastJinjaItem as OpenAiChatAssistantMessage)["tool_calls"]?.length === 0 + ) + ) { + eraseRenderedJinjaFromId = lastJinjaItem.content; + messages.push({ + role: this.userRoleName, + content: idsGenerator.generateId() + } as OpenAiChatMessage); + } - return tryMatrix({ - convertSystemMessagesToUserMessagesFormat: - getConvertUnsupportedSystemMessagesToUserMessagesTryOptions( - this.convertUnsupportedSystemMessagesToUserMessages - ), - wipeFunctionCallIds: [true, "align", false], - setFunctionNameInResponse: setFunctionNameInResponse == null - ? [false] - : [false, setFunctionNameInResponse] - }, ({convertSystemMessagesToUserMessagesFormat, wipeFunctionCallIds, setFunctionNameInResponse}) => { - return render(convertSystemMessagesToUserMessagesFormat, wipeFunctionCallIds, setFunctionNameInResponse); - }); - }; - const idsGenerator = new UniqueIdGenerator( - this.template + this.modelRoleName + this.userRoleName + this.systemRoleName + - (this.convertUnsupportedSystemMessagesToUserMessages?.format ?? "") - ); + let res = this._jinjaTemplate.render({ + ...( + this.additionalRenderParameters == null + ? {} + : structuredClone(this.additionalRenderParameters) + ), + ...additionalParams, + messages, + ...removeUndefinedFields({tools}) + }); + + if (eraseRenderedJinjaFromId != null) { + const eraseIndex = res.lastIndexOf(eraseRenderedJinjaFromId); + if (eraseIndex >= 0) + res = res.slice(0, eraseIndex + eraseRenderedJinjaFromId.length); + } + + // attempt to remove the ID pattern from the output + if (wipeFunctionCallIds === "align") + res = res + .replaceAll(/,\s*"(tool_call_id|call_id|id)":\s*"fc_1_0001"/g, "") + .replaceAll(/"(tool_call_id|call_id|id)":\s*"fc_1_0001"\s*,/g, ""); + + return res; + }; + return tryMatrix({ + convertSystemMessagesToUserMessagesFormat: + getConvertUnsupportedSystemMessagesToUserMessagesTryOptions( + this.convertUnsupportedSystemMessagesToUserMessages + ), + wipeFunctionCallIds: [true, "align", false], + setFunctionNameInResponse: setFunctionNameInResponse == null + ? [false] + : [false, setFunctionNameInResponse] + }, ({convertSystemMessagesToUserMessagesFormat, wipeFunctionCallIds, setFunctionNameInResponse}) => { + return render(convertSystemMessagesToUserMessagesFormat, wipeFunctionCallIds, setFunctionNameInResponse); + }); + }; + const idsGenerator = new UniqueIdGenerator( + this.template + this.modelRoleName + this.userRoleName + this.systemRoleName + + (this.convertUnsupportedSystemMessagesToUserMessages?.format ?? "") + ); + + if (functionCallSettings == null && functionCallMessageTemplate !== "noJinja") { + try { this._wrapFunctionParamsInsideMapKey = detectNeedToWrapFunctionArgumentsWithMap({idsGenerator, renderTemplate}); const extractedSettings = extractFunctionCallSettingsFromJinjaTemplate({ idsGenerator, @@ -386,10 +404,31 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { throw new Error("failed to extract function call settings from the Jinja template"); } + const extractedSegmentSettings = extractSegmentSettingsFromTokenizerAndChatTemplate({ + chatTemplate: this.template, + tokenizer, + renderRawJinjaTemplate: (params: Record) => { + return this._jinjaTemplate.render({ + ...( + this.additionalRenderParameters == null + ? {} + : structuredClone(this.additionalRenderParameters) + ), + ...params + }); + }, + idsGenerator, + enableReasoning: this.reasoning + }); this.settings = { ...this.settings, - functions: functionCallSettings ?? ChatWrapper.defaultSettings.functions + functions: functionCallSettings ?? ChatWrapper.defaultSettings.functions, + segments: { + ...this.settings.segments, + ...extractedSegmentSettings.settings + } }; + this.keepOnlyLastThought = keepOnlyLastThought ?? extractedSegmentSettings.keepOnlyLastThought ?? false; } /** @@ -565,7 +604,9 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { } = this._transformChatHistory(history, {convertSystemMessagesToUserMessagesFormat, availableFunctions, documentFunctionParams}); const generateMessagesWithEmbeddedTools = (chatHistory: readonly ChatHistoryItem[]) => ({ - messages: chatHistory.map((item): IntermediateOpenAiMessage => { + messages: chatHistory.map((item, index): IntermediateOpenAiMessage => { + const isLastItem = index === chatHistory.length - 1; + if (item.type === "system") return { role: "system", @@ -579,7 +620,13 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { else if (item.type === "model") return { role: "assistant", - content: this.generateModelResponseText(item.response) + content: this.generateModelResponseText( + (!this.keepOnlyLastThought || isLastItem) + ? item.response + : item.response.filter((response) => ( + !isChatModelResponseSegment(response) || response.segmentType !== "thought") + ) + ) }; void (item satisfies never); @@ -648,6 +695,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { } as const; const idToContent = new Map(); const modelMessageIds = new Set(); + const lastModelMessageIds = new Set(); const messageIds = new Set(); for (const intermediateMessage of intermediateMessages) { @@ -669,8 +717,11 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { content: id } as OpenAiChatMessage); - if (intermediateMessage.role === "assistant" || intermediateMessage.role === "tool") + if (intermediateMessage.role === "assistant" || intermediateMessage.role === "tool") { modelMessageIds.add(id); + lastModelMessageIds.add(id); + } else if (intermediateMessage.role === "user") + lastModelMessageIds.clear(); } const bosTokenId = idsGenerator.generateId(); @@ -713,6 +764,11 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { "bos_token": bosTokenId, "eos_token": eosTokenId, "eot_token": eotTokenId, + ...( + this.reasoning == null + ? {} + : {"enable_thinking": this.reasoning} + ), ...options }) )); @@ -781,21 +837,59 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { const {splitJinjaParts, stopGenerationJinjaParts} = renderJinjaAndSplitIntoParts(); const messageIdsLeftToProcess = new Set(messageIds); - const contextText = LlamaText( - splitJinjaParts.map((part) => { - if (typeof part === "string") - return new SpecialTokensText(part); // things that are not message content can be tokenized with special tokens - - const message = idToContent.get(part.separator); - - if (message == null) - throw new Error(`Message with id "${part.separator}" not found`); + const thoughSegmentPrefix = getLlamaTextOnlyText(this.settings.segments?.thought?.prefix); + const thoughSegmentSuffix = getLlamaTextOnlyText(this.settings.segments?.thought?.suffix); + let inLastModelResponseSection: boolean | null = ( + thoughSegmentPrefix == null || + thoughSegmentSuffix == null || + this.settings.segments?.thought?.openOnResponseStart !== true + ) + ? null + : false; + const llamaTextContent: Array = []; + for (let i = 0; i < splitJinjaParts.length; i++) { + const part = splitJinjaParts[i]!; + + if (typeof part === "string") { + // things that are not message content can be tokenized with special tokens + llamaTextContent.push(new SpecialTokensText(part)); + continue; + } - messageIdsLeftToProcess.delete(part.separator); + const message = idToContent.get(part.separator); + + if (message == null) + throw new Error(`Message with id "${part.separator}" not found`); + + messageIdsLeftToProcess.delete(part.separator); + + if (inLastModelResponseSection === false && lastModelMessageIds.has(part.separator)) + inLastModelResponseSection = true; + + // remove empty thinking blocks added by the template if the last model response is supposed to always have a thought + // segment rendered by the chat wrapper + if (inLastModelResponseSection === true) { + const lastPart = llamaTextContent.at(-1); + if (lastPart instanceof SpecialTokensText && thoughSegmentPrefix != null && thoughSegmentSuffix != null) { + const thoughPrefixIndex = lastPart.value.indexOf(thoughSegmentPrefix); + const thoughSuffixIndex = thoughPrefixIndex >= 0 + ? lastPart.value.indexOf(thoughSegmentSuffix, thoughPrefixIndex + thoughSegmentPrefix.length) + : -1; + + if (thoughPrefixIndex >= 0 && thoughSuffixIndex >= 0) { + const thoughtText = lastPart.value.slice(thoughPrefixIndex + thoughSegmentPrefix.length, thoughSuffixIndex); + if (thoughtText.trim() === "") + llamaTextContent[llamaTextContent.length - 1] = new SpecialTokensText( + lastPart.value.slice(0, thoughPrefixIndex) + + lastPart.value.slice(thoughSuffixIndex + thoughSegmentSuffix.length) + ); + } + } + } - return message; - }) - ); + llamaTextContent.push(message); + } + const contextText = LlamaText(llamaTextContent); if (messageIdsLeftToProcess.size !== 0) throw new Error("Some input messages are not present in the generated Jinja template output"); @@ -919,6 +1013,24 @@ function getConvertUnsupportedSystemMessagesToUserMessagesTryOptions( return [undefined, convertUnsupportedSystemMessagesToUserMessages.format]; } +function getLlamaTextOnlyText(llamaText: LlamaText | string | undefined): string | undefined { + if (llamaText == null || typeof llamaText === "string") + return llamaText; + + const texts: string[] = []; + + for (const value of llamaText.values) { + if (typeof value === "string") + texts.push(value); + else if (value instanceof SpecialTokensText) + texts.push(value.value); + else + return undefined; + } + + return texts.join(""); +} + const chatHistoriesForSanityTest: ChatHistoryItem[][] = [ [{ type: "system", diff --git a/src/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.ts b/src/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.ts index 1d18e631..3fc1e035 100644 --- a/src/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.ts +++ b/src/chatWrappers/generic/utils/extractFunctionCallSettingsFromJinjaTemplate.ts @@ -426,7 +426,7 @@ export function extractFunctionCallSettingsFromJinjaTemplate({ const callPrefixLength = findCommonEndLength(modelMessage1ToFunc1Name.text, func1ParamsToFunc2Name.text); const callPrefixText = func1ParamsToFunc2Name.text.slice(func1ParamsToFunc2Name.text.length - callPrefixLength); const parallelismCallPrefix = modelMessage1ToFunc1Name.text.slice(0, modelMessage1ToFunc1Name.text.length - callPrefixLength); - + const callSuffixAndParallelismBetweenCallsText = func1ParamsToFunc2Name.text.slice( 0, func1ParamsToFunc2Name.text.length - callPrefixLength diff --git a/src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts b/src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts index f4c6a949..3e564522 100644 --- a/src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts +++ b/src/chatWrappers/generic/utils/extractSegmentSettingsFromTokenizerAndChatTemplate.ts @@ -1,11 +1,35 @@ import {ChatWrapperSettings, Tokenizer} from "../../../types.js"; import {LlamaText, SpecialTokensText} from "../../../utils/LlamaText.js"; +import {tryMatrix} from "../../../utils/optionsMatrix.js"; import {removeUndefinedFields} from "../../../utils/removeNullFields.js"; +import {OpenAiChatMessage} from "../../../utils/OpenAIFormat.js"; +import {UniqueIdGenerator} from "./UniqueIdGenerator.js"; -export function extractSegmentSettingsFromTokenizerAndChatTemplate( - chatTemplate: string | undefined, tokenizer?: Tokenizer -): ChatWrapperSettings["segments"] { - function tryMatchPrefixSuffixPair(tryMatchGroups: [prefix: string, suffix: string][]) { +const knownThinkingSegmentControls = new Map([ + ["", ""], // DeepSeek, QwQ + ["", ""], // EXAONE Deep + ["[THINK]", "[/THINK]"], // Mistral + ["<|START_THINKING|>", "<|END_THINKING|>"], // Command R7B + ["<|begin_of_thought|>", "<|end_of_thought|>"] // JoyAI +]); + +export function extractSegmentSettingsFromTokenizerAndChatTemplate({ + chatTemplate, + tokenizer, + renderRawJinjaTemplate, + idsGenerator, + enableReasoning +}: { + chatTemplate: string | undefined, + tokenizer: Tokenizer | undefined, + renderRawJinjaTemplate(params: Record): string, + idsGenerator: UniqueIdGenerator, + enableReasoning: boolean | null +}): { + settings: ChatWrapperSettings["segments"], + keepOnlyLastThought?: boolean +} { + function tryMatchPrefixSuffixPair(tryMatchGroups: Iterable<[prefix: string, suffix: string]>) { if (chatTemplate != null) { for (const [prefix, suffix] of tryMatchGroups) { if ( @@ -74,15 +98,349 @@ export function extractSegmentSettingsFromTokenizerAndChatTemplate( return undefined; } - return removeUndefinedFields({ - thought: tryMatchPrefixSuffixPair([ - ["", ""], // DeepSeek, QwQ - ["", ""], // EXAONE Deep - ["[THINK]", "[/THINK]"], // Mistral - ["<|START_THINKING|>", "<|END_THINKING|>"], // Command R7B - ["<|begin_of_thought|>", "<|end_of_thought|>"] // JoyAI - ]) - }); + function extractThoughtSettingsFromRendering(): ( + { + thoughtSegment: Exclude["thought"] | undefined, + keepPastReasoning: boolean | undefined + } + ) { + if (chatTemplate == null) + return { + thoughtSegment: undefined, + keepPastReasoning: undefined + }; + + const systemMessage = idsGenerator.generateId(); + const userMessage1 = idsGenerator.generateId(); + const userMessage2 = idsGenerator.generateId(); + const modelResponse1 = idsGenerator.generateId(); + const modelResponse2 = idsGenerator.generateId(); + const modelResponse3 = idsGenerator.generateId(); + const modelReasoning2 = idsGenerator.generateId(); + const modelReasoning3 = idsGenerator.generateId(); + + const bosTokenId = idsGenerator.generateId(); + const eosTokenId = idsGenerator.generateId(); + const eotTokenId = idsGenerator.generateId(); + + const renderTemplate = (messages: OpenAiChatMessage[], params?: Record) => tryMatrix({ + skipSystemPrompt: [false, true] + }, ({skipSystemPrompt}) => { + let messageToRender = messages; + if (skipSystemPrompt) + messageToRender = messageToRender.slice(1); + + return renderRawJinjaTemplate({ + messages: messageToRender, + "bos_token": bosTokenId, + "eos_token": eosTokenId, + "eot_token": eotTokenId, + ...params + }); + }); + + const baseMessages: OpenAiChatMessage[] = [{ + role: "system", + content: systemMessage + }, { + role: "user", + content: userMessage1 + }]; + const longBaseMessages: OpenAiChatMessage[] = [...baseMessages, { + role: "assistant", + content: modelResponse1 + }, { + role: "user", + content: userMessage2 + }]; + + const messagesWithModelResponse: OpenAiChatMessage[] = [...baseMessages, { + role: "assistant", + content: modelResponse2 + }]; + const messagesWithModelResponseLongBase: OpenAiChatMessage[] = [...longBaseMessages, { + role: "assistant", + content: modelResponse2 + }]; + + const messagesWithModelReasoning: OpenAiChatMessage[] = [...baseMessages, { + role: "assistant", + content: modelResponse2, + "reasoning_content": modelReasoning2 + }]; + const messagesWithModelReasoningLongBase: OpenAiChatMessage[] = [...longBaseMessages, { + role: "assistant", + content: modelResponse2, + "reasoning_content": modelReasoning2 + }]; + + function extractControls() { + const {responseOnly, withReasoning} = tryMatrix({ + enableThinking: [true, null], + variation: ["simple", "separateReasoning", "nullContent", "reasoningFirst", "reasoningFirstNullMessage"] + }, ({enableThinking, variation}) => { + const thinkingParam = enableThinking === true + ? {"enable_thinking": true} + : {}; + + if (variation === "simple") + return { + responseOnly: renderTemplate(messagesWithModelResponseLongBase, thinkingParam), + withReasoning: { + long: renderTemplate(messagesWithModelReasoningLongBase, thinkingParam), + short: renderTemplate(messagesWithModelReasoning, thinkingParam) + } + }; + else if (variation === "separateReasoning" || variation === "nullContent") + return { + responseOnly: renderTemplate([...messagesWithModelResponseLongBase, { + role: "assistant", + content: "" + }], thinkingParam), + withReasoning: { + long: renderTemplate([...messagesWithModelResponseLongBase, { + role: "assistant", + ...(variation === "nullContent" ? {} : { + content: "" + }), + "reasoning_content": modelReasoning2 + }], thinkingParam), + short: renderTemplate([...messagesWithModelResponse, { + role: "assistant", + ...(variation === "nullContent" ? {} : { + content: "" + }), + "reasoning_content": modelReasoning2 + }], thinkingParam) + } + }; + else if (variation === "reasoningFirst" || variation === "reasoningFirstNullMessage") + return { + responseOnly: renderTemplate(messagesWithModelResponseLongBase, thinkingParam), + withReasoning: { + long: renderTemplate([...longBaseMessages, { + role: "assistant", + ...(variation === "reasoningFirstNullMessage" ? {} : { + content: "" + }), + "reasoning_content": modelReasoning2 + }, { + role: "assistant", + content: modelResponse2 + }], thinkingParam), + short: renderTemplate([...baseMessages, { + role: "assistant", + ...(variation === "reasoningFirstNullMessage" ? {} : { + content: "" + }), + "reasoning_content": modelReasoning2 + }, { + role: "assistant", + content: modelResponse2 + }], thinkingParam) + } + }; + + void (variation satisfies never); + throw new Error(`Unsupported variation: ${variation}`); + }); + + let reasoningSectionStartPrefix: string | undefined = undefined; + let reasoningSectionEndPrefix: string | undefined = undefined; + + if (responseOnly === withReasoning.long) + return undefined; + + const modelResponseIndex = responseOnly.indexOf(modelResponse1); + if (modelResponseIndex < 0) + return undefined; + + 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); + + const reasoningContentEndIndex = reasoningContentStartIndex + modelReasoning2.length; + const modelResponseStartIndex = withReasoning.short.indexOf(modelResponse2, reasoningContentEndIndex); + if (modelResponseStartIndex < 0) + return undefined; + + reasoningSectionEndPrefix = withReasoning.short.slice(reasoningContentEndIndex, modelResponseStartIndex); + + return { + prefix: reasoningSectionStartPrefix, + suffix: reasoningSectionEndPrefix + }; + } + + function shouldKeepPastThinking() { + const renderedOutput = tryMatrix({ + enableThinking: [true, null], + variation: ["simple", "reasoningFirst", "reasoningFirstNullMessage"] + }, ({enableThinking, variation}) => { + const thinkingParam = enableThinking === true + ? {"enable_thinking": true} + : {}; + + if (variation === "simple") + return renderTemplate([...messagesWithModelReasoning, { + role: "user", + content: userMessage2 + }, { + role: "assistant", + content: modelResponse3, + "reasoning_content": modelReasoning3 + }], thinkingParam); + else if (variation === "reasoningFirst" || variation === "reasoningFirstNullMessage") + return renderTemplate([...baseMessages, { + role: "assistant", + ...(variation === "reasoningFirstNullMessage" ? {} : { + content: "" + }), + "reasoning_content": modelReasoning2 + }, { + role: "assistant", + content: modelResponse2 + }, { + role: "user", + content: userMessage2 + }, { + role: "assistant", + ...(variation === "reasoningFirstNullMessage" ? {} : { + content: "" + }), + "reasoning_content": modelReasoning3 + }, { + role: "assistant", + content: modelResponse3 + }], thinkingParam); + + void (variation satisfies never); + throw new Error(`Unsupported variation: ${variation}`); + }); + + return renderedOutput.includes(modelReasoning2) && renderedOutput.includes(modelReasoning3); + } + + function shouldOpenThinkingSegmentOnModelResponseStart(reasoningSectionPrefix: string, reasoningSectionSuffix: string) { + if (!enableReasoning) + return false; + + const {responseOnly, withGenerationPrompt} = tryMatrix({ + enableThinking: enableReasoning + ? [true, null] + : [null] + }, ({enableThinking}) => { + const thinkingParam = (enableThinking === true || enableThinking === false) + ? {"enable_thinking": enableThinking} + : {}; + + return { + responseOnly: renderTemplate(baseMessages, thinkingParam), + withGenerationPrompt: renderTemplate(baseMessages, { + ...thinkingParam, + "add_generation_prompt": true + }) + }; + }); + + if (responseOnly === withGenerationPrompt) + return false; + + const userMessage1Index = responseOnly.indexOf(userMessage1); + if (userMessage1Index < 0) + return false; + + const withReasoningUserMessage1Index = withGenerationPrompt.indexOf(userMessage1); + if (withReasoningUserMessage1Index < 0) + return false; + + if (responseOnly.indexOf(reasoningSectionPrefix, userMessage1Index) >= 0) + return false; + + const reasoningSectionPrefixIndex = withGenerationPrompt.indexOf(reasoningSectionPrefix, withReasoningUserMessage1Index); + if (reasoningSectionPrefixIndex < 0) + return false; + + const reasoningSectionSuffixIndex = withGenerationPrompt.indexOf(reasoningSectionSuffix, reasoningSectionPrefixIndex); + if (reasoningSectionSuffixIndex >= 0) { + const reasoningSectionContent = withGenerationPrompt.slice( + reasoningSectionPrefixIndex + reasoningSectionPrefix.length, + reasoningSectionSuffixIndex + ); + + if (reasoningSectionContent.trim() === "") + return false; + } + + return true; + } + + let controls: ReturnType; + let keepPastReasoning: boolean | undefined = undefined; + let openOnResponseStart: boolean | undefined = undefined; + try { + controls = extractControls(); + + if (controls == null || controls.prefix.trim() === "") + return { + thoughtSegment: undefined, + keepPastReasoning: undefined + }; + } catch (err) { + return { + thoughtSegment: undefined, + keepPastReasoning: undefined + }; + } + + try { + keepPastReasoning = shouldKeepPastThinking(); + } catch (err) { + // do nothing + } + + try { + openOnResponseStart = controls != null && shouldOpenThinkingSegmentOnModelResponseStart(controls.prefix, controls.suffix); + } catch (err) { + // do nothing + } + + const thoughtSuffix = controls.suffix.trim() === "" + ? ( + knownThinkingSegmentControls.get(controls.prefix) ?? + knownThinkingSegmentControls.get(controls.prefix.trim()) + ) + : controls.suffix; + + return { + thoughtSegment: { + prefix: LlamaText(new SpecialTokensText(controls.prefix)), + suffix: thoughtSuffix != null + ? LlamaText(new SpecialTokensText(thoughtSuffix)) + : undefined, + openOnResponseStart + }, + keepPastReasoning + }; + } + + const extractedFromRendering = extractThoughtSettingsFromRendering(); + + return { + settings: removeUndefinedFields({ + thought: extractedFromRendering.thoughtSegment ?? tryMatchPrefixSuffixPair(knownThinkingSegmentControls) + }), + keepOnlyLastThought: !extractedFromRendering.keepPastReasoning + }; } function hasAll(text: string, matches: string[]) { diff --git a/src/chatWrappers/utils/isJinjaTemplateEquivalentToSpecializedChatWrapper.ts b/src/chatWrappers/utils/isJinjaTemplateEquivalentToSpecializedChatWrapper.ts index 20d9e35b..fd82af7a 100644 --- a/src/chatWrappers/utils/isJinjaTemplateEquivalentToSpecializedChatWrapper.ts +++ b/src/chatWrappers/utils/isJinjaTemplateEquivalentToSpecializedChatWrapper.ts @@ -136,6 +136,11 @@ function checkEquivalence( (specializedChatWrapper as Writable).settings = originalSpecializedSettings; } + if (jinjaChatWrapper.settings.segments?.thought?.openOnResponseStart === true && + specializedChatWrapper.settings.segments?.thought?.openOnResponseStart !== true + ) + return false; + if (!compareContextTexts(jinjaRes.contextText, specializedWrapperRes.contextText, tokenizer)) return false; diff --git a/src/chatWrappers/utils/resolveChatWrapper.ts b/src/chatWrappers/utils/resolveChatWrapper.ts index b1faa330..bbbca3b3 100644 --- a/src/chatWrappers/utils/resolveChatWrapper.ts +++ b/src/chatWrappers/utils/resolveChatWrapper.ts @@ -24,6 +24,7 @@ import {SeedChatWrapper} from "../SeedChatWrapper.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([ @@ -88,6 +89,7 @@ export type ResolveChatWrapperOptions = { type?: "auto" | SpecializedChatWrapperTypeName | TemplateChatWrapperTypeName, bosString?: string | null, + architecture?: GgufArchitectureType, filename?: string, fileInfo?: GgufFileInfo, tokenizer?: Tokenizer, @@ -182,6 +184,7 @@ export type ResolveChatWrapperWithModelOptions = { * * const chatWrapper = resolveChatWrapper({ * bosString: model.tokens.bosString, + * architecture: model.architecture, * filename: model.filename, * fileInfo: model.fileInfo, * tokenizer: model.tokenizer @@ -199,6 +202,7 @@ export function resolveChatWrapper( ...(modelOptions ?? {}), customWrapperSettings: modelOptions?.customWrapperSettings as ResolveChatWrapperOptions["customWrapperSettings"], bosString: options.tokens.bosString, + architecture: options.fileInfo?.metadata?.general?.architecture, filename: options.filename, fileInfo: options.fileInfo, tokenizer: options.tokenizer @@ -207,6 +211,7 @@ export function resolveChatWrapper( const { type = "auto", bosString, + architecture: archOption, filename, fileInfo, tokenizer, @@ -216,6 +221,8 @@ export function resolveChatWrapper( noJinja = false } = options; + const architecture = archOption ?? fileInfo?.metadata?.general?.architecture; + function createSpecializedChatWrapper( specializedChatWrapper: T, defaultSettings: ConstructorParameters[0] = {} @@ -293,7 +300,8 @@ export function resolveChatWrapper( const isCompatible = Wrapper._checkModelCompatibility({ tokenizer, - fileInfo + fileInfo, + architecture }); if (!isCompatible) @@ -315,12 +323,11 @@ export function resolveChatWrapper( : undefined; const testChatWrapperSettings = { - ...(wrapperSettings ?? {}), ...(testConfig ?? {}) }; const applyChatWrapperSettings = { - ...(wrapperSettings ?? {}), - ...(applyConfig ?? {}) + ...(applyConfig ?? {}), + ...(wrapperSettings ?? {}) }; const chatWrapper = new (Wrapper as any)(testChatWrapperSettings); @@ -359,9 +366,9 @@ export function resolveChatWrapper( } for (const modelNames of getModelLinageNames(fileInfo?.metadata)) { - if (includesText(modelNames, ["llama 3.2", "llama-3.2", "llama3.2"]) && Llama3_2LightweightChatWrapper._checkModelCompatibility({tokenizer, fileInfo})) + if (includesText(modelNames, ["llama 3.2", "llama-3.2", "llama3.2"]) && Llama3_2LightweightChatWrapper._checkModelCompatibility({tokenizer, fileInfo, architecture})) return createSpecializedChatWrapper(Llama3_2LightweightChatWrapper); - else if (includesText(modelNames, ["llama 3.1", "llama-3.1", "llama3.1"]) && Llama3_1ChatWrapper._checkModelCompatibility({tokenizer, fileInfo})) + else if (includesText(modelNames, ["llama 3.1", "llama-3.1", "llama3.1"]) && Llama3_1ChatWrapper._checkModelCompatibility({tokenizer, fileInfo, architecture})) return createSpecializedChatWrapper(Llama3_1ChatWrapper); else if (includesText(modelNames, ["llama 3", "llama-3", "llama3"])) return createSpecializedChatWrapper(Llama3ChatWrapper); @@ -395,7 +402,7 @@ export function resolveChatWrapper( addSpaceBeforeEos: modelJinjaTemplate.includes("' ' + eos_token") }); else if (modelJinjaTemplate.includes("<|start_header_id|>") && modelJinjaTemplate.includes("<|end_header_id|>")) { - if (Llama3_1ChatWrapper._checkModelCompatibility({tokenizer, fileInfo})) + if (Llama3_1ChatWrapper._checkModelCompatibility({tokenizer, fileInfo, architecture})) return createSpecializedChatWrapper(Llama3_1ChatWrapper); else return createSpecializedChatWrapper(Llama3ChatWrapper); @@ -455,16 +462,14 @@ export function resolveChatWrapper( } } - if (fileInfo != null) { - const arch = fileInfo.metadata.general?.architecture; - - if (arch === "llama") + if (architecture != null) { + if (architecture === "llama") return createSpecializedChatWrapper(GeneralChatWrapper); - else if (arch === "falcon") + else if (architecture === "falcon") return createSpecializedChatWrapper(FalconChatWrapper); - else if (arch === "gemma" || arch === "gemma2") + else if (architecture === "gemma" || architecture === "gemma2") return createSpecializedChatWrapper(GemmaChatWrapper); - else if (arch === "gemma4") + else if (architecture === "gemma4") return createSpecializedChatWrapper(Gemma4ChatWrapper); } diff --git a/src/cli/commands/inspect/commands/InspectEstimateCommand.ts b/src/cli/commands/inspect/commands/InspectEstimateCommand.ts index 8f152a02..c6b735fd 100644 --- a/src/cli/commands/inspect/commands/InspectEstimateCommand.ts +++ b/src/cli/commands/inspect/commands/InspectEstimateCommand.ts @@ -347,10 +347,10 @@ export const InspectEstimateCommand: CommandModule const resolvedKvCacheKeyType = kvCacheKeyType === "currentQuant" ? ggufInsights.dominantTensorType ?? GgmlType.F16 - : resolveGgmlTypeOption(kvCacheKeyType) ?? GgmlType.F16; + : resolveGgmlTypeOption(kvCacheKeyType, llama) ?? GgmlType.F16; const resolvedKvCacheValueType = kvCacheValueType === "currentQuant" ? ggufInsights.dominantTensorType ?? GgmlType.F16 - : resolveGgmlTypeOption(kvCacheValueType) ?? GgmlType.F16; + : resolveGgmlTypeOption(kvCacheValueType, llama) ?? GgmlType.F16; if (resolvedKvCacheKeyType != GgmlType.F16 || resolvedKvCacheValueType != GgmlType.F16) console.info(`${chalk.yellow("KV cache:")} ${GgmlType[resolvedKvCacheKeyType] + " " + GgmlType[resolvedKvCacheValueType]}`); @@ -609,7 +609,7 @@ function renderDiffPercentageWithColors(percentage: number, { } = {}): string { if (nanIsZero && Number.isNaN(percentage)) percentage = 0; - + const percentageText = percentage.toFixed(2).padStart(5, "0") + "%"; const absPercentage = Math.abs(percentage); @@ -916,7 +916,7 @@ async function runTestWorkerLogic() { batchSize, failedCreationRemedy: false }); - + if (evaluateText != null && evaluateText != "") { const sequence = context.getSequence(); await sequence.evaluateWithoutGeneratingNewTokens(model.tokenize(evaluateText)); diff --git a/src/cli/utils/resolveCommandGgufPath.ts b/src/cli/utils/resolveCommandGgufPath.ts index cd50ddca..4a6e25a9 100644 --- a/src/cli/utils/resolveCommandGgufPath.ts +++ b/src/cli/utils/resolveCommandGgufPath.ts @@ -32,10 +32,10 @@ export async function resolveCommandGgufPath(ggufPath: string | undefined, llama useMmap, kvCacheKeyType: kvCacheKeyType === "currentQuant" ? "currentQuant" - : resolveGgmlTypeOption(kvCacheKeyType), + : resolveGgmlTypeOption(kvCacheKeyType, llama), kvCacheValueType: kvCacheValueType === "currentQuant" ? "currentQuant" - : resolveGgmlTypeOption(kvCacheValueType) + : resolveGgmlTypeOption(kvCacheValueType, llama) }); const resolvedModelDestination = resolveModelDestination(ggufPath); @@ -134,7 +134,7 @@ export async function resolveCommandGgufPath(ggufPath: string | undefined, llama fileStats.map(async ({stats, info}) => { if (stats == null) return; - + if (stats.size !== info.totalSize) await fs.remove(info.filePath); }) diff --git a/src/evaluator/LlamaChat/LlamaChat.ts b/src/evaluator/LlamaChat/LlamaChat.ts index affbb324..8fdd0793 100644 --- a/src/evaluator/LlamaChat/LlamaChat.ts +++ b/src/evaluator/LlamaChat/LlamaChat.ts @@ -712,6 +712,22 @@ export class LlamaChat { await loadContextWindow(); generateResponseState.isRerender = false; + if (generateResponseState.isFirstEvaluation && generateResponseState.onModelResponseStateShouldOpenThoughtSegment()) { + if (!tookInitialCheckpoint && this.sequence.needsCheckpoints) { + await generateResponseState.alignCurrentSequenceStateWithCurrentTokens(false); + await generateResponseState.evaluateWithoutGeneratingNewTokens(); + + await this.sequence.takeCheckpoint(); + tookInitialCheckpoint = true; + } + + generateResponseState.openThoughtSegmentOnModelResponseStartIfNeeded(); + generateResponseState.canAvoidReloadingHistory = false; + generateResponseState.isRerender = shouldHandlePrefixTriggers; + await loadContextWindow(); + generateResponseState.isRerender = false; + } + generateResponseState.addStopGenerationTriggersFromChatWrapper(); if (generateResponseState.generatedTokens === 0) { @@ -2018,6 +2034,52 @@ class GenerateResponseState 0) + return false; + + const currentResponseSegmentsStack = SegmentHandler.getStackFromModelResponse(lastModelResponseItem.response); + if (currentResponseSegmentsStack.includes("thought")) + return false; + + if (this.abortOnNonText) + return false; + else + return true; + } + + public openThoughtSegmentOnModelResponseStartIfNeeded() { + if (this.chatWrapper.settings.segments?.thought?.openOnResponseStart !== true) + return false; + + const lastModelResponseItem = this.resolvedHistory.at(-1); + if (lastModelResponseItem == null || lastModelResponseItem.type !== "model") + return false; + + if (lastModelResponseItem.response.length > 0) + return false; + + const currentResponseSegmentsStack = SegmentHandler.getStackFromModelResponse(lastModelResponseItem.response); + if (currentResponseSegmentsStack.includes("thought")) + return false; + + if (this.abortOnNonText) + // we won't force-open a though segment if we are aborting on non-text, + // as it would never allow a textual generation even if the model would choose it otherwise + return false; + else { + this.segmentHandler.openSegment("thought"); + return true; + } + } + public ensureReopenedThoughtSegmentAfterFunctionCallsIfNeeded() { if (this.chatWrapper.settings.segments?.thought?.reopenAfterFunctionCalls !== true) return; @@ -3193,21 +3255,24 @@ class GenerateResponseState = new Set(); private _ownedSegmentsStackLength: number = 0; + private _contextWindowOwnedSegmentsStackLength: number = 0; private readonly _segments: RawSegment[] = []; private readonly _segmentsStartTokenTrail: Token[] = []; private readonly _segmentTokenCounts: Map; @@ -3781,6 +3847,7 @@ class SegmentHandler 0) @@ -4061,6 +4131,9 @@ class SegmentHandler this._ownedSegmentsStackLength, + start: this._segmentsStack.length > this._contextWindowOwnedSegmentsStackLength, startTime: now }); else { @@ -4217,7 +4293,7 @@ class SegmentHandler this._ownedSegmentsStackLength, + start: this._segmentsStack.length > this._contextWindowOwnedSegmentsStackLength, startTime: now }); else diff --git a/src/evaluator/LlamaContext/LlamaContext.ts b/src/evaluator/LlamaContext/LlamaContext.ts index cde97350..6471ee42 100644 --- a/src/evaluator/LlamaContext/LlamaContext.ts +++ b/src/evaluator/LlamaContext/LlamaContext.ts @@ -1,5 +1,5 @@ import path from "path"; -import {acquireLock, AsyncDisposeAggregator, DisposeAggregator, DisposedError, EventRelay, Lock, withLock} from "lifecycle-utils"; +import {acquireLock, AsyncDisposeAggregator, DisposedError, EventRelay, Lock, registerFinalizer, withLock} from "lifecycle-utils"; import {removeNullFields} from "../../utils/removeNullFields.js"; import {Token} from "../../types.js"; import {AddonContext, AddonModelLora, BatchLogitIndex} from "../../bindings/AddonTypes.js"; @@ -87,6 +87,7 @@ export class LlamaContext { /** @internal */ private readonly _disposeAggregator = new AsyncDisposeAggregator(); /** @internal */ private readonly _modelPreventDisposalHandle: DisposalPreventionHandle; /** @internal */ private readonly _loraAdapters = new Set(); + /** @internal */ public readonly _sequenceGcRegistry: FinalizationRegistry; /** @internal */ public _vramConsumptionMarking?: MemoryMarking; /** @internal */ public _ramConsumptionMarking?: MemoryMarking; /** @internal */ private _nextGeneratedSequenceId = 0; @@ -184,17 +185,19 @@ export class LlamaContext { this._reclaimUnusedSequenceId = this._reclaimUnusedSequenceId.bind(this); this._freeReservedThreads = this._freeReservedThreads.bind(this); + this._sequenceGcRegistry = new FinalizationRegistry(this._reclaimUnusedSequenceId); this._disposeAggregator.add(() => { this._disposed = true; }); this._disposeAggregator.add(this._onReclaimUnusedSequenceId); this._disposeAggregator.add(this.onDispose.dispatchEvent); - this._disposeAggregator.add( - this.model.onDispose.createListener( - disposeContextIfReferenced.bind(null, new WeakRef(this)) - ) + + const onModelDisposeListener = this.model.onDispose.createListener( + disposeContextIfReferenced.bind(null, new WeakRef(this)) ); + this._disposeAggregator.add(onModelDisposeListener); + this._disposeAggregator.add(registerFinalizer(this, onModelDisposeListener)); this._disposeAggregator.add(async () => { await this._backendContextDisposeGuard.acquireDisposeLock(); @@ -749,13 +752,17 @@ export class LlamaContext { if (this._disposed) return; - void withLock([this as LlamaContext, "context"], async () => { + return withLock([this as LlamaContext, "context"], async () => { if (this._disposed) return; - this._ctx.disposeSequence(sequenceId); - this._unusedSequenceIds.push(sequenceId); - this._onReclaimUnusedSequenceId.dispatchEvent(); + try { + this._ctx.disposeSequence(sequenceId); + this._unusedSequenceIds.push(sequenceId); + this._onReclaimUnusedSequenceId.dispatchEvent(); + } catch (err) { + this._llama._log(LlamaLogLevel.warn, `Failed to reclaim unused sequence ID ${sequenceId}: ${err}`); + } }); } @@ -901,10 +908,10 @@ export class LlamaContext { : Boolean(flashAttentionOption); const kvCacheKeyType = options.experimentalKvCacheKeyType === "currentQuant" ? _model.fileInsights.dominantTensorType ?? _model.defaultContextKvCacheKeyType - : resolveGgmlTypeOption(options.experimentalKvCacheKeyType) ?? _model.defaultContextKvCacheKeyType; + : resolveGgmlTypeOption(options.experimentalKvCacheKeyType, _model._llama) ?? _model.defaultContextKvCacheKeyType; const kvCacheValueType = options.experimentalKvCacheValueType === "currentQuant" ? _model.fileInsights.dominantTensorType ?? _model.defaultContextKvCacheValueType - : resolveGgmlTypeOption(options.experimentalKvCacheValueType) ?? _model.defaultContextKvCacheValueType; + : resolveGgmlTypeOption(options.experimentalKvCacheValueType, _model._llama) ?? _model.defaultContextKvCacheValueType; const swaFullCache = options.swaFullCache ?? _model.defaultContextSwaFullCache; const loraOptions = typeof options.lora === "string" ? {adapters: [{filePath: options.lora}]} satisfies LlamaContextOptions["lora"] @@ -1067,14 +1074,13 @@ export class LlamaContext { export class LlamaContextSequence { /** @internal */ private readonly _sequenceId: number; - /** @internal */ private readonly _gcRegistry: FinalizationRegistry; /** @internal */ private readonly _context: LlamaContext; /** @internal */ private readonly _contextShift: Required; /** @internal */ private readonly _tokenPredictor?: TokenPredictor; /** @internal */ private readonly _checkpoints = new LlamaContextSequenceCheckpoints(); /** @internal */ private readonly _checkpointOptions: Required; /** @internal */ private readonly _tokenMeter: TokenMeter; - /** @internal */ private readonly _disposeAggregator = new DisposeAggregator(); + /** @internal */ private readonly _disposeAggregator = new AsyncDisposeAggregator({parallel: true}); /** @internal */ private readonly _lock = {}; /** @internal */ private _resetTokenPredictor: boolean = false; /** @internal */ private _tokenPredictorOwner: {} = {}; @@ -1112,34 +1118,34 @@ export class LlamaContextSequence { interval: checkpoints?.interval ?? defaultCheckpointOptions.interval, maxMemory: checkpoints?.maxMemory ?? defaultCheckpointOptions.maxMemory }; - this._gcRegistry = new FinalizationRegistry(this._context._reclaimUnusedSequenceId); - this._gcRegistry.register(this, sequenceId, this); - this._disposeAggregator.add(() => this._gcRegistry.unregister(this)); + this._context._sequenceGcRegistry.register(this, sequenceId, this); + this._disposeAggregator.add(() => void this._context._sequenceGcRegistry.unregister(this)); this._disposeAggregator.add(this.onDispose.dispatchEvent); - this._disposeAggregator.add( - this.model.onDispose.createListener( - disposeContextSequenceIfReferenced.bind(null, new WeakRef(this)) - ) + const onContextDisposeListener = this.context.onDispose.createListener( + disposeContextSequenceIfReferenced.bind(null, new WeakRef(this)) ); - this._disposeAggregator.add(() => { + this._disposeAggregator.add(onContextDisposeListener); + this._disposeAggregator.add(registerFinalizer(this, onContextDisposeListener)); + + this._disposeAggregator.add(async () => { this._checkpoints.clearAllCheckpoints(); - this._context._reclaimUnusedSequenceId(this._sequenceId); + await this._context._reclaimUnusedSequenceId(this._sequenceId); }); if (this._tokenPredictor != null) - this._disposeAggregator.add(this._tokenPredictor); + this._disposeAggregator.add(() => void this._tokenPredictor?.dispose()); this._takeIntervalCheckpointIfNeededAfterBatch = this._takeIntervalCheckpointIfNeededAfterBatch.bind(this); } - public dispose() { + public async dispose() { if (this._disposed) return; - this._disposeAggregator.dispose(); + await this._disposeAggregator.dispose(); this._contextTokens.length = 0; @@ -1147,10 +1153,18 @@ export class LlamaContextSequence { } /** @hidden */ - public [Symbol.dispose]() { + public [Symbol.asyncDispose]() { return this.dispose(); } + /** + * @deprecated Use `[Symbol.asyncDispose]()` instead + * @hidden + */ + public [Symbol.dispose]() { + void this.dispose(); + } + public get disposed() { return this._disposed; } diff --git a/src/evaluator/LlamaModel/LlamaModel.ts b/src/evaluator/LlamaModel/LlamaModel.ts index 32fae3ae..999019d5 100644 --- a/src/evaluator/LlamaModel/LlamaModel.ts +++ b/src/evaluator/LlamaModel/LlamaModel.ts @@ -1,6 +1,6 @@ import process from "process"; import path from "path"; -import {acquireLock, AsyncDisposeAggregator, DisposedError, EventRelay, withLock} from "lifecycle-utils"; +import {acquireLock, AsyncDisposeAggregator, DisposedError, EventRelay, withLock, registerFinalizer} from "lifecycle-utils"; import {removeNullFields} from "../../utils/removeNullFields.js"; import {Token, Tokenizer} from "../../types.js"; import {AddonModel, AddonModelLora, ModelTypeDescription} from "../../bindings/AddonTypes.js"; @@ -73,7 +73,7 @@ export type LlamaModelOptions = { * * When using mmap, you might notice a delay the first time you actually use the model, * which is caused by the OS itself loading the model into memory. - * + * * When this option is set to `"auto"`, mmap may be disabled in scenarios where doing so allows more layers to be offloaded to the GPU. * * Defaults to `"auto"` if the current system supports it. @@ -317,11 +317,11 @@ export class LlamaModel { this._disposedState.disposed = true; }); this._disposeAggregator.add(this.onDispose.dispatchEvent); - this._disposeAggregator.add( - this._llama.onDispose.createListener( - disposeModelIfReferenced.bind(null, new WeakRef(this)) - ) + const onLlamaDisposeListener = this._llama.onDispose.createListener( + disposeModelIfReferenced.bind(null, new WeakRef(this)) ); + this._disposeAggregator.add(onLlamaDisposeListener); + this._disposeAggregator.add(registerFinalizer(this, onLlamaDisposeListener)); this._disposeAggregator.add(async () => { await this._backendModelDisposeGuard.acquireDisposeLock(); @@ -382,6 +382,10 @@ export class LlamaModel { return this._fileInsights; } + public get architecture(): GgufArchitectureType { + return this._fileInfo.metadata?.general?.architecture ?? GgufArchitectureType.unknown; + } + /** * Number of layers offloaded to the GPU. * If GPU support is disabled, this will always be `0`. @@ -392,7 +396,7 @@ export class LlamaModel { /** * Whether the model is loaded using mmap (memory-mapped file) or not. - * + * * When Direct I/O (setting the `useDirectIo` option to `true`) is used it'll override mmap and this value may be out of sync * with the actual usage of mmap for the loading of this model instance. */ @@ -810,11 +814,11 @@ export class LlamaModel { const resolvedDefaultContextSwaFullCache = modelOptions.defaultContextSwaFullCache ?? defaultContextSwaFullCache; const resolvedDefaultContextKvCacheKeyType = experimentalDefaultContextKvCacheKeyType === "currentQuant" ? ggufInsights.dominantTensorType ?? GgmlType.F16 - : resolveGgmlTypeOption(experimentalDefaultContextKvCacheKeyType) ?? GgmlType.F16; + : resolveGgmlTypeOption(experimentalDefaultContextKvCacheKeyType, _llama) ?? GgmlType.F16; const resolvedDefaultContextKvCacheValueType = experimentalDefaultContextKvCacheValueType === "currentQuant" ? ggufInsights.dominantTensorType ?? GgmlType.F16 - : resolveGgmlTypeOption(experimentalDefaultContextKvCacheValueType) ?? GgmlType.F16; - + : resolveGgmlTypeOption(experimentalDefaultContextKvCacheValueType, _llama) ?? GgmlType.F16; + let gpuLayers: number; let resolvedUseMmap: boolean; let resourceRequirementsEstimation: GgufInsightsResourceRequirements; @@ -830,7 +834,7 @@ export class LlamaModel { figuringGpuLayersValueLoadPercentage.percentagePerStepModelMemorySize ) ); - + const layersResolutionStartTime = Date.now(); const layersResolution = await ggufInsights.configurationResolver.resolveModelGpuLayersV2(modelOptions.gpuLayers, { ignoreMemorySafetyChecks: modelOptions.ignoreMemorySafetyChecks, @@ -855,7 +859,7 @@ export class LlamaModel { modelOptions.onLoadProgress?.(layersResolutionLoadedPercentage); }, - + _simulatorSession: simulatorSession }); const layersResolutionEndTime = Date.now(); @@ -872,7 +876,7 @@ export class LlamaModel { resourceRequirementsEstimation = await ggufInsights.estimateModelResourceRequirementsV2({ gpuLayers, useMmap: resolvedUseMmap, - + _simulatorSession: simulatorSession }); } finally { diff --git a/src/gguf/insights/GgufInsights.ts b/src/gguf/insights/GgufInsights.ts index 6fba4522..21902ce6 100644 --- a/src/gguf/insights/GgufInsights.ts +++ b/src/gguf/insights/GgufInsights.ts @@ -1,4 +1,4 @@ -import {acquireLock, withLock} from "lifecycle-utils"; +import {acquireLock, AsyncDisposeAggregator, withLock, registerFinalizer} from "lifecycle-utils"; import bytes from "bytes"; import {Llama} from "../../bindings/Llama.js"; import {doesLlamaBackendNeedAddonInitLock, LlamaLocks, LlamaLogLevel} from "../../bindings/types.js"; @@ -11,6 +11,7 @@ import {getReadablePath} from "../../cli/utils/getReadablePath.js"; import {padSafeContextSize} from "../../evaluator/LlamaContext/utils/padSafeContextSize.js"; import {removeNullFields, removeUndefinedFields} from "../../utils/removeNullFields.js"; import {LruCache} from "../../utils/LruCache.js"; +import {DisposalPreventionHandle, DisposeGuard} from "../../utils/DisposeGuard.js"; import {GgufInsightsConfigurationResolver} from "./GgufInsightsConfigurationResolver.js"; import {GgufInsightsTokens} from "./GgufInsightsTokens.js"; import type {Promisable} from "../../utils/transformPromisable.js"; @@ -111,6 +112,21 @@ export class GgufInsights { return this._modelSize; } + /** The total number of parameters in the model */ + public get totalParameters() { + let totalParameters = 0n; + + for (const tensor of this._ggufFileInfo.fullTensorInfo ?? []) { + let tensorParameters = 1n; + for (const dim of tensor.dimensions ?? []) + tensorParameters *= BigInt(dim); + + totalParameters += tensorParameters; + } + + return Number(totalParameters); + } + public get flashAttentionSupported() { // source: `llama_new_context_with_model` in `llama.cpp` @@ -351,7 +367,7 @@ export class GgufInsights { const simulatorSource = await this._resolveSimulatorSource(); if (simulatorSource == null) return null; - + let resourceRequirements: GgufInsightsResourceRequirements; try { resourceRequirements = await simulatorSession.estimateModelResources({ @@ -362,7 +378,7 @@ export class GgufInsights { } catch (error: any) { throw new Error("Failed simulating model resource usage. Falling back to estimation heuristic. Error: " + (error?.message ?? String(error))); } - + this._exactModelResourceRequirementsCache.set(cacheKey, resourceRequirements); return {...resourceRequirements}; } finally { @@ -827,7 +843,7 @@ export class GgufInsights { const simulatorSource = await this._resolveSimulatorSource(); if (simulatorSource == null) return null; - + let contextResources: GgufInsightsResourceRequirements; try { const paddedContextSize = padSafeContextSize(contextSize, "up"); @@ -861,7 +877,7 @@ export class GgufInsights { cpuRam: contextResources.cpuRam, gpuVram: contextResources.gpuVram } satisfies GgufInsightsResourceRequirements; - + this._exactContextResourceRequirementsCache.set(cacheKey, resourceRequirements); return {...resourceRequirements}; } finally { @@ -1186,12 +1202,22 @@ export class GgufInsights { export class GgufInsightsSimulatorSession { private readonly _llama: Llama; - private readonly _modelPromises: LruCache>; + private readonly _modelHandlePromises: LruCache>; private _disposed = false; public constructor(llama: Llama, lruCacheSize: number = 10) { this._llama = llama; - this._modelPromises = new LruCache(lruCacheSize); + this._modelHandlePromises = new LruCache(lruCacheSize, { + async onDelete(key, value) { + try { + const modelHandle = await value; + await Promise.resolve(); // wait for a tick to allow any pending consumers to acquire dispose prevention handles + await modelHandle.dispose(); + } catch (err) { + // do nothing + } + } + }); } public async estimateModelResources({ @@ -1203,16 +1229,28 @@ export class GgufInsightsSimulatorSession { gpuLayers: number, useMmap?: boolean }): Promise { - const model = await this._getModel({source: modelSource, gpuLayers, useMmap}); - const memoryBreakdown = model.getMemoryBreakdown(); - if (this._llama._shouldLog(LlamaLogLevel.debug)) - this._llama._log(LlamaLogLevel.debug, "Simulating model resource usage. " + [ - `gpuLayers=${gpuLayers}`, - `useMmap=${useMmap}`, - `memoryBreakdownCpuRam=${bytes(memoryBreakdown.cpuRam)}`, - `memoryBreakdownGpuVram=${bytes(memoryBreakdown.gpuVram)}` - ].join(" ")); - return memoryBreakdown; + const modelHandle = await this._getModelHandle({source: modelSource, gpuLayers, useMmap}); + + let preventDisposalHandle: DisposalPreventionHandle; + try { + preventDisposalHandle = modelHandle.disposeGuard.createPreventDisposalHandle(); + } catch (err) { + throw new Error("Model is disposed"); + } + + try { + const memoryBreakdown = modelHandle.model.getMemoryBreakdown(); + if (this._llama._shouldLog(LlamaLogLevel.debug)) + this._llama._log(LlamaLogLevel.debug, "Simulating model resource usage. " + [ + `gpuLayers=${gpuLayers}`, + `useMmap=${useMmap}`, + `memoryBreakdownCpuRam=${bytes(memoryBreakdown.cpuRam)}`, + `memoryBreakdownGpuVram=${bytes(memoryBreakdown.gpuVram)}` + ].join(" ")); + return memoryBreakdown; + } finally { + preventDisposalHandle.dispose(); + } } public async estimateContextResources({ @@ -1240,53 +1278,65 @@ export class GgufInsightsSimulatorSession { kvCacheKeyType?: GgmlType, kvCacheValueType?: GgmlType }): Promise { - const model = await this._getModel({source: modelSource, gpuLayers, useMmap}); - const context = new this._llama._bindings.AddonContext(model, removeUndefinedFields({ - contextSize, - batchSize, - sequences, - embeddings: isEmbeddingContext, - flashAttention: flashAttention === "auto" - ? "auto" - : flashAttention, - kvCacheKeyType, - kvCacheValueType, - swaFullCache - } satisfies AddonContextParams)); + const modelHandle = await this._getModelHandle({source: modelSource, gpuLayers, useMmap}); + + let preventDisposalHandle: DisposalPreventionHandle; + try { + preventDisposalHandle = modelHandle.disposeGuard.createPreventDisposalHandle(); + } catch (err) { + throw new Error("Model is disposed"); + } try { - const loadingLock = doesLlamaBackendNeedAddonInitLock(this._llama.gpu) - ? await acquireLock([this._llama._memoryLock, LlamaLocks.addonInit]) - : undefined; - const disposeLogLevelOverride = this._llama._createLogLevelOverride(LlamaLogLevel.error); + const context = new this._llama._bindings.AddonContext(modelHandle.model, removeUndefinedFields({ + contextSize, + batchSize, + sequences, + embeddings: isEmbeddingContext, + flashAttention: flashAttention === "auto" + ? "auto" + : flashAttention, + kvCacheKeyType, + kvCacheValueType, + swaFullCache + } satisfies AddonContextParams)); + try { - const contextLoaded = await context.init(); - if (!contextLoaded) - throw new Error("Failed to create context"); + const loadingLock = doesLlamaBackendNeedAddonInitLock(this._llama.gpu) + ? await acquireLock([this._llama._memoryLock, LlamaLocks.addonInit]) + : undefined; + const disposeLogLevelOverride = this._llama._createLogLevelOverride(LlamaLogLevel.error); + try { + const contextLoaded = await context.init(); + if (!contextLoaded) + throw new Error("Failed to create context"); + } finally { + disposeLogLevelOverride(); + loadingLock?.dispose(); + } + + const memoryBreakdown = context.getMemoryBreakdown(); + if (this._llama._shouldLog(LlamaLogLevel.debug)) + this._llama._log(LlamaLogLevel.debug, "Simulating context resource usage. " + [ + `gpuLayers=${gpuLayers}`, + `contextSize=${contextSize.toLocaleString("en-US", {notation: "compact"})}`, + `batchSize=${batchSize}`, + `sequences=${sequences}`, + `isEmbeddingContext=${isEmbeddingContext}`, + `flashAttention=${flashAttention}`, + `swaFullCache=${swaFullCache}`, + `kvCacheKeyType=${kvCacheKeyType}`, + `kvCacheValueType=${kvCacheValueType}`, + `useMmap=${useMmap}`, + `memoryBreakdownCpuRam=${bytes(memoryBreakdown.cpuRam)}`, + `memoryBreakdownGpuVram=${bytes(memoryBreakdown.gpuVram)}` + ].join(" ")); + return memoryBreakdown; } finally { - disposeLogLevelOverride(); - loadingLock?.dispose(); + await context.dispose(); } - - const memoryBreakdown = context.getMemoryBreakdown(); - if (this._llama._shouldLog(LlamaLogLevel.debug)) - this._llama._log(LlamaLogLevel.debug, "Simulating context resource usage. " + [ - `gpuLayers=${gpuLayers}`, - `contextSize=${contextSize.toLocaleString("en-US", {notation: "compact"})}`, - `batchSize=${batchSize}`, - `sequences=${sequences}`, - `isEmbeddingContext=${isEmbeddingContext}`, - `flashAttention=${flashAttention}`, - `swaFullCache=${swaFullCache}`, - `kvCacheKeyType=${kvCacheKeyType}`, - `kvCacheValueType=${kvCacheValueType}`, - `useMmap=${useMmap}`, - `memoryBreakdownCpuRam=${bytes(memoryBreakdown.cpuRam)}`, - `memoryBreakdownGpuVram=${bytes(memoryBreakdown.gpuVram)}` - ].join(" ")); - return memoryBreakdown; } finally { - await context.dispose(); + preventDisposalHandle.dispose(); } } @@ -1300,18 +1350,19 @@ export class GgufInsightsSimulatorSession { this._disposed = true; - const modelPromises = [...this._modelPromises.values()].map((modelPromise) => modelPromise.catch(() => void 0)); - this._modelPromises.clear(); - const loadedModels = (await Promise.all(modelPromises)).filter((model) => model != null); + const modelHandlePromises = [...this._modelHandlePromises.values()] + .map((modelHandlePromise) => modelHandlePromise.catch(() => void 0)); + this._modelHandlePromises.clear(); + const loadedModelHandles = (await Promise.all(modelHandlePromises)).filter((model) => model != null); - await Promise.all(loadedModels.map((model) => model.dispose().catch(() => void 0))); + await Promise.all(loadedModelHandles.map((modelHandle) => modelHandle.dispose())); } public get disposed() { return this._disposed; } - private async _getModel({ + private async _getModelHandle({ source, gpuLayers, useMmap = this._llama.supportsMmap @@ -1323,25 +1374,36 @@ export class GgufInsightsSimulatorSession { if (this._disposed) throw new Error("simulator session is disposed"); - const cacheKey = String(gpuLayers) + ":" + String(useMmap); - const existingModelPromise = this._modelPromises.get(cacheKey); - if (existingModelPromise != null) - return await existingModelPromise; - - if (this._llama._shouldLog(LlamaLogLevel.debug)) - this._llama._log(LlamaLogLevel.debug, `Loading model for simulator session. gpuLayers=${gpuLayers} useMmap=${useMmap}`); - const modelPromise = this._loadModel({ - source, - gpuLayers, - useMmap - }); - this._modelPromises.set(cacheKey, modelPromise); + let preventDisposalHandle: DisposalPreventionHandle; + try { + preventDisposalHandle = this._llama._backendDisposeGuard.createPreventDisposalHandle(); + } catch (err) { + throw new Error("Llama instance is disposed"); + } try { - return await modelPromise; - } catch (error) { - this._modelPromises.delete(cacheKey); - throw error; + const cacheKey = String(gpuLayers) + ":" + String(useMmap); + const existingModelPromise = this._modelHandlePromises.get(cacheKey); + if (existingModelPromise != null) + return await existingModelPromise; + + if (this._llama._shouldLog(LlamaLogLevel.debug)) + this._llama._log(LlamaLogLevel.debug, `Loading model for simulator session. gpuLayers=${gpuLayers} useMmap=${useMmap}`); + const modelHandlePromise = this._loadModel({ + source, + gpuLayers, + useMmap + }); + this._modelHandlePromises.set(cacheKey, modelHandlePromise); + + try { + return await modelHandlePromise; + } catch (error) { + this._modelHandlePromises.delete(cacheKey); + throw error; + } + } finally { + preventDisposalHandle.dispose(); } } @@ -1377,10 +1439,51 @@ export class GgufInsightsSimulatorSession { loadingLock?.dispose(); } - return model; + return new SimulatorModelHandle(this._llama, model); + } +} + +class SimulatorModelHandle { + public readonly model: AddonModel; + public readonly disposeGuard: DisposeGuard; + + private readonly _llamaPreventDisposalHandle: DisposalPreventionHandle; + private readonly _disposeAggregator = new AsyncDisposeAggregator(); + + public constructor(llama: Llama, model: AddonModel) { + this.model = model; + this.disposeGuard = new DisposeGuard([llama._backendDisposeGuard]); + + this._llamaPreventDisposalHandle = llama._backendDisposeGuard.createPreventDisposalHandle(); + this._disposeAggregator.add(registerFinalizer(model, this._llamaPreventDisposalHandle)); + + const onLlamaDisposeListener = llama.onDispose.createListener( + disposeSimulatorModelHandleIfReferenced.bind(null, new WeakRef(this)) + ); + this._disposeAggregator.add(onLlamaDisposeListener); + this._disposeAggregator.add(registerFinalizer(model, onLlamaDisposeListener)); + + this._disposeAggregator.add(this._dispose.bind(this)); + } + + public async dispose() { + await this._disposeAggregator.dispose(); + } + + private async _dispose() { + await this.disposeGuard.acquireDisposeLock(); + + await this.model.dispose().catch(() => void 0); + + await this._llamaPreventDisposalHandle.dispose(); } } +function disposeSimulatorModelHandleIfReferenced(modelHandleRef: WeakRef) { + return modelHandleRef.deref()?.dispose() + .catch(() => void 0); +} + function parseTensorName(tensorName?: string): { layerNumber: number | undefined } { diff --git a/src/gguf/types/GgufMetadataTypes.ts b/src/gguf/types/GgufMetadataTypes.ts index 3d53e764..1faae49c 100644 --- a/src/gguf/types/GgufMetadataTypes.ts +++ b/src/gguf/types/GgufMetadataTypes.ts @@ -198,7 +198,8 @@ export enum GgufFileType { MOSTLY_TQ2_0 = 37, MOSTLY_MXFP4_MOE = 38, MOSTLY_NVFP4 = 39, - MOSTLY_Q1_0 = 40 + MOSTLY_Q1_0 = 40, + MOSTLY_Q2_0 = 41 } @@ -238,6 +239,11 @@ export type GgufMetadataGeneral([ ["Q1_0", GgufFileType.MOSTLY_Q1_0], + ["Q2_0", GgufFileType.MOSTLY_Q2_0], ["Q4_0", GgufFileType.MOSTLY_Q4_0], ["Q4_1", GgufFileType.MOSTLY_Q4_1], ["MXFP4_MOE", GgufFileType.MOSTLY_MXFP4_MOE], @@ -41,3 +42,14 @@ export const ggufQuantNames = new Map([ ["F32", GgufFileType.ALL_F32], ["COPY", GgufFileType.ALL_F32] ]); + +export const ggufFileQuantNames = Object.freeze([ + ...ggufQuantNames.keys(), + "Q2_K_XL", + "Q3_K_XL", + "Q4_K_XL", + "Q5_K_XL", + "Q6_K_XL", + "Q7_K_XL", + "Q8_K_XL" +].filter((name) => name !== "COPY")); diff --git a/src/index.ts b/src/index.ts index c70f7175..7dad0945 100644 --- a/src/index.ts +++ b/src/index.ts @@ -97,6 +97,7 @@ import { import {type ModelDownloadEndpoints} from "./utils/modelDownloadEndpoints.js"; import {jsonDumps} from "./chatWrappers/utils/jsonDumps.js"; import {experimentalChunkDocument} from "./evaluator/utils/chunkDocument.js"; +import {ggufFileQuantNames} from "./gguf/utils/ggufQuantNames.js"; import { type ChatHistoryItem, type ChatModelFunctionCall, type ChatModelSegmentType, type ChatModelSegment, type ChatModelFunctions, @@ -346,5 +347,6 @@ export { type CombinedModelDownloaderOptions, jsonDumps, type OverridesObject, - experimentalChunkDocument + experimentalChunkDocument, + ggufFileQuantNames }; diff --git a/src/types.ts b/src/types.ts index 630da6c1..7034a5b6 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1,6 +1,7 @@ import {GbnfJsonSchema, GbnfJsonSchemaToType} from "./utils/gbnfJson/types.js"; import {LlamaText, BuiltinSpecialTokenValue, LlamaTextJSON} from "./utils/LlamaText.js"; import type {GgufFileInfo} from "./gguf/types/GgufFileInfoTypes.js"; +import type {GgufArchitectureType} from "./gguf/types/GgufMetadataTypes.js"; export type Token = number & { __token: never @@ -111,6 +112,7 @@ export type ChatWrapperSettings = { /** Chain of Thought text segment */ readonly thought?: ChatWrapperSettingsSegment & { + openOnResponseStart?: boolean, reopenAfterFunctionCalls?: boolean }, @@ -135,7 +137,8 @@ export type ChatWrapperGenerateContextStateOptions = { export type ChatWrapperCheckModelCompatibilityParams = { tokenizer?: Tokenizer, - fileInfo?: GgufFileInfo + fileInfo?: GgufFileInfo, + architecture?: GgufArchitectureType }; export type ChatWrapperGeneratedContextState = diff --git a/src/utils/LruCache.ts b/src/utils/LruCache.ts index 7f44cd17..894f7024 100644 --- a/src/utils/LruCache.ts +++ b/src/utils/LruCache.ts @@ -57,6 +57,6 @@ export class LruCache { } public delete(key: Key) { - this._cache.delete(key); + return this._cache.delete(key); } } diff --git a/src/utils/OpenAIFormat.ts b/src/utils/OpenAIFormat.ts index c1a839ab..c3d7426c 100644 --- a/src/utils/OpenAIFormat.ts +++ b/src/utils/OpenAIFormat.ts @@ -608,7 +608,8 @@ export type OpenAiChatAssistantMessage = { name: string, arguments: string } - }> + }>, + "reasoning_content"?: string }; export type OpenAiChatToolMessage = { role: "tool", diff --git a/src/utils/ThreadsSplitter.ts b/src/utils/ThreadsSplitter.ts index 7f08e29d..7b8b1c03 100644 --- a/src/utils/ThreadsSplitter.ts +++ b/src/utils/ThreadsSplitter.ts @@ -8,6 +8,9 @@ export class ThreadsSplitter { private _totalWantedThreads: number = 0; public maxThreads: number; + /** @internal */ public readonly _wantedThreadsGcRegistry: FinalizationRegistry; + /** @internal */ public readonly _demandedThreadsGcRegistry: FinalizationRegistry; + /** * Set to `0` to disable the limit * @param maxThreads @@ -17,6 +20,9 @@ export class ThreadsSplitter { this._removeWantedThreads = this._removeWantedThreads.bind(this); this._removeThreadDemand = this._removeThreadDemand.bind(this); + + this._wantedThreadsGcRegistry = new FinalizationRegistry(this._removeWantedThreads); + this._demandedThreadsGcRegistry = new FinalizationRegistry(this._removeThreadDemand); } public createConsumer(wantedThreads: number, minThreads: number = 1) { @@ -126,8 +132,6 @@ export class ThreadsSplitterConsumer { private readonly _threadsSplitter: ThreadsSplitter; private readonly _wantedThreads: number; private readonly _demandedThreads: number; - private readonly _wantedThreadsGcRegistry: FinalizationRegistry; - private readonly _demandedThreadsGcRegistry: FinalizationRegistry; private _usedThreads: number = 0; private _disposed: boolean = false; @@ -139,11 +143,8 @@ export class ThreadsSplitterConsumer { this._threadsSplitter._addWantedThreads(this._wantedThreads); this._threadsSplitter._addThreadDemand(this._demandedThreads); - this._wantedThreadsGcRegistry = new FinalizationRegistry(this._threadsSplitter._removeWantedThreads); - this._wantedThreadsGcRegistry.register(this, this._wantedThreads, this); - - this._demandedThreadsGcRegistry = new FinalizationRegistry(this._threadsSplitter._removeThreadDemand); - this._demandedThreadsGcRegistry.register(this, this._demandedThreads, this); + this._threadsSplitter._wantedThreadsGcRegistry.register(this, this._wantedThreads, this); + this._threadsSplitter._demandedThreadsGcRegistry.register(this, this._demandedThreads, this); } public [Symbol.dispose]() { @@ -159,8 +160,8 @@ export class ThreadsSplitterConsumer { this._threadsSplitter._removeWantedThreads(this._wantedThreads); this._threadsSplitter._removeThreadDemand(this._demandedThreads); - this._wantedThreadsGcRegistry.unregister(this); - this._demandedThreadsGcRegistry.unregister(this); + this._threadsSplitter._wantedThreadsGcRegistry.unregister(this); + this._threadsSplitter._demandedThreadsGcRegistry.unregister(this); } public getAllocationToConsume(): Promisable<[threadsToUse: number, usageHandle: DisposableHandle]> { diff --git a/src/utils/createModelDownloader.ts b/src/utils/createModelDownloader.ts index 1c48018d..50e7ede8 100644 --- a/src/utils/createModelDownloader.ts +++ b/src/utils/createModelDownloader.ts @@ -1,6 +1,6 @@ import process from "process"; import path from "path"; -import {DownloadEngineMultiDownload, DownloadEngineNodejs, downloadFile, downloadSequence} from "ipull"; +import {DownloadEngineMultiDownload, DownloadEngineNodejs, downloadFile, downloadSequence, DownloadStatus} from "ipull"; import fs from "fs-extra"; import chalk from "chalk"; import {createSplitPartFilename, resolveSplitGgufParts} from "../gguf/utils/resolveSplitGgufParts.js"; @@ -45,7 +45,7 @@ export type ModelDownloaderOptions = ({ */ showCliProgress?: boolean, - onProgress?: (status: {totalSize: number, downloadedSize: number}) => void, + onProgress?: (status: {totalSize: number, downloadedSize: number, estimatedTimeLeft: number, averageSpeed: number}) => void, /** * If true, the downloader will skip the download if the file already exists, and its size matches the size of the remote file. @@ -294,6 +294,31 @@ export class ModelDownloader { .reduce((acc, transferredBytes) => acc + transferredBytes, 0); } + public get estimatedTimeLeft() { + let maxTimeLeft: number | undefined = undefined; + for (const downloader of this._specificFileDownloaders) { + const timeLeft = downloader.status.timeLeft; + if (timeLeft == null) + continue; + + if (maxTimeLeft == null || timeLeft > maxTimeLeft) + maxTimeLeft = timeLeft; + } + + return maxTimeLeft ?? Infinity; + } + + public get averageSpeed() { + const speeds = this._specificFileDownloaders + .filter((downloader) => downloader.status.downloadStatus === DownloadStatus.Active) + .map((downloader) => downloader.status.speed); + + if (speeds.length === 0) + return 0; + + return speeds.reduce((res, speed) => res + speed, 0) / speeds.length; + } + /** * Info about all the files that will be saved to the download directory, * including their filenames, full paths, total sizes and downloaded sizes. @@ -368,7 +393,9 @@ export class ModelDownloader { private _onDownloadProgress() { this._onProgress?.({ totalSize: this.totalSize, - downloadedSize: this.downloadedSize + downloadedSize: this.downloadedSize, + estimatedTimeLeft: this.estimatedTimeLeft, + averageSpeed: this.averageSpeed }); } diff --git a/src/utils/getTempDir.ts b/src/utils/getTempDir.ts index 5f93121d..aee63b4f 100644 --- a/src/utils/getTempDir.ts +++ b/src/utils/getTempDir.ts @@ -101,18 +101,17 @@ function onExit() { } } +const fsHandleFinalizationRegistry = new FinalizationRegistry(removePathUsageSync); export class FsPathHandle { public readonly path: string; - private _finalizationRegistry: FinalizationRegistry; private _disposed: boolean = false; public constructor(dirPath: string) { this.path = dirPath; - this._finalizationRegistry = new FinalizationRegistry(removePathUsageSync); addPathUsage(this.path); - this._finalizationRegistry.register(this, this.path, this); + fsHandleFinalizationRegistry.register(this, this.path, this); } public async dispose() { @@ -120,7 +119,7 @@ export class FsPathHandle { return; this._disposed = true; - this._finalizationRegistry.unregister(this); + fsHandleFinalizationRegistry.unregister(this); await removePathUsage(this.path, true); } @@ -133,7 +132,7 @@ export class FsPathHandle { return; this._disposed = true; - this._finalizationRegistry.unregister(this); + fsHandleFinalizationRegistry.unregister(this); removePathUsage(this.path, false); } } diff --git a/src/utils/getTypeScriptTypeStringForGbnfJsonSchema.ts b/src/utils/getTypeScriptTypeStringForGbnfJsonSchema.ts index 5cb0c0c5..c4b6fff3 100644 --- a/src/utils/getTypeScriptTypeStringForGbnfJsonSchema.ts +++ b/src/utils/getTypeScriptTypeStringForGbnfJsonSchema.ts @@ -143,7 +143,7 @@ function _getTypeScriptTypeStringForGbnfJsonSchema( "\n ", valueTypes .map((value) => value.split("\n").join("\n ")) - .join(",\n ") + .join(",\n") .trimStart(), "\n" ].join("") diff --git a/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts b/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts index 3a76cc82..69c8934f 100644 --- a/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts +++ b/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts @@ -1595,7 +1595,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/qwen3.5-0.8b/checkpoints.test.ts b/test/modelDependent/qwen3.5-0.8b/checkpoints.test.ts index 62b58c95..ae33c5d2 100644 --- a/test/modelDependent/qwen3.5-0.8b/checkpoints.test.ts +++ b/test/modelDependent/qwen3.5-0.8b/checkpoints.test.ts @@ -51,26 +51,29 @@ describe("qwen3.5 0.8b", () => { The second word is "secret"." `); - expect(chatSession.sequence.tokenMeter.usedInputTokens).toMatchInlineSnapshot("389"); - expect(chatSession.sequence.lastCheckpointIndex).toMatchInlineSnapshot("425"); - expect(chatSession.sequence.nextTokenIndex).toMatchInlineSnapshot("437"); + expect(chatSession.sequence.tokenMeter.usedInputTokens).toMatchInlineSnapshot("391"); + expect(chatSession.sequence.lastCheckpointIndex).toMatchInlineSnapshot("550"); + expect(chatSession.sequence.nextTokenIndex).toMatchInlineSnapshot("585"); const initialMeterState = chatSession.sequence.tokenMeter.getState(); const res2 = await chatSession.prompt("Explain what this word means. short", { ...promptOptions, - maxTokens: 15 + maxTokens: 12, + budgets: { + thoughtTokens: 4 + } }); const diffMeterState = chatSession.sequence.tokenMeter.diff(initialMeterState); - expect(res2).to.toMatchInlineSnapshot(` + expect(res2.toLowerCase()).to.toMatchInlineSnapshot(` " - "Secret" means something that is hidden or not known" + "secret" means something hidden or" `); - expect(diffMeterState.usedInputTokens).toMatchInlineSnapshot("90"); + expect(diffMeterState.usedInputTokens).toMatchInlineSnapshot("94"); expect(diffMeterState.usedInputTokens).to.be.lessThanOrEqual(95); expect(chatSession.sequence.lastCheckpointIndex).toMatchInlineSnapshot("448"); - expect(chatSession.sequence.nextTokenIndex).toMatchInlineSnapshot("463"); + expect(chatSession.sequence.nextTokenIndex).toMatchInlineSnapshot("464"); }); test("disposing the context asynchronously works", {timeout: 1000 * 60 * 60 * 2}, async () => { diff --git a/test/standalone/chatWrappers/FunctionaryChatWrapper.test.ts b/test/standalone/chatWrappers/FunctionaryChatWrapper.test.ts index a330babe..a562aeab 100644 --- a/test/standalone/chatWrappers/FunctionaryChatWrapper.test.ts +++ b/test/standalone/chatWrappers/FunctionaryChatWrapper.test.ts @@ -243,7 +243,7 @@ describe("FunctionaryChatWrapper", () => { type notifyOwner2 = (_: /* Type: notification */ { // Notification message message: string, - + // Sub notifications subNotifications: (/* notification type */ any)[] }) => any; @@ -253,15 +253,15 @@ describe("FunctionaryChatWrapper", () => { // Some message // minimum length: 3, maximum length: 10 message: string, - + // Some words // maximum items: 5 words: [string, string, ...string[]], - + // Some headers // minimum number of properties: 4, maximum number of properties: 12 headers: {[key: string]: string}, - + // Some mappings // minimum number of properties: 4, maximum number of properties: 12 mappings: {a: boolean, b: number, c: string | null} & {[key: string]: string} @@ -452,7 +452,7 @@ describe("FunctionaryChatWrapper", () => { type notifyOwner2 = (_: /* Type: notification */ { // Notification message message: string, - + // Sub notifications subNotifications: (/* notification type */ any)[] }) => any; @@ -462,15 +462,15 @@ describe("FunctionaryChatWrapper", () => { // Some message // minimum length: 3, maximum length: 10 message: string, - + // Some words // maximum items: 5 words: [string, string, ...string[]], - + // Some headers // minimum number of properties: 4, maximum number of properties: 12 headers: {[key: string]: string}, - + // Some mappings // minimum number of properties: 4, maximum number of properties: 12 mappings: {a: boolean, b: number, c: string | null} & {[key: string]: string} diff --git a/test/standalone/chatWrappers/Gemma4ChatWrapper.test.ts b/test/standalone/chatWrappers/Gemma4ChatWrapper.test.ts index 9aeb8e5f..6ea1c4a8 100644 --- a/test/standalone/chatWrappers/Gemma4ChatWrapper.test.ts +++ b/test/standalone/chatWrappers/Gemma4ChatWrapper.test.ts @@ -216,7 +216,8 @@ describe("Gemma4ChatWrapper", () => { LlamaText([ new SpecialToken("BOS"), new SpecialTokensText("<|turn>system - <|think|>"), + <|think|> + "), "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible. If a question does not make any sense, or is not factually coherent, explain why instead of answering something incorrectly. If you don't know the answer to a question, don't share false information.", new SpecialTokensText(" @@ -252,7 +253,8 @@ describe("Gemma4ChatWrapper", () => { LlamaText([ new SpecialToken("BOS"), new SpecialTokensText("<|turn>system - <|think|>"), + <|think|> + "), "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible. If a question does not make any sense, or is not factually coherent, explain why instead of answering something incorrectly. If you don't know the answer to a question, don't share false information.", new SpecialTokensText("<|tool>"), @@ -306,7 +308,8 @@ describe("Gemma4ChatWrapper", () => { LlamaText([ new SpecialToken("BOS"), new SpecialTokensText("<|turn>system - <|think|>"), + <|think|> + "), "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible. If a question does not make any sense, or is not factually coherent, explain why instead of answering something incorrectly. If you don't know the answer to a question, don't share false information.", new SpecialTokensText(" @@ -334,7 +337,8 @@ describe("Gemma4ChatWrapper", () => { LlamaText([ new SpecialToken("BOS"), new SpecialTokensText("<|turn>system - <|think|>"), + <|think|> + "), "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible. If a question does not make any sense, or is not factually coherent, explain why instead of answering something incorrectly. If you don't know the answer to a question, don't share false information.", new SpecialTokensText(" diff --git a/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts b/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts index 0658e6af..83d5b88a 100644 --- a/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts +++ b/test/standalone/chatWrappers/generic/JinjaTemplateChatWrapper.test.ts @@ -2,7 +2,7 @@ import {describe, expect, test} from "vitest"; import {Template} from "@huggingface/jinja"; import {ChatHistoryItem, ChatModelFunctions, JinjaTemplateChatWrapper} from "../../../../src/index.js"; import {defaultChatSystemPrompt} from "../../../../src/config.js"; -import {LlamaText} from "../../../../src/utils/LlamaText.js"; +import {LlamaText, SpecialTokensText} from "../../../../src/utils/LlamaText.js"; import {fromChatHistoryToIntermediateOpenAiMessages, fromIntermediateToCompleteOpenAiMessages} from "../../../../src/utils/OpenAIFormat.js"; import {removeUndefinedFields} from "../../../../src/utils/removeNullFields.js"; @@ -200,6 +200,167 @@ const llama3_1ChatJinjaTemplate = ` {%- endif %} `.slice(1, -1); +const qwen3_6Template = ` +{%- set image_count = namespace(value=0) %} +{%- set video_count = namespace(value=0) %} +{%- macro render_content(content, do_vision_count, is_system_content=false) %} + {%- if content is string %} + {{- content }} + {%- elif content is iterable and content is not mapping %} + {%- for item in content %} + {%- if 'image' in item or 'image_url' in item or item.type == 'image' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain images.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set image_count.value = image_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Picture ' ~ image_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|image_pad|><|vision_end|>' }} + {%- elif 'video' in item or item.type == 'video' %} + {%- if is_system_content %} + {{- raise_exception('System message cannot contain videos.') }} + {%- endif %} + {%- if do_vision_count %} + {%- set video_count.value = video_count.value + 1 %} + {%- endif %} + {%- if add_vision_id %} + {{- 'Video ' ~ video_count.value ~ ': ' }} + {%- endif %} + {{- '<|vision_start|><|video_pad|><|vision_end|>' }} + {%- elif 'text' in item %} + {{- item.text }} + {%- else %} + {{- raise_exception('Unexpected item type in content.') }} + {%- endif %} + {%- endfor %} + {%- elif content is none or content is undefined %} + {{- '' }} + {%- else %} + {{- raise_exception('Unexpected content type.') }} + {%- endif %} +{%- endmacro %} +{%- if not messages %} + {{- raise_exception('No messages provided.') }} +{%- endif %} +{%- set num_sys = 0 %} +{%- set merged_system = '' %} +{%- if messages[0].role == 'system' or messages[0].role == 'developer' %} + {%- set first = render_content(messages[0].content, false, true)|trim %} + {%- if messages|length > 1 and (messages[1].role == 'system' or messages[1].role == 'developer') %} + {%- set second = render_content(messages[1].content, false, true)|trim %} + {%- set merged_system = first + '\n' + second %} + {%- set num_sys = 2 %} + {%- else %} + {%- set merged_system = first %} + {%- set num_sys = 1 %} + {%- endif %} +{%- endif %} +{%- if tools and tools is iterable and tools is not mapping %} + {{- '<|im_start|>system\n' }} + {{- "# Tools\n\nYou have access to the following functions:\n\n" }} + {%- for tool in tools %} + {{- "\n" }} + {{- tool | tojson }} + {%- endfor %} + {{- "\n" }} + {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n\n\n\nvalue_1\n\n\nThis is the value for the second parameter\nthat can span\nmultiple lines\n\n\n\n\n\nReminder:\n- Function calls MUST follow the specified format: an inner block must be nested within XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n' }} + {%- if merged_system %} + {{- '\n\n' + merged_system }} + {%- endif %} + {{- '<|im_end|>\n' }} +{%- else %} + {%- if merged_system %} + {{- '<|im_start|>system\n' + merged_system + '<|im_end|>\n' }} + {%- endif %} +{%- endif %} +{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %} +{%- for message in messages[::-1] %} + {%- set index = (messages|length - 1) - loop.index0 %} + {%- if ns.multi_step_tool and message.role == "user" %} + {%- set content = render_content(message.content, false)|trim %} + {%- if not(content.startswith('') and content.endswith('')) %} + {%- set ns.multi_step_tool = false %} + {%- set ns.last_query_index = index %} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- for message in messages %} + {%- if loop.index0 >= num_sys and message.role != "system" and message.role != "developer" %} + {%- set content = render_content(message.content, true)|trim %} + {%- if message.role == "user" %} + {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }} + {%- elif message.role == "assistant" %} + {%- set reasoning_content = '' %} + {%- if message.reasoning_content is string %} + {%- set reasoning_content = message.reasoning_content %} + {%- else %} + {%- if '' in content %} + {%- set reasoning_content = content.split('')[0].rstrip('\n').split('')[-1].lstrip('\n') %} + {%- set content = content.split('')[-1].lstrip('\n') %} + {%- endif %} + {%- endif %} + {%- set reasoning_content = reasoning_content|trim %} + {%- if (preserve_thinking is defined and preserve_thinking is true) or (loop.index0 > ns.last_query_index) %} + {{- '<|im_start|>' + message.role + '\n\n' + reasoning_content + '\n\n\n' + content }} + {%- else %} + {{- '<|im_start|>' + message.role + '\n' + content }} + {%- endif %} + {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %} + {%- for tool_call in message.tool_calls %} + {%- if tool_call.function is defined %} + {%- set tool_call = tool_call.function %} + {%- endif %} + {%- if loop.first %} + {%- if content|trim %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n\n' }} + {%- endif %} + {%- else %} + {{- '\n\n\n' }} + {%- endif %} + {%- if tool_call.arguments is mapping %} + {%- for args_name in tool_call.arguments %} + {%- set args_value = tool_call.arguments[args_name] %} + {{- '\n' }} + {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %} + {{- args_value }} + {{- '\n\n' }} + {%- endfor %} + {%- endif %} + {{- '\n' }} + {%- endfor %} + {%- endif %} + {{- '<|im_end|>\n' }} + {%- elif message.role == "tool" %} + {%- if loop.previtem and loop.previtem.role != "tool" %} + {{- '<|im_start|>user' }} + {%- endif %} + {{- '\n\n' }} + {{- content }} + {{- '\n' }} + {%- if not loop.last and loop.nextitem.role != "tool" %} + {{- '<|im_end|>\n' }} + {%- elif loop.last %} + {{- '<|im_end|>\n' }} + {%- endif %} + {%- endif %} + {%- endif %} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} + {%- if enable_thinking is defined and enable_thinking is false %} + {{- '\n\n\n\n' }} + {%- else %} + {{- '\n' }} + {%- endif %} +{%- endif %} +{#- Unsloth fixes - developer role, tool calling #} +`.slice(1, -1); + describe("JinjaTemplateChatWrapper", () => { const template1 = "{{ bos_token }}" + @@ -804,10 +965,10 @@ describe("JinjaTemplateChatWrapper", () => { function func8(params: { // The main message message: string, - + // The feeling feeling: "good" | "bad", - + // The number of words. // For example, 6 words: number @@ -883,10 +1044,10 @@ describe("JinjaTemplateChatWrapper", () => { function func8(params: { // The main message message: string, - + // The feeling feeling: "good" | "bad", - + // The number of words. // For example, 6 words: number @@ -961,10 +1122,10 @@ describe("JinjaTemplateChatWrapper", () => { function func8(params: { // The main message message: string, - + // The feeling feeling: "good" | "bad", - + // The number of words. // For example, 6 words: number @@ -1534,6 +1695,41 @@ describe("JinjaTemplateChatWrapper", () => { }); }); + describe("thought segment extraction", () => { + test("Qwen 3.6 resolves thinking segment settings properly", {timeout: 1000 * 60 * 60 * 2}, async () => { + const chatWrapper = new JinjaTemplateChatWrapper({ + template: qwen3_6Template + }); + expect(chatWrapper.keepOnlyLastThought).to.be.eql(true); + expect(chatWrapper.settings.segments?.thought).to.eql({ + openOnResponseStart: true, + prefix: LlamaText(new SpecialTokensText("\n")), + suffix: LlamaText(new SpecialTokensText("\n\n\n")) + }); + }); + + test("basic template resolves thinking segment settings properly", {timeout: 1000 * 60 * 60 * 2}, async () => { + const template = ( + "{%- for m in messages %}" + + "{{- '<|im_start|>' + m.role + '\\n' }}" + + "{%- if m.reasoning_content %}" + + "{{- '\\n' + m.reasoning_content + '\\n\\n' }}" + + "{%- endif %}" + + "{{- m.content + '<|im_end|>\\n' }}" + + "{%- endfor %}" + + "{%- if add_generation_prompt %}" + + "{{- '<|im_start|>assistant\\n\\n' }}" + + "{%- endif %}" + ); + const chatWrapper = new JinjaTemplateChatWrapper({template}); + expect(chatWrapper.settings.segments?.thought).to.eql({ + openOnResponseStart: true, + prefix: LlamaText(new SpecialTokensText("\n")), + suffix: LlamaText(new SpecialTokensText("\n\n")) + }); + }); + }); + test("Fails when messages are not present in the render output", () => { try { new JinjaTemplateChatWrapper({ diff --git a/test/standalone/chatWrappers/utils/jinjaTemplates.ts b/test/standalone/chatWrappers/utils/jinjaTemplates.ts index 9636e12c..654702e1 100644 --- a/test/standalone/chatWrappers/utils/jinjaTemplates.ts +++ b/test/standalone/chatWrappers/utils/jinjaTemplates.ts @@ -1,6 +1,6 @@ // source: https://huggingface.co/openai/gpt-oss-20b/blob/main/chat_template.jinja export const harmonyJinjaTemplate = ` -{# +{# In addition to the normal inputs of \`messages\` and \`tools\`, this template also accepts the following kwargs: - "builtin_tools": A list, can contain "browser" and/or "python". @@ -2172,6 +2172,394 @@ export const gemma4JinjaTemplate2 = ` {%- endif -%} `.slice(1); +export const gemma4JinjaTemplate3 = ` +{%- macro format_parameters(properties, required, filter_keys=false) -%} + {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in properties | dictsort -%} + {%- set add_comma = false -%} + {%- if not filter_keys or key not in standard_keys -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {{ key }}:{ + {%- if value['description'] -%} + description:<|"|>{{ value['description'] }}<|"|> + {%- set add_comma = true -%} + {%- endif -%} + {%- if value['type'] | upper == 'STRING' -%} + {%- if value['enum'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + enum:{{ format_argument(value['enum']) }} + {%- endif -%} + {%- elif value['type'] | upper == 'ARRAY' -%} + {%- if value['items'] is mapping and value['items'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + items:{ + {%- set ns_items = namespace(found_first=false) -%} + {%- for item_key, item_value in value['items'] | dictsort -%} + {%- if item_value is not none -%} + {%- if ns_items.found_first %},{% endif -%} + {%- set ns_items.found_first = true -%} + {%- if item_key == 'properties' -%} + properties:{ + {%- if item_value is mapping -%} + {{- format_parameters(item_value, value['items']['required'] | default([])) -}} + {%- endif -%} + } + {%- elif item_key == 'required' -%} + required:[ + {%- for req_item in item_value -%} + <|"|>{{- req_item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- elif item_key == 'type' -%} + {%- if item_value is string -%} + type:{{ format_argument(item_value | upper) }} + {%- else -%} + type:{{ format_argument(item_value | map('upper') | list) }} + {%- endif -%} + {%- else -%} + {{ item_key }}:{{ format_argument(item_value) }} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + } + {%- endif -%} + {%- endif -%} + {%- if value['nullable'] %} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + nullable:true + {%- endif -%} + {%- if value['type'] | upper == 'OBJECT' -%} + {%- if value['properties'] is defined and value['properties'] is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value['properties'], value['required'] | default([])) -}} + } + {%- elif value is mapping -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + properties:{ + {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}} + } + {%- endif -%} + {%- if value['required'] -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + required:[ + {%- for item in value['required'] | default([]) -%} + <|"|>{{- item -}}<|"|> + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + ] + {%- endif -%} + {%- endif -%} + {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%} + type:<|"|>{{ value['type'] | upper }}<|"|>} + {%- endif -%} + {%- endfor -%} +{%- endmacro -%} +{%- macro format_function_declaration(tool_data) -%} + declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|> + {%- set params = tool_data['function']['parameters'] -%} + {%- if params -%} + ,parameters:{ + {%- if params['properties'] -%} + properties:{ {{- format_parameters(params['properties'], params['required']) -}} }, + {%- endif -%} + {%- if params['required'] -%} + required:[ + {%- for item in params['required'] -%} + <|"|>{{- item -}}<|"|> + {{- ',' if not loop.last -}} + {%- endfor -%} + ], + {%- endif -%} + {%- if params['type'] -%} + type:<|"|>{{- params['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + {%- if 'response' in tool_data['function'] -%} + {%- set response_declaration = tool_data['function']['response'] -%} + ,response:{ + {%- if response_declaration['description'] -%} + description:<|"|>{{- response_declaration['description'] -}}<|"|>, + {%- endif -%} + {%- if response_declaration['type'] | upper == 'OBJECT' -%} + type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>} + {%- endif -%} + {%- endif -%} + } +{%- endmacro -%} +{%- macro format_argument(argument, escape_keys=True) -%} + {%- if argument is none -%} + {{- 'null' -}} + {%- elif argument is string -%} + {{- '<|"|>' + argument + '<|"|>' -}} + {%- elif argument is boolean -%} + {{- 'true' if argument else 'false' -}} + {%- elif argument is mapping -%} + {{- '{' -}} + {%- set ns = namespace(found_first=false) -%} + {%- for key, value in argument | dictsort -%} + {%- if ns.found_first %},{% endif -%} + {%- set ns.found_first = true -%} + {%- if escape_keys -%} + {{- '<|"|>' + key + '<|"|>' -}} + {%- else -%} + {{- key -}} + {%- endif -%} + :{{- format_argument(value, escape_keys=escape_keys) -}} + {%- endfor -%} + {{- '}' -}} + {%- elif argument is sequence -%} + {{- '[' -}} + {%- for item in argument -%} + {{- format_argument(item, escape_keys=escape_keys) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- ']' -}} + {%- else -%} + {{- argument -}} + {%- endif -%} +{%- endmacro -%} +{%- macro strip_thinking(text) -%} + {%- set ns = namespace(result='') -%} + {%- for part in text.split('') -%} + {%- if '<|channel>' in part -%} + {%- set ns.result = ns.result + part.split('<|channel>')[0] -%} + {%- else -%} + {%- set ns.result = ns.result + part -%} + {%- endif -%} + {%- endfor -%} + {{- ns.result | trim -}} +{%- endmacro -%} + +{%- macro format_tool_response_block(tool_name, response) -%} + {{- '<|tool_response>' -}} + {%- if response is mapping -%} + {{- 'response:' + tool_name + '{' -}} + {%- for key, value in response | dictsort -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- if not loop.last %},{% endif -%} + {%- endfor -%} + {{- '}' -}} + {%- else -%} + {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}} + {%- endif -%} + {{- '' -}} +{%- endmacro -%} + +{#- ===== SETUP ===== -#} +{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%} +{%- set loop_messages = messages -%} +{%- set enable_thinking = enable_thinking | default(false) -%} +{%- set preserve_thinking = preserve_thinking | default(false) -%} +{{- bos_token -}} +{#- Handle System/Tool Definitions Block -#} +{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%} + {{- '<|turn>system\n' -}} + {#- Inject Thinking token at the very top of the FIRST system turn -#} + {%- if enable_thinking -%} + {{- '<|think|>\n' -}} + {%- set ns.prev_message_type = 'think' -%} + {%- endif -%} + {%- if messages and messages[0]['role'] in ['system', 'developer'] -%} + {%- if messages[0]['content'] is string -%} + {{- messages[0]['content'] | trim -}} + {%- elif messages[0]['content'] is sequence -%} + {%- for item in messages[0]['content'] -%} + {{- item['text'] | trim + ' '-}} + {%- endfor -%} + {%- endif -%} + {%- set loop_messages = messages[1:] -%} + {%- endif -%} + {%- if tools -%} + {%- for tool in tools %} + {{- '<|tool>' -}} + {{- format_function_declaration(tool) | trim -}} + {{- '' -}} + {%- endfor %} + {%- set ns.prev_message_type = 'tool' -%} + {%- endif -%} + {{- '\n' -}} +{%- endif %} + +{#- Pre-scan: find last user message index for reasoning guard -#} +{%- set ns_turn = namespace(last_user_idx=-1) -%} +{%- for i in range(loop_messages | length) -%} + {%- if loop_messages[i]['role'] == 'user' -%} + {%- set ns_turn.last_user_idx = i -%} + {%- endif -%} +{%- endfor -%} + +{#- Loop through messages -#} +{%- for message in loop_messages -%} + {%- if message['role'] != 'tool' -%} + {%- set ns.prev_message_type = None -%} + {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%} + {#- Detect continuation using tracked state — O(1) instead of O(n) backward scan -#} + {%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%} + {%- if not continue_same_model_turn -%} + {{- '<|turn>' + role + '\n' }} + {%- endif -%} + + {#- Render reasoning/reasoning_content as thinking channel -#} + {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%} + {%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or (preserve_thinking and message.get('tool_calls')) -%} + {%- if thinking_text and thinking_gate -%} + {{- '<|channel>thought\n' + thinking_text + '\n' -}} + {%- endif -%} + + {%- if message.get('tool_calls') -%} + {%- for tool_call in message.get('tool_calls') -%} + {%- set function = tool_call['function'] -%} + {{- '<|tool_call>call:' + function['name'] + '{' -}} + {%- if function['arguments'] is mapping -%} + {%- set ns_args = namespace(found_first=false) -%} + {%- for key, value in function['arguments'] | dictsort -%} + {%- if ns_args.found_first %},{% endif -%} + {%- set ns_args.found_first = true -%} + {{- key -}}:{{- format_argument(value, escape_keys=False) -}} + {%- endfor -%} + {%- elif function['arguments'] is none -%} + {%- elif function['arguments'] is string -%} + {#- Pre-serialized args (e.g. an OpenAI JSON string). We cannot JSON-parse + portably in-template, so render non-fatally instead of erroring. Strip an + outer {...} so it composes with the DSL braces rather than double-wrapping. + Prefer passing arguments as a mapping for exact Gemma DSL. -#} + {%- set argstr = function['arguments'] | trim -%} + {%- if argstr[:1] == '{' and argstr[-1:] == '}' -%} + {{- argstr[1:-1] -}} + {%- else -%} + {{- function['arguments'] -}} + {%- endif -%} + {%- endif -%} + {{- '}' -}} + {%- endfor -%} + {%- set ns.prev_message_type = 'tool_call' -%} + {%- endif -%} + + {%- set ns_tr_out = namespace(flag=false) -%} + {%- if message.get('tool_responses') -%} + {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#} + {%- for tool_response in message.get('tool_responses') -%} + {{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endfor -%} + {%- elif message.get('tool_calls') -%} + {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#} + {%- set ns_tool_scan = namespace(stopped=false) -%} + {%- for k in range(loop.index0 + 1, loop_messages | length) -%} + {%- if ns_tool_scan.stopped -%} + {%- elif loop_messages[k]['role'] != 'tool' -%} + {%- set ns_tool_scan.stopped = true -%} + {%- else -%} + {%- set follow = loop_messages[k] -%} + {#- Resolve tool_call_id to function name -#} + {%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%} + {%- for tc in message.get('tool_calls') -%} + {%- if tc.get('id') == follow.get('tool_call_id') -%} + {%- set ns_tname.name = tc['function']['name'] -%} + {%- endif -%} + {%- endfor -%} + {#- Handle content as string or content-parts array -#} + {%- set tool_body = follow.get('content') -%} + {%- if tool_body is string -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- elif tool_body is sequence and tool_body is not string -%} + {%- set ns_txt = namespace(s='') -%} + {%- for part in tool_body -%} + {%- if part.get('type') == 'text' -%} + {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%} + {%- endif -%} + {%- endfor -%} + {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}} + {%- for part in tool_body -%} + {%- if part.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif part.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif part.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {{- format_tool_response_block(ns_tname.name, tool_body) -}} + {%- endif -%} + {%- set ns_tr_out.flag = true -%} + {%- set ns.prev_message_type = 'tool_response' -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + + {%- set captured_content -%} + {%- if message.get('content') is string -%} + {%- if role == 'model' -%} + {{- strip_thinking(message['content']) -}} + {%- else -%} + {{- message['content'] | trim -}} + {%- endif -%} + {%- elif message.get('content') is sequence -%} + {%- for item in message['content'] -%} + {%- if item.get('type') == 'text' -%} + {%- if role == 'model' -%} + {{- strip_thinking(item['text']) -}} + {%- else -%} + {{- item['text'] | trim -}} + {%- endif -%} + {%- elif item.get('type') in ['image', 'image_url'] -%} + {{- '<|image|>' -}} + {%- elif item.get('type') in ['audio', 'input_audio'] -%} + {{- '<|audio|>' -}} + {%- elif item.get('type') == 'video' -%} + {{- '<|video|>' -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endset -%} + + {{- captured_content -}} + {%- set has_content = captured_content | trim | length > 0 -%} + + {#- Forward-scan: find next non-tool message role for continuation detection -#} + {%- set next_nt = namespace(role=None, found=false) -%} + {%- for j in range(loop.index0 + 1, loop_messages | length) -%} + {%- if not next_nt.found -%} + {%- if loop_messages[j]['role'] != 'tool' -%} + {%- set next_nt.role = loop_messages[j]['role'] -%} + {%- set next_nt.found = true -%} + {%- endif -%} + {%- endif -%} + {%- endfor -%} + + {%- set continues_into_next = ( + role == 'model' + and next_nt.role == 'assistant' + and (not message.get('tool_calls') or ns_tr_out.flag) + ) -%} + + {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%} + {{- '<|tool_response>' -}} + {%- elif continues_into_next -%} + {%- elif not (ns_tr_out.flag and not has_content and not next_nt.found) -%} + {{- '\n' -}} + {%- endif -%} + + {#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#} + {%- set ns.prev_non_tool_role = message['role'] -%} + {%- endif -%} +{%- endfor -%} + +{%- if add_generation_prompt -%} + {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%} + {{- '<|turn>model\n' -}} + {%- elif ns.prev_message_type == 'tool_response' and enable_thinking -%} + {{- '<|channel>thought\n' -}} + {%- endif -%} +{%- endif -%} +`.slice(1, -1); + export const lfm2_5JinjaTemplate = ` {{- bos_token -}} {%- set preserve_thinking = preserve_thinking | default(false) -%} diff --git a/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts b/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts index e68f9024..39da3d5d 100644 --- a/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts +++ b/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts @@ -6,7 +6,7 @@ import { } from "../../../../src/index.js"; import { harmonyJinjaTemplate, harmonyJinjaTemplate2, harmonyJinjaTemplate3, harmonyJinjaTemplate4, harmonyJinjaTemplate5, - gemma4JinjaTemplate1, gemma4JinjaTemplate2 + gemma4JinjaTemplate1, gemma4JinjaTemplate2, gemma4JinjaTemplate3 } from "./jinjaTemplates.js"; @@ -811,6 +811,19 @@ describe("resolveChatWrapper", () => { expect(chatWrapper).to.be.instanceof(Gemma4ChatWrapper); }); + test("should resolve to specialized Gemma4ChatWrapper 3", () => { + const chatWrapper = resolveChatWrapper({ + customWrapperSettings: { + jinjaTemplate: { + template: gemma4JinjaTemplate3 + } + }, + fallbackToOtherWrappersOnJinjaError: false + }); + + expect(chatWrapper).to.be.instanceof(Gemma4ChatWrapper); + }); + test("should resolve to specialized GeneralChatWrapper", () => { const chatWrapper = resolveChatWrapper({ customWrapperSettings: {