diff --git a/llama/addon/AddonContext.cpp b/llama/addon/AddonContext.cpp index 113a3cf3..b0305b18 100644 --- a/llama/addon/AddonContext.cpp +++ b/llama/addon/AddonContext.cpp @@ -724,18 +724,31 @@ Napi::Value AddonContext::GetEmbedding(const Napi::CallbackInfo& info) { } int32_t inputTokensLength = info[0].As().Int32Value(); - int32_t maxVectorSize = (info.Length() > 1 && info[1].IsNumber()) ? info[1].As().Int32Value() : 0; + const double maxVectorSize = (info.Length() > 1 && info[1].IsNumber()) ? info[1].As().DoubleValue() : 0; if (inputTokensLength <= 0) { Napi::Error::New(info.Env(), "Invalid input tokens length").ThrowAsJavaScriptException(); return info.Env().Undefined(); } - const int n_embd = llama_model_n_embd(model->model); + if (!std::isfinite(maxVectorSize) || maxVectorSize < 0 || std::floor(maxVectorSize) != maxVectorSize) { + Napi::Error::New(info.Env(), "Invalid maximum embedding vector size").ThrowAsJavaScriptException(); + return info.Env().Undefined(); + } + const enum llama_pooling_type pooling_type = llama_pooling_type(ctx); + const int64_t n_embd = pooling_type == LLAMA_POOLING_TYPE_RANK + ? static_cast(llama_model_n_cls_out(model->model)) + : llama_model_n_embd_out(model->model); + + if (n_embd <= 0) { + Napi::Error::New(info.Env(), "Invalid embedding vector size").ThrowAsJavaScriptException(); + return info.Env().Undefined(); + } + const auto* embeddings = pooling_type == LLAMA_POOLING_TYPE_NONE ? NULL : llama_get_embeddings_seq(ctx, 0); if (embeddings == NULL) { - embeddings = llama_get_embeddings_ith(ctx, inputTokensLength - 1); + embeddings = llama_get_embeddings_ith(ctx, -1); } if (embeddings == NULL) { @@ -743,7 +756,7 @@ Napi::Value AddonContext::GetEmbedding(const Napi::CallbackInfo& info) { return info.Env().Undefined(); } - size_t resultSize = maxVectorSize == 0 ? n_embd : std::min(n_embd, maxVectorSize); + const size_t resultSize = maxVectorSize == 0 ? n_embd : std::min(n_embd, maxVectorSize); Napi::Float64Array result = Napi::Float64Array::New(info.Env(), resultSize); for (size_t i = 0; i < resultSize; i++) { result[i] = embeddings[i]; diff --git a/llama/addon/AddonJinjaRenderer.cpp b/llama/addon/AddonJinjaRenderer.cpp new file mode 100644 index 00000000..6bc1c98a --- /dev/null +++ b/llama/addon/AddonJinjaRenderer.cpp @@ -0,0 +1,189 @@ +#include "AddonJinjaRenderer.h" +#include "jinja/parser.h" + +#include +#include +#include +#include +#include + +namespace { + Napi::Array ownEnumerableKeys(const Napi::Object& object) { + napi_value keys; + const auto status = napi_get_all_property_names(object.Env(), object, napi_key_own_only, + static_cast(napi_key_enumerable | napi_key_skip_symbols), napi_key_numbers_to_strings, &keys); + if (status != napi_ok) { + throw Napi::Error::New(object.Env()); + } + return Napi::Array(object.Env(), keys); + } + + struct InputConverter { + std::string errorPath; + + jinja::value convert(const Napi::Value& value) { + if (value.IsUndefined()) { + return jinja::mk_val(); + } else if (value.IsNull()) { + return jinja::mk_val(); + } else if (value.IsBoolean()) { + return jinja::mk_val(value.As().Value()); + } else if (value.IsString()) { + auto result = jinja::mk_val(); + result->val_str.parts.push_back({false, value.As().Utf8Value()}); + return result; + } else if (value.IsNumber()) { + const double number = value.As().DoubleValue(); + constexpr double minInteger = static_cast(std::numeric_limits::min()); + if (!std::isfinite(number) || number < minInteger || number >= -minInteger) { + throw std::invalid_argument("Number is outside the native Jinja numeric range"); + } + if (std::trunc(number) == number) { + return jinja::mk_val(static_cast(number)); + } + return jinja::mk_val(number); + } else if (!value.IsObject() || value.IsFunction()) { + throw std::runtime_error("Unsupported JavaScript value; expected a string, number, boolean, null, undefined, array or object"); + } + + jinja::value result; + if (value.IsArray()) { + const auto array = value.As(); + const auto length = array.Length(); + auto converted = jinja::mk_val(); + converted->val_arr.reserve(length); + for (uint32_t i = 0; i < length; i++) { + Napi::HandleScope scope(value.Env()); + try { + converted->push_back(convert(array.Get(i))); + } catch (...) { + errorPath.insert(0, "[" + std::to_string(i) + "]"); + throw; + } + } + result = std::move(converted); + } else { + const auto object = value.As(); + const auto keys = ownEnumerableKeys(object); + const auto length = keys.Length(); + auto converted = jinja::mk_val(); + converted->val_obj.reserve(length); + converted->unordered.reserve(length); + for (uint32_t i = 0; i < length; i++) { + Napi::HandleScope scope(value.Env()); + const auto key = keys.Get(i).As(); + const auto name = key.Utf8Value(); + try { + converted->insert(name, convert(object.Get(key))); + } catch (...) { + errorPath.insert(0, "[\"" + name + "\"]"); + throw; + } + } + result = std::move(converted); + } + return result; + } + }; + + // like gather_string_parts_recursive() from llama/llama.cpp/common/jinja/runtime.h but without copying arrays or string parts + void appendResult(const jinja::value& value, std::string& output) { + if (jinja::is_val(value)) { + for (const auto& part : value->val_str.parts) { + output.append(part.val); + } + } else if (jinja::is_val(value)) { + for (const auto& item : value->as_array()) { + appendResult(item, output); + } + } else if (jinja::is_val(value) || jinja::is_val(value) || + jinja::is_val(value)) { + output.append(value->as_string().str()); + } + } + + void throwError(Napi::Env env, const std::string& message, const std::exception& error) { + auto result = Napi::Error::New(env, "AddonJinjaRenderer: " + message + ": " + error.what()); + const auto* jsError = dynamic_cast(&error); + if (jsError != nullptr) { + result.Set("cause", jsError->Value()); + } + result.ThrowAsJavaScriptException(); + } +} + +AddonJinjaRenderer::AddonJinjaRenderer(const Napi::CallbackInfo& info) : Napi::ObjectWrap(info) { + try { + if (info.Length() != 1 || !info[0].IsString()) { + throw std::invalid_argument("Constructor expects a template string"); + } + auto text = info[0].As().Utf8Value(); + lexerResult = jinja::lexer().tokenize(text); + program.emplace(jinja::parse_from_tokens(lexerResult)); + source = std::make_shared(std::move(text)); + } catch (const std::exception& error) { + throwError(info.Env(), "failed to compile template", error); + } catch (...) { + Napi::Error::New(info.Env(), "AddonJinjaRenderer: unknown error compiling template").ThrowAsJavaScriptException(); + } +} + +Napi::Value AddonJinjaRenderer::render(const Napi::CallbackInfo& info) { + if (source.get() == nullptr) { + Napi::Error::New(info.Env(), "AddonJinjaRenderer: failed to render template: Renderer is not initialized").ThrowAsJavaScriptException(); + return info.Env().Undefined(); + } + + InputConverter converter; + jinja::value input; + try { + if (info.Length() > 1 || (info.Length() == 1 && !info[0].IsUndefined() && + (!info[0].IsObject() || info[0].IsArray() || info[0].IsFunction()))) { + throw std::invalid_argument("render expects an optional object of template variables"); + } + if (info.Length() == 1 && !info[0].IsUndefined()) { + input = converter.convert(info[0]); + } + } catch (const std::exception& error) { + throwError(info.Env(), "failed to convert input at items" + converter.errorPath, error); + return info.Env().Undefined(); + } catch (...) { + Napi::Error::New(info.Env(), "AddonJinjaRenderer: failed to convert input at items" + converter.errorPath + ": unknown native error").ThrowAsJavaScriptException(); + return info.Env().Undefined(); + } + + try { + jinja::context context; + context.src = source; + if (input.get() != nullptr) { + for (const auto& item : input->as_ordered_object()) { + context.set_val(item.first, item.second); + } + input.reset(); + } + + if (!program.has_value()) { + program.emplace(jinja::parse_from_tokens(lexerResult)); + } + + jinja::runtime runtime(context); + const auto result = runtime.execute(*program); + std::string output; + appendResult(result, output); + return Napi::String::New(info.Env(), output); + } catch (const std::exception& error) { + // upstream filter blocks can move AST nodes before throwing + program.reset(); + throwError(info.Env(), "failed to render template", error); + } catch (...) { + program.reset(); + Napi::Error::New(info.Env(), "AddonJinjaRenderer: failed to render template: unknown native error").ThrowAsJavaScriptException(); + } + return info.Env().Undefined(); +} + +void AddonJinjaRenderer::init(Napi::Object exports) { + exports.Set("AddonJinjaRenderer", DefineClass(exports.Env(), "AddonJinjaRenderer", { + InstanceMethod("render", &AddonJinjaRenderer::render) + })); +} diff --git a/llama/addon/AddonJinjaRenderer.h b/llama/addon/AddonJinjaRenderer.h new file mode 100644 index 00000000..6e9824bb --- /dev/null +++ b/llama/addon/AddonJinjaRenderer.h @@ -0,0 +1,22 @@ +#pragma once + +#include "napi.h" +#include "jinja/lexer.h" +#include "jinja/runtime.h" + +#include +#include + +class AddonJinjaRenderer : public Napi::ObjectWrap { + private: + jinja::lexer_result lexerResult; + std::shared_ptr source; + std::optional program; + + public: + AddonJinjaRenderer(const Napi::CallbackInfo& info); + + Napi::Value render(const Napi::CallbackInfo& info); + + static void init(Napi::Object exports); +}; diff --git a/llama/addon/AddonModel.cpp b/llama/addon/AddonModel.cpp index 356a30e3..aa3f1b8e 100644 --- a/llama/addon/AddonModel.cpp +++ b/llama/addon/AddonModel.cpp @@ -312,6 +312,7 @@ AddonModel::AddonModel(const Napi::CallbackInfo& info) : Napi::ObjectWrap(info) { data = new AddonModelData(); model_params = llama_model_default_params(); + model_params.lazy_mode = LLAMA_LAZY_MODE_OFF; modelPath = info[0].As().Utf8Value(); @@ -349,6 +350,18 @@ AddonModel::AddonModel(const Napi::CallbackInfo& info) : model_params.no_alloc = options.Get("noAlloc").As().Value(); } + if (options.Has("lazyMode")) { + auto lazyMode = options.Get("lazyMode"); + + if (lazyMode.IsString() && (lazyMode.As().Utf8Value() == "auto")) { + model_params.lazy_mode = LLAMA_LAZY_MODE_AUTO; + } else if (lazyMode.IsBoolean() && lazyMode.As().Value()) { + model_params.lazy_mode = LLAMA_LAZY_MODE_ON; + } else if (lazyMode.IsBoolean() && !lazyMode.As().Value()) { + model_params.lazy_mode = LLAMA_LAZY_MODE_OFF; + } + } + if (options.Has("onLoadProgress")) { auto onLoadProgressJSCallback = options.Get("onLoadProgress").As(); if (onLoadProgressJSCallback.IsFunction()) { @@ -627,7 +640,7 @@ Napi::Value AddonModel::GetEmbeddingVectorSize(const Napi::CallbackInfo& info) { return info.Env().Undefined(); } - return Napi::Number::From(info.Env(), llama_model_n_embd(model)); + return Napi::Number::From(info.Env(), llama_model_n_embd_out(model)); } Napi::Value AddonModel::GetTotalSize(const Napi::CallbackInfo& info) { diff --git a/llama/addon/addon.cpp b/llama/addon/addon.cpp index f8ce97af..e54b59b2 100644 --- a/llama/addon/addon.cpp +++ b/llama/addon/addon.cpp @@ -8,6 +8,7 @@ #include "AddonGgufMetadata.h" #include "AddonGrammar.h" #include "AddonGrammarEvaluationState.h" +#include "AddonJinjaRenderer.h" #include "AddonModel.h" #include "AddonModelLora.h" #include "AddonSampler.h" @@ -19,6 +20,7 @@ #include "globals/getSwapInfo.h" #include "globals/getSystemMemoryInfo.h" #include "globals/addonEnv.h" +#include "llama-arch.h" std::mutex backendMutex; bool backendInitialized = false; @@ -161,6 +163,77 @@ Napi::Value addonGetConsts(const Napi::CallbackInfo& info) { return consts; } +Napi::Value addonGetAllArchs(const Napi::CallbackInfo& info) { + auto allArchs = llm_arch_all(); + Napi::Array archs = Napi::Array::New(info.Env(), allArchs.size()); + + for (size_t i = 0; i < allArchs.size(); ++i) { + auto archName = llm_arch_name(allArchs[i]); + if (archName == nullptr) { + archs[i] = info.Env().Undefined(); + continue; + } + + archs[i] = Napi::String::New(info.Env(), archName); + } + + return archs; +} + +Napi::Value addonGetIsArchSupported(const Napi::CallbackInfo& info) { + if (info.Length() == 0 || !info[0].IsString()) { + return Napi::Boolean::New(info.Env(), false); + } + + std::string archName = info[0].As().Utf8Value(); + if (archName.empty()) { + return Napi::Boolean::New(info.Env(), false); + } + + auto arch = llm_arch_from_string(archName); + if (arch == LLM_ARCH_UNKNOWN) { + return Napi::Boolean::New(info.Env(), false); + } + + return Napi::Boolean::New(info.Env(), true); +} + +Napi::Value addonGetIsArchRecurrent(const Napi::CallbackInfo& info) { + if (info.Length() == 0 || !info[0].IsString()) { + return info.Env().Undefined(); + } + + std::string archName = info[0].As().Utf8Value(); + if (archName.empty()) { + return info.Env().Undefined(); + } + + auto arch = llm_arch_from_string(archName); + if (arch == LLM_ARCH_UNKNOWN) { + return info.Env().Undefined(); + } + + return Napi::Boolean::New(info.Env(), llm_arch_is_recurrent(arch)); +} + +Napi::Value addonGetIsArchHybrid(const Napi::CallbackInfo& info) { + if (info.Length() == 0 || !info[0].IsString()) { + return info.Env().Undefined(); + } + + std::string archName = info[0].As().Utf8Value(); + if (archName.empty()) { + return info.Env().Undefined(); + } + + auto arch = llm_arch_from_string(archName); + if (arch == LLM_ARCH_UNKNOWN) { + return info.Env().Undefined(); + } + + return Napi::Boolean::New(info.Env(), llm_arch_is_hybrid(arch)); +} + class AddonBackendLoadWorker : public Napi::AsyncWorker { public: AddonBackendLoadWorker(const Napi::Env& env) @@ -358,6 +431,10 @@ Napi::Object registerCallback(Napi::Env env, Napi::Object exports) { Napi::PropertyDescriptor::Function("getGgmlGraphOverheadCustom", addonGetGgmlGraphOverheadCustom), Napi::PropertyDescriptor::Function("getGgmlType", addonGetGgmlType), Napi::PropertyDescriptor::Function("getConsts", addonGetConsts), + Napi::PropertyDescriptor::Function("getAllArchs", addonGetAllArchs), + Napi::PropertyDescriptor::Function("getIsArchSupported", addonGetIsArchSupported), + Napi::PropertyDescriptor::Function("getIsArchRecurrent", addonGetIsArchRecurrent), + Napi::PropertyDescriptor::Function("getIsArchHybrid", addonGetIsArchHybrid), Napi::PropertyDescriptor::Function("setLogger", setLogger), Napi::PropertyDescriptor::Function("setLoggerLogLevel", setLoggerLogLevel), Napi::PropertyDescriptor::Function("setLoggerLogLevelOverride", setLoggerLogLevelOverride), @@ -379,6 +456,7 @@ Napi::Object registerCallback(Napi::Env env, Napi::Object exports) { AddonModelLora::init(exports); AddonGrammar::init(exports); AddonGrammarEvaluationState::init(exports); + AddonJinjaRenderer::init(exports); AddonContext::init(exports); AddonContextSequenceCheckpoint::init(exports); AddonSampler::init(exports); diff --git a/llama/addon/globals/getGpuInfo.cpp b/llama/addon/globals/getGpuInfo.cpp index 9ded4cba..6076a47a 100644 --- a/llama/addon/globals/getGpuInfo.cpp +++ b/llama/addon/globals/getGpuInfo.cpp @@ -133,7 +133,7 @@ Napi::Value getGpuType(const Napi::CallbackInfo& info) { } else if (backendName == "Vulkan") { return Napi::String::New(info.Env(), "vulkan"); } - + // else if ( // backendName == "CUDA" || backendName == "ROCm" || backendName == "MUSA") { // return Napi::String::New(info.Env(), "cuda"); diff --git a/src/bindings/AddonTypes.ts b/src/bindings/AddonTypes.ts index 9c4911be..92fab613 100644 --- a/src/bindings/AddonTypes.ts +++ b/src/bindings/AddonTypes.ts @@ -9,6 +9,7 @@ export type AddonModelParams = { useDirectIo?: boolean, useMlock?: boolean, checkTensors?: boolean, + lazyMode?: "auto" | boolean, overridesList?: Array<[key: string, value: number | bigint | boolean | string, type: 0 | 1 | undefined]> }; @@ -57,6 +58,9 @@ export type BindingModule = { new (model: AddonModel, grammar: AddonGrammar): AddonGrammarEvaluationState, new (existingState: AddonGrammarEvaluationState): AddonGrammarEvaluationState }, + AddonJinjaRenderer: { + new (template: string): AddonJinjaRenderer + }, AddonSampler: { new (model: AddonModel): AddonSampler, acceptGrammarEvaluationStateToken(grammarEvaluationState: AddonGrammarEvaluationState, token: Token): void, @@ -81,6 +85,10 @@ export type BindingModule = { llamaPosSize: number, llamaSeqIdSize: number }, + getAllArchs(): string[], + getIsArchSupported(architecture: string): boolean, + getIsArchRecurrent(architecture: string): boolean | undefined, + getIsArchHybrid(architecture: string): boolean | undefined, setLogger(logger: (level: number, message: string) => void): void, setLoggerLogLevel(level: number): void, setLoggerLogLevelOverride(level: number | undefined): void, @@ -215,6 +223,10 @@ export type AddonGrammar = { isTextCompatible(testText: string): boolean }; +export type AddonJinjaRenderer = { + render(items?: Record): string +}; + export type AddonGrammarEvaluationState = "AddonGrammarEvaluationState" & { readonly __brand: never }; diff --git a/src/bindings/utils/compileLLamaCpp.ts b/src/bindings/utils/compileLLamaCpp.ts index f7588304..d84fd8f9 100644 --- a/src/bindings/utils/compileLLamaCpp.ts +++ b/src/bindings/utils/compileLLamaCpp.ts @@ -117,7 +117,7 @@ export async function compileLlamaCpp(buildOptions: BuildOptions, compileOptions ) cmakeToolchainOptions.set("GGML_VULKAN_SHADERS_GEN_TOOLCHAIN", toolchainFile); - if (buildOptions.gpu === "metal" && process.platform === "darwin" && !cmakeCustomOptions.has("GGML_METAL")) + if (buildOptions.gpu === "metal" && platform === "mac" && !cmakeCustomOptions.has("GGML_METAL")) cmakeCustomOptions.set("GGML_METAL", "1"); else if (!cmakeCustomOptions.has("GGML_METAL")) cmakeCustomOptions.set("GGML_METAL", "OFF"); diff --git a/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts b/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts index 236831c4..0fd54613 100644 --- a/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts +++ b/src/chatWrappers/generic/JinjaTemplateChatWrapper.ts @@ -15,6 +15,7 @@ import {jsonDumps} from "../utils/jsonDumps.js"; import {tryMatrix} from "../../utils/optionsMatrix.js"; import {getStandardizedChatWrapperSegmentDefinition} from "../../utils/getStandardizedChatWrapperSegmentDefinition.js"; import {replaceRegularTextInLlamaText} from "../utils/replaceRegularTextInLlamaText.js"; +import {LruCache} from "../../utils/LruCache.js"; import {ChatHistoryFunctionCallMessageTemplate, parseFunctionCallMessageTemplate} from "./utils/chatHistoryFunctionCallMessageTemplate.js"; import { templateSegmentOptionsToChatWrapperSettings, TemplateChatWrapperSegmentsOptions @@ -25,6 +26,8 @@ import { } from "./utils/extractFunctionCallSettingsFromJinjaTemplate.js"; import {squashChatHistoryItems} from "./utils/squashChatHistoryItems.js"; import {extractSegmentSettingsFromTokenizerAndChatTemplate} from "./utils/extractSegmentSettingsFromTokenizerAndChatTemplate.js"; +import type {Llama} from "../../bindings/Llama.js"; +import type {AddonJinjaRenderer} from "../../bindings/AddonTypes.js"; export type JinjaTemplateChatWrapperOptions = { template: string, @@ -128,7 +131,13 @@ export type JinjaTemplateChatWrapperOptions = { _requireFunctionCallSettingsExtraction?: boolean, /** @internal */ - _functionCallExtractionExamineNonFirst?: boolean + _functionCallExtractionExamineNonFirst?: boolean, + + /** @internal */ + _cachedJinjaEngine?: CachedJinjaEngine, + + /** @internal */ + _templateCacheKeys?: object[] }; export type JinjaTemplateChatWrapperOptionsConvertMessageFormat = { @@ -184,7 +193,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { public readonly keepOnlyLastThought: boolean; public readonly additionalRenderParameters?: Record; - /** @internal */ private readonly _jinjaTemplate: Template; + /** @internal */ private readonly _jinjaTemplate: JinjaRenderer; /** @internal */ private readonly _usingJinjaFunctionCallTemplate: boolean = false; /** @internal */ private readonly _stringifyFunctionParams: boolean = false; /** @internal */ private readonly _wrapFunctionParamsInsideMapKey?: string; @@ -213,7 +222,9 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { segments, tokenizer, _requireFunctionCallSettingsExtraction = false, - _functionCallExtractionExamineNonFirst = false + _functionCallExtractionExamineNonFirst = false, + _cachedJinjaEngine = CachedJinjaEngine._create(), + _templateCacheKeys = [] } = options; if (template == null) @@ -233,7 +244,7 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { if (this.convertUnsupportedSystemMessagesToUserMessages != null && !this.convertUnsupportedSystemMessagesToUserMessages.format.includes("{{message}}")) throw new Error('convertUnsupportedSystemMessagesToUserMessages format must include "{{message}}"'); - this._jinjaTemplate = new Template(this.template); + this._jinjaTemplate = _cachedJinjaEngine.getFor(_templateCacheKeys, this.template); this.settings = { ...ChatWrapper.defaultSettings, @@ -987,6 +998,102 @@ export class JinjaTemplateChatWrapper extends ChatWrapper { } } +export class CachedJinjaEngine { + private _llama?: Llama; + private _cache: LruCache = new LruCache(40); + private _weakCache = new WeakMap(); + + private constructor(_llama?: Llama) { + this._llama = _llama; + } + + public getFor(keys: Array, template: string): JinjaRenderer { + let renderer: JinjaRenderer | undefined; + for (const key of keys) { + renderer = this._weakCache.get(key); + if (renderer != null) + break; + } + + if (renderer == null) { + renderer = this._cache.get(template); + + if (renderer == null) + renderer = JinjaRenderer._create(this._llama, template); + } + + for (const key of keys) + this._weakCache.set(key, renderer); + + this._cache.set(renderer._template, renderer); + + return renderer; + } + + /** @internal */ + public static _create(_llama?: Llama) { + return new CachedJinjaEngine(_llama); + } +} + +export class JinjaRenderer { + /** @internal */ public _template: string; + /** @internal */ private _llama?: Llama; + /** @internal */ private _jsRenderer?: null | Template; + /** @internal */ private _jsRendererError?: unknown; + /** @internal */ private _nativeRenderer?: null | AddonJinjaRenderer; + /** @internal */ private _nativeRendererInitError?: unknown; + + private constructor(llama: Llama | undefined, template: string) { + this._llama = llama; + this._template = template; + } + + public render(items?: Record): string { + try { + if (this._jsRenderer === undefined) + this._jsRenderer = new Template(this._template); + } catch (error) { + this._jsRenderer = null; + this._jsRendererError = error; + } + + if (this._jsRenderer != null) + return this._jsRenderer.render(items); + + try { + if (this._nativeRenderer === undefined && this._llama != null) + this._nativeRenderer = new this._llama._bindings.AddonJinjaRenderer(this._template); + } catch (error) { + this._nativeRenderer = null; + this._nativeRendererInitError = error; + } + + if (this._nativeRenderer != null) + return this._nativeRenderer.render(items); + + if (this._jsRendererError != null && this._nativeRendererInitError != null) + throw new AggregateError( + [this._jsRendererError, this._nativeRendererInitError], + "Jinja renderer failed. " + + String((this._jsRendererError as Error)?.message ?? this._jsRendererError) + ". " + + String((this._nativeRendererInitError as Error)?.message ?? this._nativeRendererInitError), + {cause: this._jsRendererError} + ); + else if (this._jsRendererError != null) + throw this._jsRendererError; + else if (this._nativeRendererInitError != null) + throw this._nativeRendererInitError; + + throw new Error("Failed to render Jinja template"); + } + + /** @internal */ + public static _create(llama: Llama | undefined, template: string) { + return new JinjaRenderer(llama, template); + } +} + function resolveConvertUnsupportedSystemMessagesToUserMessagesOption( convertUnsupportedSystemMessagesToUserMessages?: JinjaTemplateChatWrapperOptions["convertUnsupportedSystemMessagesToUserMessages"] ): JinjaTemplateChatWrapperOptionsConvertMessageFormat | undefined { diff --git a/src/chatWrappers/utils/resolveChatWrapper.ts b/src/chatWrappers/utils/resolveChatWrapper.ts index 27ba453f..1abf03c6 100644 --- a/src/chatWrappers/utils/resolveChatWrapper.ts +++ b/src/chatWrappers/utils/resolveChatWrapper.ts @@ -8,7 +8,7 @@ import {FunctionaryChatWrapper} from "../FunctionaryChatWrapper.js"; import {AlpacaChatWrapper} from "../AlpacaChatWrapper.js"; import {GemmaChatWrapper} from "../GemmaChatWrapper.js"; import {Gemma4ChatWrapper} from "../Gemma4ChatWrapper.js"; -import {JinjaTemplateChatWrapper, JinjaTemplateChatWrapperOptions} from "../generic/JinjaTemplateChatWrapper.js"; +import {CachedJinjaEngine, JinjaTemplateChatWrapper, JinjaTemplateChatWrapperOptions} from "../generic/JinjaTemplateChatWrapper.js"; import {TemplateChatWrapper} from "../generic/TemplateChatWrapper.js"; import {getConsoleLogPrefix} from "../../utils/getConsoleLogPrefix.js"; import {Llama3_1ChatWrapper} from "../Llama3_1ChatWrapper.js"; @@ -26,6 +26,7 @@ import {GgufArchitectureType} from "../../gguf/types/GgufMetadataTypes.js"; import {isJinjaTemplateEquivalentToSpecializedChatWrapper} from "./isJinjaTemplateEquivalentToSpecializedChatWrapper.js"; import {getModelLinageNames} from "./getModelLinageNames.js"; import type {GgufFileInfo} from "../../gguf/types/GgufFileInfoTypes.js"; +import type {Llama} from "../../bindings/Llama.js"; export const specializedChatWrapperTypeNames = Object.freeze([ @@ -115,7 +116,10 @@ export type ResolveChatWrapperOptions = { * * Defaults to `false`. */ - noJinja?: boolean + noJinja?: boolean, + + /** Optional Llama instance to optimize some internal operations */ + llama?: Llama }; export type ResolveChatWrapperWithModelOptions = { @@ -208,7 +212,8 @@ export function resolveChatWrapper( architecture: options.fileInfo?.metadata?.general?.architecture, filename: options.filename, fileInfo: options.fileInfo, - tokenizer: options.tokenizer + tokenizer: options.tokenizer, + llama: options._llama }) ?? new GeneralChatWrapper(); const { @@ -221,7 +226,8 @@ export function resolveChatWrapper( customWrapperSettings, warningLogs = true, fallbackToOtherWrappersOnJinjaError = true, - noJinja = false + noJinja = false, + llama } = options; const architecture = archOption ?? fileInfo?.metadata?.general?.architecture; @@ -287,10 +293,15 @@ export function resolveChatWrapper( const modelJinjaTemplate = customWrapperSettings?.jinjaTemplate?.template ?? fileInfo?.metadata?.tokenizer?.chat_template; if (modelJinjaTemplate != null && modelJinjaTemplate.trim() !== "") { + const cachedJinjaEngine = CachedJinjaEngine._create(llama); const jinjaTemplateChatWrapperOptions: JinjaTemplateChatWrapperOptions = { tokenizer, ...(customWrapperSettings?.jinjaTemplate ?? {}), - template: modelJinjaTemplate + template: modelJinjaTemplate, + _cachedJinjaEngine: cachedJinjaEngine, + _templateCacheKeys: tokenizer == null + ? [options] + : [tokenizer, options] }; const chatWrapperNamesToCheck = orderChatWrapperNamesByAssumedCompatibilityWithModel( diff --git a/src/cli/commands/ChatCommand.ts b/src/cli/commands/ChatCommand.ts index 55bd27e6..b178e96d 100644 --- a/src/cli/commands/ChatCommand.ts +++ b/src/cli/commands/ChatCommand.ts @@ -85,6 +85,7 @@ type ChatCommand = { timing: boolean, mmap?: boolean, useDirectIo: boolean, + lazyMode?: "auto" | boolean, printTimings: boolean }; @@ -404,6 +405,24 @@ export const ChatCommand: CommandModule = { default: false, description: "Use Direct I/O usage when available" }) + .option("lazyMode", { + type: "string", + alias: ["lazy"], + + // yargs types don't support passing `false` as a choice, although it is supported by yargs + choices: ["auto", true, false] as const as string[], + coerce: (value) => { + if (value === "false") + return false; + else if (value === "true") + return true; + else if (value === "auto") + return "auto"; + + return value; + }, + description: "Lazily read tensors from the file on demand when they are needed, rather than loading all tensors upfront. Only works when using mmap" + }) .option("printTimings", { alias: "pt", type: "boolean", @@ -419,7 +438,7 @@ export const ChatCommand: CommandModule = { repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, maxTokens, reasoningBudget, noHistory, environmentFunctions, tokenPredictionDraftModel, tokenPredictionModelContextSize, debug, numa, meter, timing, mmap, useDirectIo, - printTimings + lazyMode, printTimings }) { try { await RunChat({ @@ -429,7 +448,7 @@ export const ChatCommand: CommandModule = { gpuLayers, lastTokensRepeatPenalty, repeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, maxTokens, reasoningBudget, noHistory, environmentFunctions, tokenPredictionDraftModel, tokenPredictionModelContextSize, - debug, numa, meter, timing, mmap, useDirectIo, printTimings + debug, numa, meter, timing, mmap, useDirectIo, lazyMode, printTimings }); } catch (err) { await new Promise((accept) => setTimeout(accept, 0)); // wait for logs to finish printing @@ -447,7 +466,7 @@ async function RunChat({ threads, temperature, minP, topK, topP, seed, xtc, gpuLayers, lastTokensRepeatPenalty, repeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, maxTokens, reasoningBudget, noHistory, environmentFunctions, tokenPredictionDraftModel, - tokenPredictionModelContextSize, debug, numa, meter, timing, mmap, useDirectIo, printTimings + tokenPredictionModelContextSize, debug, numa, meter, timing, mmap, useDirectIo, lazyMode, printTimings }: ChatCommand) { if (contextSize === -1) contextSize = undefined; if (gpuLayers === -1) gpuLayers = undefined; @@ -548,6 +567,7 @@ async function RunChat({ defaultContextSwaFullCache: swaFullCache, useMmap, useDirectIo, + lazyMode, ignoreMemorySafetyChecks: gpuLayers != null, onLoadProgress(loadProgress: number) { progressUpdater.setProgress(loadProgress); @@ -585,6 +605,7 @@ async function RunChat({ defaultContextSwaFullCache: swaFullCache, useMmap, useDirectIo, + lazyMode, onLoadProgress(loadProgress: number) { progressUpdater.setProgress(loadProgress); }, @@ -691,6 +712,7 @@ async function RunChat({ draftContext, useMmap, useDirectIo, + lazyMode, printBos: true, printEos: true, logBatchSize, diff --git a/src/cli/commands/CompleteCommand.ts b/src/cli/commands/CompleteCommand.ts index 2f1c1da3..a96444e0 100644 --- a/src/cli/commands/CompleteCommand.ts +++ b/src/cli/commands/CompleteCommand.ts @@ -67,6 +67,7 @@ type CompleteCommand = { timing: boolean, mmap?: boolean, useDirectIo: boolean, + lazyMode?: "auto" | boolean, printTimings: boolean }; @@ -324,6 +325,24 @@ export const CompleteCommand: CommandModule = { default: false, description: "Use Direct I/O usage when available" }) + .option("lazyMode", { + type: "string", + alias: ["lazy"], + + // yargs types don't support passing `false` as a choice, although it is supported by yargs + choices: ["auto", true, false] as const as string[], + coerce: (value) => { + if (value === "false") + return false; + else if (value === "true") + return true; + else if (value === "auto") + return "auto"; + + return value; + }, + description: "Lazily read tensors from the file on demand when they are needed, rather than loading all tensors upfront. Only works when using mmap" + }) .option("printTimings", { alias: "pt", type: "boolean", @@ -337,7 +356,7 @@ export const CompleteCommand: CommandModule = { topP, seed, xtc, gpuLayers, repeatPenalty, lastTokensRepeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, maxTokens, tokenPredictionDraftModel, tokenPredictionModelContextSize, - debug, numa, meter, timing, mmap, useDirectIo, printTimings + debug, numa, meter, timing, mmap, useDirectIo, lazyMode, printTimings }) { try { await RunCompletion({ @@ -346,7 +365,8 @@ export const CompleteCommand: CommandModule = { threads, temperature, minP, topK, topP, seed, xtc, gpuLayers, lastTokensRepeatPenalty, repeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, maxTokens, - tokenPredictionDraftModel, tokenPredictionModelContextSize, debug, numa, meter, timing, mmap, useDirectIo, printTimings + tokenPredictionDraftModel, tokenPredictionModelContextSize, debug, numa, meter, timing, mmap, useDirectIo, lazyMode, + printTimings }); } catch (err) { await new Promise((accept) => setTimeout(accept, 0)); // wait for logs to finish printing @@ -363,7 +383,8 @@ async function RunCompletion({ threads, temperature, minP, topK, topP, seed, xtc, gpuLayers, lastTokensRepeatPenalty, repeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, - tokenPredictionDraftModel, tokenPredictionModelContextSize, maxTokens, debug, numa, meter, timing, mmap, useDirectIo, printTimings + tokenPredictionDraftModel, tokenPredictionModelContextSize, maxTokens, debug, numa, meter, timing, mmap, useDirectIo, lazyMode, + printTimings }: CompleteCommand) { if (contextSize === -1) contextSize = undefined; if (gpuLayers === -1) gpuLayers = undefined; @@ -455,6 +476,7 @@ async function RunCompletion({ defaultContextSwaFullCache: swaFullCache, useMmap, useDirectIo, + lazyMode, ignoreMemorySafetyChecks: gpuLayers != null, onLoadProgress(loadProgress: number) { progressUpdater.setProgress(loadProgress); @@ -492,6 +514,7 @@ async function RunCompletion({ defaultContextSwaFullCache: swaFullCache, useMmap, useDirectIo, + lazyMode, onLoadProgress(loadProgress: number) { progressUpdater.setProgress(loadProgress); }, @@ -571,6 +594,7 @@ async function RunCompletion({ draftContext, useMmap, useDirectIo, + lazyMode, minTitleLength: "Complete".length + 1, logBatchSize, tokenMeterEnabled: meter, diff --git a/src/cli/commands/InfillCommand.ts b/src/cli/commands/InfillCommand.ts index 13ea8e31..2a022a6a 100644 --- a/src/cli/commands/InfillCommand.ts +++ b/src/cli/commands/InfillCommand.ts @@ -69,6 +69,7 @@ type InfillCommand = { timing: boolean, mmap?: boolean, useDirectIo: boolean, + lazyMode?: "auto" | boolean, printTimings: boolean }; @@ -334,6 +335,24 @@ export const InfillCommand: CommandModule = { default: false, description: "Use Direct I/O usage when available" }) + .option("lazyMode", { + type: "string", + alias: ["lazy"], + + // yargs types don't support passing `false` as a choice, although it is supported by yargs + choices: ["auto", true, false] as const as string[], + coerce: (value) => { + if (value === "false") + return false; + else if (value === "true") + return true; + else if (value === "auto") + return "auto"; + + return value; + }, + description: "Lazily read tensors from the file on demand when they are needed, rather than loading all tensors upfront. Only works when using mmap" + }) .option("printTimings", { alias: "pt", type: "boolean", @@ -347,7 +366,7 @@ export const InfillCommand: CommandModule = { topP, seed, xtc, gpuLayers, repeatPenalty, lastTokensRepeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, maxTokens, tokenPredictionDraftModel, tokenPredictionModelContextSize, - debug, numa, meter, timing, mmap, useDirectIo, printTimings + debug, numa, meter, timing, mmap, useDirectIo, lazyMode, printTimings }) { try { await RunInfill({ @@ -356,7 +375,8 @@ export const InfillCommand: CommandModule = { threads, temperature, minP, topK, topP, seed, xtc, gpuLayers, lastTokensRepeatPenalty, repeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, maxTokens, - tokenPredictionDraftModel, tokenPredictionModelContextSize, debug, numa, meter, timing, mmap, useDirectIo, printTimings + tokenPredictionDraftModel, tokenPredictionModelContextSize, debug, numa, meter, timing, mmap, useDirectIo, lazyMode, + printTimings }); } catch (err) { await new Promise((accept) => setTimeout(accept, 0)); // wait for logs to finish printing @@ -372,7 +392,8 @@ async function RunInfill({ kvCacheKeyType, kvCacheValueType, swaFullCache, maxRam, maxVram, threads, temperature, minP, topK, topP, seed, xtc, gpuLayers, lastTokensRepeatPenalty, repeatPenalty, penalizeRepeatingNewLine, repeatFrequencyPenalty, repeatPresencePenalty, dryRepeatPenaltyStrength, dryRepeatPenaltyBase, dryRepeatPenaltyAllowedLength, dryRepeatPenaltyLastTokens, - tokenPredictionDraftModel, tokenPredictionModelContextSize, maxTokens, debug, numa, meter, timing, mmap, useDirectIo, printTimings + tokenPredictionDraftModel, tokenPredictionModelContextSize, maxTokens, debug, numa, meter, timing, mmap, useDirectIo, lazyMode, + printTimings }: InfillCommand) { if (contextSize === -1) contextSize = undefined; if (gpuLayers === -1) gpuLayers = undefined; @@ -478,6 +499,7 @@ async function RunInfill({ defaultContextSwaFullCache: swaFullCache, useMmap, useDirectIo, + lazyMode, ignoreMemorySafetyChecks: gpuLayers != null, onLoadProgress(loadProgress: number) { progressUpdater.setProgress(loadProgress); @@ -515,6 +537,7 @@ async function RunInfill({ defaultContextSwaFullCache: swaFullCache, useMmap, useDirectIo, + lazyMode, onLoadProgress(loadProgress: number) { progressUpdater.setProgress(loadProgress); }, @@ -594,6 +617,7 @@ async function RunInfill({ draftContext, useMmap, useDirectIo, + lazyMode, logBatchSize, tokenMeterEnabled: meter, resolvedMaxRam, diff --git a/src/cli/commands/inspect/commands/InspectMeasureCommand.ts b/src/cli/commands/inspect/commands/InspectMeasureCommand.ts index 13ed01b1..bb66ff47 100644 --- a/src/cli/commands/inspect/commands/InspectMeasureCommand.ts +++ b/src/cli/commands/inspect/commands/InspectMeasureCommand.ts @@ -48,6 +48,7 @@ type InspectMeasureCommand = { memory: "vram" | "ram" | "all", mmap?: boolean, useDirectIo: boolean, + lazyMode?: "auto" | boolean, printHeaderBeforeEachLayer?: boolean, evaluateText?: string, repeatEvaluateText?: number @@ -186,6 +187,24 @@ export const InspectMeasureCommand: CommandModule default: false, description: "Use Direct I/O usage when available" }) + .option("lazyMode", { + type: "string", + alias: ["lazy"], + + // yargs types don't support passing `false` as a choice, although it is supported by yargs + choices: ["auto", true, false] as const as string[], + coerce: (value) => { + if (value === "false") + return false; + else if (value === "true") + return true; + else if (value === "auto") + return "auto"; + + return value; + }, + description: "Lazily read tensors from the file on demand when they are needed, rather than loading all tensors upfront. Only works when using mmap" + }) .option("printHeaderBeforeEachLayer", { alias: "ph", type: "boolean", @@ -207,8 +226,8 @@ export const InspectMeasureCommand: CommandModule async handler({ modelPath: ggufPath, header: headerArg, gpu, minLayers, maxLayers, minContextSize, maxContextSize, flashAttention, embedding, kvCacheKeyType, kvCacheValueType, swaFullCache, maxRam, maxVram, - batchSize, measures = 10, memory: measureMemoryType, mmap, useDirectIo, printHeaderBeforeEachLayer = true, evaluateText, - repeatEvaluateText + batchSize, measures = 10, memory: measureMemoryType, mmap, useDirectIo, lazyMode, printHeaderBeforeEachLayer = true, + evaluateText, repeatEvaluateText }: InspectMeasureCommand) { if (maxLayers === -1) maxLayers = undefined; if (maxContextSize === -1) maxContextSize = undefined; @@ -316,6 +335,7 @@ export const InspectMeasureCommand: CommandModule const done = await measureModel({ modelPath: resolvedGgufPath, useMmap, + lazyMode, useDirectIo, gpu: gpu == null ? undefined @@ -630,13 +650,14 @@ const detectedFileName = path.basename(__filename); const expectedFileName = "InspectMeasureCommand"; async function measureModel({ - modelPath, useMmap, useDirectIo, gpu, tests, initialMaxContextSize, maxContextSize, minContextSize, maxGpuLayers, minGpuLayers, - flashAttention, embedding, kvCacheKeyType, kvCacheValueType, swaFullCache, maxRam, maxVram, batchSize, evaluateText, + modelPath, useMmap, lazyMode, useDirectIo, gpu, tests, initialMaxContextSize, maxContextSize, minContextSize, maxGpuLayers, + minGpuLayers, flashAttention, embedding, kvCacheKeyType, kvCacheValueType, swaFullCache, maxRam, maxVram, batchSize, evaluateText, exitAfterMeasurement = false, onInfo }: { modelPath: string, useMmap?: "auto" | boolean, + lazyMode?: "auto" | boolean, useDirectIo?: boolean, gpu?: BuildGpu | "auto", tests: number, @@ -670,6 +691,7 @@ async function measureModel({ modelRamUsage: number, contextSize?: number, useMmap: boolean, + lazyMode: "auto" | boolean, contextVramUsage?: number, contextRamUsage?: number, contextStateSize?: number, @@ -754,6 +776,7 @@ async function measureModel({ type: "start", modelPath, useMmap, + lazyMode, useDirectIo, tests, initialMaxContextSize, @@ -806,6 +829,7 @@ async function measureModel({ modelRamUsage: message.modelRamUsage, contextSize: message.contextSize, useMmap: message.useMmap, + lazyMode: message.lazyMode, contextVramUsage: message.contextVramUsage, contextRamUsage: message.contextRamUsage, contextStateSize: message.contextStateSize, @@ -945,6 +969,7 @@ async function runTestWorkerLogic() { ? context.contextSize : context._llamaContext.contextSize, useMmap: model.useMmap, + lazyMode: model.lazyMode, contextVramUsage: postContextVramUsage - preContextVramUsage, contextRamUsage: postContextRamUsage - preContextRamUsage, contextStateSize: context instanceof LlamaContext @@ -984,10 +1009,11 @@ async function runTestWorkerLogic() { } async function testWithGpuLayers({ - modelPath, useMmap, useDirectIo, gpuLayers, tests, startContextSize, maxContextSize, minContextSize, flashAttention, embedding, - kvCacheKeyType, kvCacheValueType, swaFullCache, batchSize, evaluateText, exitAfterMeasurement = false, isFirstLoad + modelPath, useMmap, lazyMode, useDirectIo, gpuLayers, tests, startContextSize, maxContextSize, minContextSize, flashAttention, + embedding, kvCacheKeyType, kvCacheValueType, swaFullCache, batchSize, evaluateText, exitAfterMeasurement = false, isFirstLoad }: { - modelPath: string, useMmap?: "auto" | boolean, useDirectIo?: boolean, gpuLayers: number, tests: number, startContextSize?: number, + modelPath: string, useMmap?: "auto" | boolean, lazyMode?: "auto" | boolean, + useDirectIo?: boolean, gpuLayers: number, tests: number, startContextSize?: number, maxContextSize?: number, minContextSize?: number, flashAttention?: boolean, embedding?: boolean, kvCacheKeyType?: GgmlType, kvCacheValueType?: GgmlType, swaFullCache?: boolean, batchSize?: number, evaluateText?: string, exitAfterMeasurement?: boolean, @@ -1003,6 +1029,7 @@ async function runTestWorkerLogic() { model = await llama.loadModel({ modelPath, useMmap, + lazyMode, useDirectIo, gpuLayers, defaultContextFlashAttention: flashAttention, @@ -1030,6 +1057,7 @@ async function runTestWorkerLogic() { type: "stats", gpuLayers: model.gpuLayers, useMmap: model.useMmap, + lazyMode: model.lazyMode, modelVramUsage: postModelVramUsage - preModelVramUsage, modelRamUsage: postModelRamUsage - preModelRamUsage, totalVramUsage: postModelVramUsage, @@ -1092,6 +1120,7 @@ async function runTestWorkerLogic() { const measurementsDone = await testWithGpuLayers({ modelPath: message.modelPath, useMmap: message.useMmap, + lazyMode: message.lazyMode, useDirectIo: message.useDirectIo, gpuLayers, tests: message.tests, @@ -1194,6 +1223,7 @@ type ParentToChildMessage = { type: "start", modelPath: string, useMmap?: "auto" | boolean, + lazyMode?: "auto" | boolean, useDirectIo?: boolean, tests: number, maxGpuLayers: number, @@ -1224,6 +1254,7 @@ type ChildToParentMessage = { modelRamUsage: number, contextSize?: number, useMmap: boolean, + lazyMode: "auto" | boolean, contextVramUsage?: number, contextRamUsage?: number, contextStateSize?: number, diff --git a/src/cli/utils/printCommonInfoLines.ts b/src/cli/utils/printCommonInfoLines.ts index 42e4eb5d..44e71501 100644 --- a/src/cli/utils/printCommonInfoLines.ts +++ b/src/cli/utils/printCommonInfoLines.ts @@ -12,6 +12,7 @@ export async function printCommonInfoLines({ minTitleLength = 0, useMmap, useDirectIo, + lazyMode, logBatchSize = false, tokenMeterEnabled = false, printBos = false, @@ -25,6 +26,7 @@ export async function printCommonInfoLines({ minTitleLength?: number, useMmap?: "auto" | boolean, useDirectIo?: boolean, + lazyMode?: "auto" | boolean, logBatchSize?: boolean, tokenMeterEnabled?: boolean, printBos?: boolean, @@ -124,6 +126,18 @@ export async function printCommonInfoLines({ : (useDirectIo || useDirectIo == null) ? "enabled" : "disabled" + }, { + show: lazyMode != null, + title: "Lazy mode", + value: !model._llama.supportsMmap + ? "mmap unsupported" + : model.useMmap === false + ? "mmap disabled" + : lazyMode === "auto" + ? "auto" + : lazyMode === true + ? "enabled" + : "disabled" }, { show: printBos, title: "BOS", diff --git a/src/evaluator/LlamaContext/LlamaContext.ts b/src/evaluator/LlamaContext/LlamaContext.ts index 6471ee42..c363cd10 100644 --- a/src/evaluator/LlamaContext/LlamaContext.ts +++ b/src/evaluator/LlamaContext/LlamaContext.ts @@ -293,6 +293,17 @@ export class LlamaContext { return this._totalSequences - this._nextGeneratedSequenceId + this._unusedSequenceIds.length; } + /** Assumed memory footprint of the context in bytes */ + public get memoryUsage(): { + ram: number, + vram: number + } { + return { + ram: this._ramConsumptionMarking?.size ?? 0, + vram: this._vramConsumptionMarking?.size ?? 0 + }; + } + /** * Before calling this method, make sure to call `sequencesLeft` to check if there are any sequences left. * When there are no sequences left, this method will throw an error. diff --git a/src/evaluator/LlamaModel/LlamaModel.ts b/src/evaluator/LlamaModel/LlamaModel.ts index 999019d5..533fa66a 100644 --- a/src/evaluator/LlamaModel/LlamaModel.ts +++ b/src/evaluator/LlamaModel/LlamaModel.ts @@ -110,6 +110,21 @@ export type LlamaModelOptions = { */ checkTensors?: boolean, + /** + * Lazily read tensors from the file on demand when they are needed, rather than loading all tensors upfront. + * Only works when mmap ({@link useMmap `useMmap`}) is enabled. + * + * This will cause the inference to potentially start slower the first time a tensor is accessed, + * but can significantly reduce the total amount of memory used by the model. + * + * - `true`: for supported tensors, read them on demand then they are needed + * - `"auto"`: for supported tensors, only read on demand ones that are larger than 4GiB + * - `false`: do not read tensors on demand, load all tensors upfront + * + * Defaults to `false`. + */ + lazyMode?: "auto" | boolean, + /** * Enable flash attention by default for contexts created with this model. * Only works with models that support flash attention. @@ -216,6 +231,7 @@ export class LlamaModel { /** @internal */ private readonly _fileInsights: GgufInsights; /** @internal */ private readonly _gpuLayers: number; /** @internal */ public readonly _useMmap: boolean; + /** @internal */ private readonly _lazyMode: "auto" | boolean; /** @internal */ private readonly _vocabOnly: boolean; /** @internal */ private readonly _filename?: string; /** @internal */ private readonly _disposedState: DisposedState = {disposed: false}; @@ -239,8 +255,8 @@ export class LlamaModel { public readonly onDispose = new EventRelay(); private constructor({ - modelPath, gpuLayers, vocabOnly = false, useMmap, useDirectIo, useMlock = false, checkTensors, onLoadProgress, loadSignal, - metadataOverrides + modelPath, gpuLayers, vocabOnly = false, useMmap, useDirectIo, useMlock = false, checkTensors, lazyMode, onLoadProgress, + loadSignal, metadataOverrides }: LlamaModelOptions & { gpuLayers: number, useMmap: boolean @@ -276,6 +292,13 @@ export class LlamaModel { this._gpuLayers = gpuLayers; this._useMmap = useMmap ?? false; this._vocabOnly = vocabOnly ?? false; + this._lazyMode = !useMmap + ? false + : lazyMode == null + ? false + : (typeof lazyMode === "boolean" || lazyMode === "auto") + ? lazyMode + : false; this._backendModelDisposeGuard = new DisposeGuard([this._llama._backendDisposeGuard]); this._llamaPreventDisposalHandle = this._llama._backendDisposeGuard.createPreventDisposalHandle(); this._defaultContextFlashAttentionOptionEnabled = _defaultContextFlashAttentionOptionEnabled; @@ -295,6 +318,7 @@ export class LlamaModel { ? useMlock : undefined, checkTensors: checkTensors ?? false, + lazyMode: this._lazyMode, onLoadProgress: onLoadProgress == null ? undefined : (loadPercentage: number) => { @@ -404,6 +428,10 @@ export class LlamaModel { return this._useMmap; } + public get lazyMode(): "auto" | boolean { + return this._lazyMode; + } + /** * Total model size in memory in bytes. * @@ -435,6 +463,17 @@ export class LlamaModel { return this._defaultContextKvCacheValueType; } + /** Assumed memory footprint of the model in bytes */ + public get memoryUsage(): { + ram: number, + vram: number + } { + return { + ram: this._ramConsumptionMarking?.size ?? 0, + vram: this._vramConsumptionMarking?.size ?? 0 + }; + } + /** * Transform text into tokens that can be fed to the model * @param text - the text to tokenize diff --git a/src/evaluator/LlamaRankingContext.ts b/src/evaluator/LlamaRankingContext.ts index e8a4ef88..63f6fa8a 100644 --- a/src/evaluator/LlamaRankingContext.ts +++ b/src/evaluator/LlamaRankingContext.ts @@ -63,12 +63,39 @@ export type LlamaRankingContextOptions = { ignoreMemorySafetyChecks?: boolean }; +export type RankingOptions = { + /** + * When the given document it too big that it exceeds the context size, this option determines how to handle it. + * + * - `"throw"`: throw an error + * - `"maxChunk"`: split the document into smaller chunks that would fit the context, rank all of them, + * and return the highest ranking score among the chunks. By default, sequential chunks will overlap by at least 50%. + * + * Default to `"throw"`. + */ + onOverflow?: "throw" | "maxChunk" | { + type: "throw" + } | { + type: "maxChunk", + + /** + * The percentage of overlap between sequential chunks when splitting the document. + * + * Defaults to `0.5`. + */ + overlapPercentage?: number + } +}; + +const defaultOverlapPercentage = 0.5; + /** * @see [Reranking Documents](https://node-llama-cpp.withcat.ai/guide/embedding#reranking) tutorial */ export class LlamaRankingContext { /** @internal */ private readonly _llamaContext: LlamaContext; /** @internal */ private readonly _template: string | undefined; + /** @internal */ private readonly _templateDocumentInstances?: number; /** @internal */ private readonly _sequence: LlamaContextSequence; /** @internal */ private readonly _disposeAggregator = new AsyncDisposeAggregator(); @@ -83,6 +110,9 @@ export class LlamaRankingContext { }) { this._llamaContext = _llamaContext; this._template = _template; + this._templateDocumentInstances = _template == null + ? undefined + : _template.split("{{document}}").length - 1; this._sequence = this._llamaContext.getSequence(); this._disposeAggregator.add( @@ -102,17 +132,25 @@ export class LlamaRankingContext { * A ranking score is a number between 0 and 1 representing the probability that the document is relevant to the query. * @returns a ranking score between 0 and 1 representing the probability that the document is relevant to the query. */ - public async rank(query: Token[] | string | LlamaText, document: Token[] | string | LlamaText) { - const resolvedInput = this._getEvaluationInput(query, document); + public async rank( + query: Token[] | string | LlamaText, + document: Token[] | string | LlamaText, + options?: RankingOptions + ): Promise { + const resolvedQuery = tokenizeInput(query, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); + const resolvedDocument = tokenizeInput(document, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); + const resolvedInput = this._chunkEvaluatedInputs(resolvedQuery, resolvedDocument, options); - if (resolvedInput.length > this._llamaContext.contextSize) + if (resolvedInput[0] == null) + throw new Error("Failed to generate a valid input for ranking."); + else if (resolvedInput.length === 1 && resolvedInput[0].length >= this._llamaContext.contextSize) throw new Error( "The input length exceed the context size. " + - `Try to increase the context size to at least ${resolvedInput.length + 1} ` + + `Try to increase the context size to at least ${resolvedInput[0].length + 1} ` + "or use another model that supports longer contexts." ); - return this._evaluateRankingForInput(resolvedInput); + return getMaxScore(await this._evaluateChunks(resolvedInput)); } /** @@ -121,21 +159,34 @@ export class LlamaRankingContext { * A ranking score is a number between 0 and 1 representing the probability that the document is relevant to the query. * @returns an array of ranking scores between 0 and 1 representing the probability that the document is relevant to the query. */ - public async rankAll(query: Token[] | string | LlamaText, documents: Array): Promise { - const resolvedTokens = documents.map((document) => this._getEvaluationInput(query, document)); - const maxInputTokensLength = resolvedTokens.reduce((max, tokens) => Math.max(max, tokens.length), 0); - - if (maxInputTokensLength > this._llamaContext.contextSize) + public async rankAll( + query: Token[] | string | LlamaText, + documents: Array, + options?: RankingOptions + ): Promise { + const resolvedQuery = tokenizeInput(query, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); + const resolvedInputs = documents.map((document) => this._chunkEvaluatedInputs( + resolvedQuery, + tokenizeInput(document, this._llamaContext.model.tokenizer, "trimLeadingSpace", false), + options + )); + const maxInputTokensLength = resolvedInputs.reduce((max, chunks) => ( + (chunks[0] == null || chunks.length !== 1) + ? max + : Math.max(max, chunks[0].length) + ), 0); + + if (maxInputTokensLength >= this._llamaContext.contextSize) throw new Error( "The input lengths of some of the given documents exceed the context size. " + `Try to increase the context size to at least ${maxInputTokensLength + 1} ` + "or use another model that supports longer contexts." ); - else if (resolvedTokens.length === 0) + else if (resolvedInputs.length === 0) return []; return await Promise.all( - resolvedTokens.map((tokens) => this._evaluateRankingForInput(tokens)) + resolvedInputs.map(async (chunks) => getMaxScore(await this._evaluateChunks(chunks))) ); } @@ -159,6 +210,13 @@ export class LlamaRankingContext { .sort((a, b) => b.score - a.score); } + /** Calculate the input length for a given query and document so you can determine whether it fits in the context size */ + public calculateInputLength(query: Token[] | string | LlamaText, document: Token[] | string | LlamaText) { + const resolvedQuery = tokenizeInput(query, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); + const resolvedDocument = tokenizeInput(document, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); + return this._getEvaluationInput(resolvedQuery, resolvedDocument).length; + } + public async dispose() { await this._disposeAggregator.dispose(); } @@ -176,17 +234,21 @@ export class LlamaRankingContext { return this._llamaContext.model; } + public get contextSize() { + return this._llamaContext.contextSize; + } + /** @internal */ - private _getEvaluationInput(query: Token[] | string | LlamaText, document: Token[] | string | LlamaText) { + private _getEvaluationInput(query: Token[], document: Token[]) { if (this._template != null) { const resolvedInput = splitText(this._template, ["{{query}}", "{{document}}"]) .flatMap((item) => { if (typeof item === "string") return this._llamaContext.model.tokenize(item, true, "trimLeadingSpace"); else if (item.separator === "{{query}}") - return tokenizeInput(query, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); + return query; else if (item.separator === "{{document}}") - return tokenizeInput(document, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); + return document; else void (item satisfies never); @@ -209,18 +271,15 @@ export class LlamaRankingContext { if (this.model.tokens.eos == null && this.model.tokens.sep == null) throw new Error("Computing rankings is not supported for this model."); - const resolvedQuery = tokenizeInput(query, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); - const resolvedDocument = tokenizeInput(document, this._llamaContext.model.tokenizer, "trimLeadingSpace", false); - - if (resolvedQuery.length === 0 && resolvedDocument.length === 0) + if (query.length === 0 && document.length === 0) return []; const resolvedInput = [ ...(this.model.tokens.bos == null ? [] : [this.model.tokens.bos]), - ...resolvedQuery, + ...query, ...(this.model.tokens.eos == null ? [] : [this.model.tokens.eos]), ...(this.model.tokens.sep == null ? [] : [this.model.tokens.sep]), - ...resolvedDocument, + ...document, ...(this.model.tokens.eos == null ? [] : [this.model.tokens.eos]) ]; @@ -257,6 +316,48 @@ export class LlamaRankingContext { }); } + private _evaluateChunks(input: Token[][]): Promise { + return Promise.all(input.map((chunk) => this._evaluateRankingForInput(chunk))); + } + + /** @internal */ + private _chunkEvaluatedInputs(query: Token[], document: Token[], options?: RankingOptions): Token[][] { + const fullInput = this._getEvaluationInput(query, document); + if (fullInput.length < this._llamaContext.contextSize || options?.onOverflow == null || options?.onOverflow === "throw" || ( + typeof options?.onOverflow === "object" && options.onOverflow.type === "throw" + )) + return [fullInput]; + + const templateDocumentInstances = Math.max(1, this._templateDocumentInstances ?? 1); + const maxChunkSize = Math.floor( + ( + this._llamaContext.contextSize - (fullInput.length - (document.length * templateDocumentInstances)) - 1 + ) / templateDocumentInstances + ); + if (maxChunkSize <= 0) + throw new Error("The document is too long to fit into the context window with the given query"); + + const overlapPercentage = (typeof options?.onOverflow === "object" && options.onOverflow.type === "maxChunk" && options.onOverflow.overlapPercentage != null) + ? Math.min(1, Math.max(0, Math.min(1, options.onOverflow.overlapPercentage))) + : defaultOverlapPercentage; + + const overlapTokens = Math.min(Math.max(0, Math.ceil(maxChunkSize * overlapPercentage)), maxChunkSize - 1); + + const result: Token[][] = []; + let start = 0; + while (start < document.length) { + const end = Math.min(start + maxChunkSize, document.length); + result.push(this._getEvaluationInput(query, document.slice(start, end))); + + if (end === document.length) + break; + + start = Math.min(end - overlapTokens, document.length - maxChunkSize); + } + + return result; + } + /** @internal */ private get _currentArchRankingAlreadyNormalized() { const architecture = this.model.fileInfo.metadata?.general?.architecture; @@ -316,3 +417,13 @@ export class LlamaRankingContext { function logitToSigmoid(logit: number) { return 1 / (1 + Math.exp(-logit)); } + +function getMaxScore(arr: number[]) { + let max = 0; + for (const num of arr) { + if (num > max) { + max = num; + } + } + return max; +} diff --git a/src/gguf/fileReaders/GgufFsFileReader.ts b/src/gguf/fileReaders/GgufFsFileReader.ts index 9080083f..6640d697 100644 --- a/src/gguf/fileReaders/GgufFsFileReader.ts +++ b/src/gguf/fileReaders/GgufFsFileReader.ts @@ -41,7 +41,7 @@ export class GgufFsFileReader extends GgufFileReader { const readOffset = GgufReadOffset.resolveReadOffset(offset); const endOffset = readOffset.offset + length; - if (endOffset >= this._buffer.length) + if (endOffset > this._buffer.length) return this._readToExpandBufferUpToOffset(endOffset) .then(() => { if (endOffset >= this._buffer.length) @@ -53,7 +53,7 @@ export class GgufFsFileReader extends GgufFileReader { private async _readToExpandBufferUpToOffset(endOffset: number, extraAllocationSize: number = defaultExtraAllocationSize) { return await withLock([this as GgufFsFileReader, "modifyBuffer"], this._signal, async () => { - if (endOffset < this._buffer.length) + if (endOffset <= this._buffer.length) return; const missingBytesBuffer = await this._readByteRange( diff --git a/src/gguf/fileReaders/GgufNetworkFetchFileReader.ts b/src/gguf/fileReaders/GgufNetworkFetchFileReader.ts index 11113420..0a6d0c12 100644 --- a/src/gguf/fileReaders/GgufNetworkFetchFileReader.ts +++ b/src/gguf/fileReaders/GgufNetworkFetchFileReader.ts @@ -21,6 +21,7 @@ export class GgufNetworkFetchFileReader extends GgufFileReader { public readonly headers: Record; public readonly tokens?: ModelFileAccessTokens; public readonly endpoints?: ModelDownloadEndpoints; + private _fileSize?: number; private readonly _signal?: AbortSignal; private _tryHeaders: Record[] | undefined = undefined; @@ -55,7 +56,7 @@ export class GgufNetworkFetchFileReader extends GgufFileReader { const readOffset = GgufReadOffset.resolveReadOffset(offset); const endOffset = readOffset.offset + length; - if (endOffset >= this._buffer.length) + if (endOffset > this._buffer.length) return this._fetchToExpandBufferUpToOffset(endOffset) .then(() => { if (endOffset >= this._buffer.length) @@ -67,7 +68,7 @@ export class GgufNetworkFetchFileReader extends GgufFileReader { private async _fetchToExpandBufferUpToOffset(endOffset: number, extraAllocationSize: number = defaultExtraAllocationSize) { await withLock([this as GgufNetworkFetchFileReader, "modifyBuffer"], this._signal, async () => { - if (endOffset < this._buffer.length) + if (endOffset <= this._buffer.length) return; const missingBytesBuffer = await retry(async (bail) => { @@ -77,6 +78,9 @@ export class GgufNetworkFetchFileReader extends GgufFileReader { if (this._signal?.aborted) { bail(this._signal.reason); throw this._signal.reason; + } else if (err instanceof FetchError && !err.canRetry) { + bail(err); + throw err; } throw err; @@ -96,6 +100,9 @@ export class GgufNetworkFetchFileReader extends GgufFileReader { const headersToTry = [this.headers, ...this._tryHeaders]; + if (this._fileSize != null && start >= this._fileSize) + throw new FetchError(`Requested byte range starting at index ${start} exceeds the file size of ${this._fileSize}`, false); + while (headersToTry.length > 0) { const headers = headersToTry.shift(); @@ -108,11 +115,20 @@ export class GgufNetworkFetchFileReader extends GgufFileReader { signal: this._signal }); - if ((response.status >= 500 || response.status === 429 || response.status === 401) && headersToTry.length > 0) + const technicalIssue = response.status >= 500 || response.status === 429 || response.status === 403; + const cannotAccess = (response.status >= 400 && response.status <= 402) || response.status === 404; + if (headersToTry.length > 0 && (technicalIssue || cannotAccess)) continue; if (!response.ok) - throw new Error(`Failed to fetch byte range: ${response.status} ${response.statusText}`); + throw new FetchError(`Failed to fetch byte range: ${response.status} ${response.statusText}`, technicalIssue); + + const fileSizeHeader = response.headers.get("content-range")?.split("/")[1] ?? response.headers.get("x-linked-size"); + if (fileSizeHeader != null) { + const fileSize = Number(fileSizeHeader); + if (Number.isSafeInteger(fileSize) && fileSize >= 0 && (this._fileSize == null || fileSize > this._fileSize)) + this._fileSize = fileSize; + } const arrayBuffer = await response.arrayBuffer(); return Buffer.from(arrayBuffer); @@ -121,3 +137,18 @@ export class GgufNetworkFetchFileReader extends GgufFileReader { throw new Error("Failed to fetch byte range: no more headers to try"); } } + +class FetchError extends Error { + public readonly canRetry: boolean; + + public constructor(message: string, canRetry: boolean) { + super(message); + this.canRetry = canRetry; + + Object.defineProperty(this, "canRetry" satisfies keyof this, { + enumerable: false, + configurable: false, + value: canRetry + }); + } +} diff --git a/src/gguf/insights/GgufInsights.ts b/src/gguf/insights/GgufInsights.ts index ff146620..534f4101 100644 --- a/src/gguf/insights/GgufInsights.ts +++ b/src/gguf/insights/GgufInsights.ts @@ -158,40 +158,22 @@ export class GgufInsights { return true; } - public get isRecurrent() { - // source: `llm_arch_is_recurrent` in `llama-arch.cpp` - switch (this._ggufFileInfo.metadata?.general?.architecture) { - case GgufArchitectureType.mamba: - case GgufArchitectureType.mamba2: - case GgufArchitectureType.rwkv6: - case GgufArchitectureType.rwkv6qwen2: - case GgufArchitectureType.rwkv7: - case GgufArchitectureType.arwkv7: - return true; - } + public get isRecurrent(): boolean { + return this._llama._bindings.getIsArchRecurrent(this._ggufFileInfo.metadata?.general?.architecture ?? "") ?? false; + } - return false; + public get isHybrid(): boolean { + return this._llama._bindings.getIsArchHybrid(this._ggufFileInfo.metadata?.general?.architecture ?? "") ?? false; } - public get isHybrid() { - // source: `llm_arch_is_hybrid` in `llama-arch.cpp` - switch (this._ggufFileInfo.metadata?.general?.architecture) { - case GgufArchitectureType.jamba: - case GgufArchitectureType.falconH1: - case GgufArchitectureType.plamo2: - case GgufArchitectureType.granitehybrid: - case GgufArchitectureType.lfm2: - case GgufArchitectureType.lfm2moe: - case GgufArchitectureType.nemotronH: - case GgufArchitectureType.nemotronHMoe: - case GgufArchitectureType.qwen3next: - case GgufArchitectureType.kimiLinear: - case GgufArchitectureType.qwen35: - case GgufArchitectureType.qwen35moe: - return true; - } + public get isSupportedByLlamaCpp(): boolean { + return this._llama._bindings.getIsArchSupported(this._ggufFileInfo.metadata?.general?.architecture ?? ""); + } - return false; + public get hasMtpWeights() { + const predictLayers = this._ggufFileInfo.architectureMetadata?.nextn_predict_layers ?? 0; + + return typeof predictLayers === "number" && predictLayers > 0 && predictLayers < this.totalLayers; } /** diff --git a/src/gguf/types/GgufMetadataTypes.ts b/src/gguf/types/GgufMetadataTypes.ts index 18040122..e38c6a9e 100644 --- a/src/gguf/types/GgufMetadataTypes.ts +++ b/src/gguf/types/GgufMetadataTypes.ts @@ -399,6 +399,7 @@ export type GgufMetadataDefaultArchitectureType = { readonly pooling_type?: GgufMetadataArchitecturePoolingType, readonly logit_scale?: number, readonly token_shift_count?: number, + readonly nextn_predict_layers?: number, readonly attention?: { readonly head_count?: number, diff --git a/src/gguf/utils/ggufQuantNames.ts b/src/gguf/utils/ggufQuantNames.ts index e3c9f453..4370bbb8 100644 --- a/src/gguf/utils/ggufQuantNames.ts +++ b/src/gguf/utils/ggufQuantNames.ts @@ -44,6 +44,7 @@ export const ggufQuantNames = new Map([ ]); export const ggufFileQuantNamesSet = new Set([ ...ggufQuantNames.keys(), + "MXFP4", "Q2_K_XL", "Q3_K_XL", "Q4_K_XL", diff --git a/src/index.ts b/src/index.ts index 541d15a3..8956de83 100644 --- a/src/index.ts +++ b/src/index.ts @@ -17,7 +17,7 @@ import {LlamaGrammarEvaluationState, LlamaGrammarEvaluationStateOptions} from ". import {LlamaContext, LlamaContextSequence} from "./evaluator/LlamaContext/LlamaContext.js"; import {LlamaEmbeddingContext, type LlamaEmbeddingContextOptions} from "./evaluator/LlamaEmbeddingContext.js"; import {LlamaEmbedding, type LlamaEmbeddingOptions, type LlamaEmbeddingJSON} from "./evaluator/LlamaEmbedding.js"; -import {LlamaRankingContext, type LlamaRankingContextOptions} from "./evaluator/LlamaRankingContext.js"; +import {LlamaRankingContext, type LlamaRankingContextOptions, type RankingOptions} from "./evaluator/LlamaRankingContext.js"; import { type LlamaContextOptions, type SequenceEvaluateOptions, type BatchingOptions, type LlamaContextSequenceRepeatPenalty, type CustomBatchingDispatchSchedule, type CustomBatchingPrioritizationStrategy, type BatchItem, type PrioritizedBatchItem, @@ -181,6 +181,7 @@ export { type LlamaEmbeddingJSON, LlamaRankingContext, type LlamaRankingContextOptions, + type RankingOptions, LlamaChatSession, defineChatSessionFunction, type LlamaChatSessionOptions, diff --git a/test/modelDependent/bgeReranker/rank.test.ts b/test/modelDependent/bgeReranker/rank.test.ts index 7ebc731d..3ad7003d 100644 --- a/test/modelDependent/bgeReranker/rank.test.ts +++ b/test/modelDependent/bgeReranker/rank.test.ts @@ -114,6 +114,122 @@ describe("bgeReranker", () => { `); }); + describe("overflow", () => { + const contextSize = 256; + const query = "Tell me a geographical fact"; + const parisIntroduction = "The volunteers arrived early to prepare the community hall for an evening reading club. They unfolded the tables, checked that every chair was steady, and moved the spare furniture into a storage cupboard. One person wiped the windows while another tested the lamps beside the comfortable armchairs. A box of donated books waited near the entrance, with handwritten notes explaining who had brought each one. The organizers sorted the books by size so that the shelves would look tidy, then made small paper labels for the empty spaces. In the kitchen, two helpers washed the cups and counted the spoons. They prepared a tray of biscuits, filled a jug with water, and put clean towels beside the sink. Someone noticed that a table leg was loose and fetched a screwdriver from the cupboard. Once the repair was finished, the group covered the tables with plain cloths and set out pencils and blank cards for visitors to write recommendations. The first guests arrived carrying coats and shopping bags. They chose seats, introduced themselves, and talked about how they found time to read during a busy week. One guest preferred reading before breakfast, while another listened to stories while washing dishes. A child drew a picture on a spare card and asked whether it could be used as a bookmark. The organizer agreed and found a ribbon to tie through a hole in the paper. When everyone had settled, the host invited each reader to share something from a book. After several people discussed recipes and craft projects, the final guest opened an encyclopedia and read a fact aloud."; + const everestIntroduction = "I spent the morning clearing a cupboard that had become difficult to close. First I carried the boxes into the living room and laid an old sheet over the carpet. The largest box contained tangled cables, spare buttons, and instruction booklets for appliances I no longer owned. I checked each cable, wound the useful ones neatly, and put them in a small basket. The buttons went into a glass jar beside my sewing kit. Under the box I found a wooden frame with a loose corner, so I cleaned the joints and applied a little glue. While it dried, I sorted a pile of notebooks into used and unused pages. Some contained shopping lists, others had sketches of furniture I once planned to build. I kept the sketches and placed the blank paper in a drawer for future notes. A tin of pencils needed sharpening, and several pens had dried out completely. By lunchtime the floor was covered with small groups of objects, each waiting for a proper place. I made a sandwich and ate it at the kitchen table before returning to the work. In the afternoon I lined the cupboard shelves with clean paper and measured the space available for baskets. The lighter boxes went on the upper shelf, with tools and household supplies below. I wrote labels on pieces of card and attached them with string so that everything would be easier to find. Finally, I vacuumed the carpet and folded the sheet away. Only a forgotten quiz book remained on the sofa. I sat down to read it and discovered the answer to one of its questions."; + const documents = [ + "The sky is clear and blue today. A few white clouds drift slowly above the houses while sunlight fills the garden. People open their windows and enjoy the warm afternoon outside.", + "Making pizza starts with kneading dough and letting it rise. The cook spreads tomato sauce over the base, adds cheese and vegetables, and bakes everything until the crust is crisp and golden.", + "Dogs love to play fetch with their owners. A ball thrown across the garden sends an excited dog running through the grass. After returning the ball, the dog waits eagerly for another throw.", + parisIntroduction + " Paris is the capital of France and stands on the banks of the Seine. The river flows through the city beneath many bridges, connecting neighborhoods with museums, parks, shops, and historic buildings.", + "After walking around the neighborhood, I stopped at home for a glass of water. I filled a bottle for the rest of the afternoon and placed it beside my backpack near the front door.", + everestIntroduction + " Mount Everest is the highest mountain above sea level in the world. It belongs to the Himalayan mountain range and lies on the border between Nepal and China. Snow and ice cover its upper slopes.", + "A warm cup of tea is pleasant on a cold winter day. I boil water, let the tea leaves steep, and carry the cup to a comfortable chair where I can read a book.", + "Painting is a form of creative expression. An artist mixes colors on a palette before applying them to a canvas. Different brushes create broad areas of color, delicate lines, and interesting textures in the picture.", + "Not everything that shines is made of gold. A shop window can display polished brass, colored glass, and silver jewelry beside golden objects. Their bright surfaces look similar even though the materials are different.", + "Cleaning the house begins with putting scattered objects back in their places. I dust the shelves, sweep the floor, and wash the dishes. Opening a window lets fresh air into the newly tidy room." + ]; + + test("rank", {timeout: 1000 * 60 * 60 * 2}, async (test) => { + if (process.platform !== "darwin" && process.arch !== "arm64") + test.skip(); // the scores are a bit different on different platforms, so skipping on other platforms due to flakiness + + const modelPath = await getModelFile("bge-reranker-v2-m3-Q8_0.gguf"); + const llama = await getTestLlama(); + + const model = await llama.loadModel({ + modelPath + }); + const rankingContext = await model.createRankingContext({ + contextSize + }); + + for (const introduction of [parisIntroduction, everestIntroduction]) + expect(rankingContext.calculateInputLength(query, introduction)).toBeGreaterThan(contextSize); + + for (const index of [3, 5]) { + const document = documents[index]!; + expect(rankingContext.calculateInputLength(query, document)).toBeGreaterThan(contextSize); + await expect(rankingContext.rank(query, document)).rejects.toThrow("exceed the context size"); + } + + const ranks = await Promise.all( + documents.map((document) => rankingContext.rank(query, document, {onOverflow: "maxChunk"})) + ); + + const firstChunkSize = contextSize - rankingContext.calculateInputLength(query, []) - 1; + for (const index of [3, 5]) { + const firstChunk = model.tokenize(documents[index]!, false, "trimLeadingSpace").slice(0, firstChunkSize); + expect(ranks[index]).toBeGreaterThan(await rankingContext.rank(query, firstChunk)); + } + + expect(ranks).toHaveLength(documents.length); + expect(simplifyRanks(ranks)).toMatchInlineSnapshot(` + [ + 0.00002039908727992137, + 0.00002039908727992137, + 0.00002039908727992137, + 0.026596993576865856, + 0.00002039908727992137, + 0.08317269649392238, + 0.00002039908727992137, + 0.00003716893710288947, + 0.00002039908727992137, + 0.00002039908727992137, + ] + `); + }); + + test("rank all", {timeout: 1000 * 60 * 60 * 2}, async (test) => { + if (process.platform !== "darwin" && process.arch !== "arm64") + test.skip(); // the scores are a bit different on different platforms, so skipping on other platforms due to flakiness + + const modelPath = await getModelFile("bge-reranker-v2-m3-Q8_0.gguf"); + const llama = await getTestLlama(); + + const model = await llama.loadModel({ + modelPath + }); + const rankingContext = await model.createRankingContext({ + contextSize + }); + + for (const introduction of [parisIntroduction, everestIntroduction]) + expect(rankingContext.calculateInputLength(query, introduction)).toBeGreaterThan(contextSize); + + for (const index of [3, 5]) + expect(rankingContext.calculateInputLength(query, documents[index]!)).toBeGreaterThan(contextSize); + + await expect(rankingContext.rankAll(query, documents)).rejects.toThrow("exceed the context size"); + + const ranks = await rankingContext.rankAll(query, documents, {onOverflow: "maxChunk"}); + + const firstChunkSize = contextSize - rankingContext.calculateInputLength(query, []) - 1; + for (const index of [3, 5]) { + const firstChunk = model.tokenize(documents[index]!, false, "trimLeadingSpace").slice(0, firstChunkSize); + expect(ranks[index]).toBeGreaterThan(await rankingContext.rank(query, firstChunk)); + } + + expect(ranks).toHaveLength(documents.length); + expect(simplifyRanks(ranks)).toMatchInlineSnapshot(` + [ + 0.00002039908727992137, + 0.00002039908727992137, + 0.00002039908727992137, + 0.026596993576865856, + 0.00002039908727992137, + 0.08317269649392238, + 0.00002039908727992137, + 0.00003716893710288947, + 0.00002039908727992137, + 0.00002039908727992137, + ] + `); + }); + }); + test("rank and sort", {timeout: 1000 * 60 * 60 * 2}, async (test) => { if (process.platform !== "darwin" && process.arch !== "arm64") test.skip(); // the scores are a bit different on different platforms, so skipping on other platforms due to flakiness diff --git a/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts b/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts index 72c6f31d..f0ec553e 100644 --- a/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts +++ b/test/modelDependent/functionary/functionaryModelGpuLayersOptions.test.ts @@ -1549,8 +1549,8 @@ describe("functionary", () => { }); expect(res.gpuLayers).to.be.gte(16); expect(res.gpuLayers).to.be.lte(24); - expect(res.gpuLayers).to.toMatchInlineSnapshot("18"); - expect(res.contextSize).to.toMatchInlineSnapshot("6144"); + expect(res.gpuLayers).to.toMatchInlineSnapshot("17"); + expect(res.contextSize).to.toMatchInlineSnapshot("8192"); expect(res.useMmap).to.toMatchInlineSnapshot("false"); } }); @@ -1593,8 +1593,8 @@ describe("functionary", () => { totalRam: s1GB * 8, freeRam: s1GB * 8 }); - expect(res.gpuLayers).to.toMatchInlineSnapshot("0"); - expect(res.contextSize).to.toMatchInlineSnapshot("8192"); + expect(res.gpuLayers).to.toMatchInlineSnapshot("9"); + expect(res.contextSize).to.toMatchInlineSnapshot("7424"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); expect(res.contextSize).to.be.gte(contextSize); } diff --git a/test/modelDependent/stableCode/stableCodeModelGpuLayersOptions.test.ts b/test/modelDependent/stableCode/stableCodeModelGpuLayersOptions.test.ts index 9a6ee495..ddc20ba2 100644 --- a/test/modelDependent/stableCode/stableCodeModelGpuLayersOptions.test.ts +++ b/test/modelDependent/stableCode/stableCodeModelGpuLayersOptions.test.ts @@ -116,7 +116,7 @@ describe("stableCode", () => { freeVram: s1GB * 3 }); expect(res.gpuLayers).to.eql(16); - expect(res.contextSize).to.toMatchInlineSnapshot("13824"); + expect(res.contextSize).to.toMatchInlineSnapshot("13312"); expect(res.useMmap).to.toMatchInlineSnapshot("false"); } try { @@ -183,7 +183,7 @@ describe("stableCode", () => { freeVram: s1GB * 6 }); expect(res.gpuLayers).to.eql(32); - expect(res.contextSize).to.toMatchInlineSnapshot("14080"); + expect(res.contextSize).to.toMatchInlineSnapshot("13824"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } try { @@ -236,7 +236,7 @@ describe("stableCode", () => { freeVram: s1GB * 6 }); expect(res.gpuLayers).to.eql(33); - expect(res.contextSize).to.toMatchInlineSnapshot("13312"); + expect(res.contextSize).to.toMatchInlineSnapshot("13056"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } try { @@ -329,7 +329,7 @@ describe("stableCode", () => { freeVram: s1GB * 4 }); expect(res.gpuLayers).to.eql(33); - expect(res.contextSize).to.toMatchInlineSnapshot("6912"); + expect(res.contextSize).to.toMatchInlineSnapshot("6656"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { @@ -338,7 +338,7 @@ describe("stableCode", () => { freeVram: s1GB * 4.4 }); expect(res.gpuLayers).to.eql(33); - expect(res.contextSize).to.toMatchInlineSnapshot("8192"); + expect(res.contextSize).to.toMatchInlineSnapshot("7936"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { @@ -347,7 +347,7 @@ describe("stableCode", () => { freeVram: s1GB * 4.8 }); expect(res.gpuLayers).to.eql(33); - expect(res.contextSize).to.toMatchInlineSnapshot("9472"); + expect(res.contextSize).to.toMatchInlineSnapshot("9216"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } }); @@ -377,7 +377,7 @@ describe("stableCode", () => { freeVram: s1GB * 0.8 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("5"); - expect(res.contextSize).to.toMatchInlineSnapshot("9984"); + expect(res.contextSize).to.toMatchInlineSnapshot("8960"); expect(res.useMmap).to.toMatchInlineSnapshot("false"); } { @@ -386,7 +386,7 @@ describe("stableCode", () => { freeVram: s1GB * 1.4 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("10"); - expect(res.contextSize).to.toMatchInlineSnapshot("8192"); + expect(res.contextSize).to.toMatchInlineSnapshot("7936"); expect(res.useMmap).to.toMatchInlineSnapshot("false"); } { @@ -440,7 +440,7 @@ describe("stableCode", () => { freeVram: s1GB * 4 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("33"); - expect(res.contextSize).to.toMatchInlineSnapshot("6912"); + expect(res.contextSize).to.toMatchInlineSnapshot("6656"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { @@ -449,7 +449,7 @@ describe("stableCode", () => { freeVram: s1GB * 4.3 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("33"); - expect(res.contextSize).to.toMatchInlineSnapshot("7936"); + expect(res.contextSize).to.toMatchInlineSnapshot("7680"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { @@ -467,7 +467,7 @@ describe("stableCode", () => { freeVram: s1GB * 4.8 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("33"); - expect(res.contextSize).to.toMatchInlineSnapshot("9472"); + expect(res.contextSize).to.toMatchInlineSnapshot("9216"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { @@ -476,7 +476,7 @@ describe("stableCode", () => { freeVram: s1GB * 5.2 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("33"); - expect(res.contextSize).to.toMatchInlineSnapshot("10752"); + expect(res.contextSize).to.toMatchInlineSnapshot("10496"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { @@ -485,7 +485,7 @@ describe("stableCode", () => { freeVram: s1GB * 5.8 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("33"); - expect(res.contextSize).to.toMatchInlineSnapshot("12800"); + expect(res.contextSize).to.toMatchInlineSnapshot("12544"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { @@ -494,7 +494,7 @@ describe("stableCode", () => { freeVram: s1GB * 6 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("33"); - expect(res.contextSize).to.toMatchInlineSnapshot("13312"); + expect(res.contextSize).to.toMatchInlineSnapshot("13056"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } }); @@ -562,7 +562,7 @@ describe("stableCode", () => { }); expect(res.gpuLayers).to.be.gte(16); expect(res.gpuLayers).to.toMatchInlineSnapshot("33"); - expect(res.contextSize).to.toMatchInlineSnapshot("6912"); + expect(res.contextSize).to.toMatchInlineSnapshot("6656"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { @@ -572,9 +572,9 @@ describe("stableCode", () => { }); expect(res.gpuLayers).to.be.gte(16); expect(res.gpuLayers).to.be.lte(24); - expect(res.gpuLayers).to.toMatchInlineSnapshot("20"); - expect(res.contextSize).to.toMatchInlineSnapshot("15360"); - expect(res.useMmap).to.toMatchInlineSnapshot("false"); + expect(res.gpuLayers).to.toMatchInlineSnapshot("24"); + expect(res.contextSize).to.toMatchInlineSnapshot("9728"); + expect(res.useMmap).to.toMatchInlineSnapshot("true"); } { const res = await resolveGpuLayers({min: 16, max: 24}, { @@ -584,7 +584,7 @@ describe("stableCode", () => { expect(res.gpuLayers).to.be.gte(16); expect(res.gpuLayers).to.be.lte(24); expect(res.gpuLayers).to.toMatchInlineSnapshot("22"); - expect(res.contextSize).to.toMatchInlineSnapshot("8448"); + expect(res.contextSize).to.toMatchInlineSnapshot("8192"); expect(res.useMmap).to.toMatchInlineSnapshot("false"); } }); @@ -609,7 +609,7 @@ describe("stableCode", () => { freeVram: s1GB * 4 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("33"); - expect(res.contextSize).to.toMatchInlineSnapshot("6912"); + expect(res.contextSize).to.toMatchInlineSnapshot("6656"); expect(res.useMmap).to.toMatchInlineSnapshot("true"); expect(res.contextSize).to.be.gte(contextSize); } @@ -631,7 +631,7 @@ describe("stableCode", () => { freeVram: s1GB * 4 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("28"); - expect(res.contextSize).to.toMatchInlineSnapshot("9216"); + expect(res.contextSize).to.toMatchInlineSnapshot("8960"); expect(res.useMmap).to.toMatchInlineSnapshot("false"); expect(res.contextSize).to.be.gte(contextSize); } @@ -642,7 +642,7 @@ describe("stableCode", () => { freeVram: s1GB * 1 }); expect(res.gpuLayers).to.toMatchInlineSnapshot("6"); - expect(res.contextSize).to.toMatchInlineSnapshot("11008"); + expect(res.contextSize).to.toMatchInlineSnapshot("9984"); expect(res.useMmap).to.toMatchInlineSnapshot("false"); expect(res.contextSize).to.be.gte(contextSize); } diff --git a/test/standalone/chatWrappers/utils/jinjaTemplates.ts b/test/standalone/chatWrappers/utils/jinjaTemplates.ts index 1facd494..d5195c0d 100644 --- a/test/standalone/chatWrappers/utils/jinjaTemplates.ts +++ b/test/standalone/chatWrappers/utils/jinjaTemplates.ts @@ -3501,3 +3501,582 @@ export const museGlimmerJinjaTemplate = String.raw` {{- "<|start|>assistant" -}} {%- endif -%} `.slice(1, -1).replaceAll("\\`", "`"); + +export const museGlimmerJinjaTemplate2 = String.raw` +{%- macro render_content(content) -%} + {%- if content is string -%} + {{- content -}} + {%- elif content is not none -%} + {%- for part in content -%} + {%- if part["type"] == "image" -%} + {{- "<|patch|>" -}} + {%- elif part["type"] == "video" -%} + {{- "<|video|>" -}} + {%- elif part["type"] == "text" -%} + {{- part["text"] -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} +{%- endmacro -%} +{%- macro render_atem(tc) -%} + {%- set args = tc.function.arguments -%} + {%- if args is not mapping -%} + {{- raise_exception("Onyx ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.") -}} + {%- endif -%} + {{- "\n\n" -}} + {%- for (k, v) in args.items() -%} + {{- "" -}} + {%- if v is boolean -%} + {%- if v -%} + {{- "true" -}} + {%- else -%} + {{- "false" -}} + {%- endif -%} + {%- elif v is none -%} + {{- "null" -}} + {%- elif v is mapping or v is iterable and v is not string -%} + {{- v | tojson -}} + {%- else -%} + {{- v -}} + {%- endif -%} + {{- "\n" -}} + {%- endfor -%} + {{- "\n" -}} +{%- endmacro -%} +{%- macro render_tool_defs(tools) -%} + {{- "In this environment you have access to a set of tools you can use to answer the user's question.\n\n" -}} + {{- "You can invoke a function by writing a \"\" block like the following:\n" -}} + {{- "\n\n$PARAMETER_VALUE\n...\n\n\n\n" -}} + {{- "String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n" -}} + {{- "Here are the functions available in JSONSchema format:\n" -}} + {{- "// Tool metadata\n" -}} + {%- set nsns = namespace(seen=[]) -%} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set tns = fn.name.split(".")[0] -%} + {%- if tns not in nsns.seen -%} + {%- set nsns.seen = nsns.seen + [tns] -%} + {%- endif -%} + {%- endfor -%} + {%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%} + {%- for tns in nsns.seen -%} + {{- "{\"name\": " + tns | tojson + ", \"description\": " + (nd[tns] if tns in nd else "") | tojson + "}\n" -}} + {%- endfor -%} + {{- "// Function schemas" -}} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {{- "\n{\"name\": " + fn.name | tojson + ", \"description\": " + fn.description | tojson + ", \"parameters\": " + fn.parameters | tojson + "}" -}} + {%- endfor -%} + {{- "\n\nHere's an example of how to call a function in the tool set:\n" -}} + {{- "(If the tool namespace is not specified, invoke the function directly as \`example_function_name\` rather than \`example_tool_name.example_function_name\`)\n\n" -}} + {{- "to=example_tool_name.example_function_name\n\n" -}} + {{- "\n\n" -}} + {{- "value_1\n" -}} + {{- "This is the value for the second parameter\nthat can span\n\"multiple\" lines\n\n" -}} + {{- "\n" -}} +{%- endmacro -%} +{%- macro render_reasoning() -%} + {%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else "high" -%} + {{- "Reasoning strength: " + rs + "." -}} +{%- endmacro -%} +{%- macro render_system_meta(tools) -%} + {%- set rns = namespace(recipients=["\"self\""], nslist=[]) -%} + {%- if tools -%} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set tns = fn.name.split(".")[0] -%} + {%- if tns not in rns.nslist -%} + {%- set rns.nslist = rns.nslist + [tns] -%} + {%- endif -%} + {%- endfor -%} + {%- for tns in rns.nslist -%} + {%- set rns.recipients = rns.recipients + ["\"" + tns + ".*\""] -%} + {%- endfor -%} + {%- endif -%} + {%- set rns.recipients = rns.recipients + ["\"user\""] -%} + {{- "# Valid recipients: " + rns.recipients | join(", ") + "." -}} +{%- endmacro -%} +{{- bos_token -}} +{%- set ns = namespace(has_system=false) -%} +{%- for m in messages -%} + {%- if m["role"] == "system" -%} + {%- set ns.has_system = true -%} + {%- endif -%} +{%- endfor -%} +{%- if not ns.has_system -%} + {{- "<|start|>system<|message|>You are a helpful AI assistant." -}} + {%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else "2026-01-04" -%} + {{- "\nKnowledge cutoff: " + kc + "." -}} + {%- if current_date is defined and current_date -%} + {{- "\nCurrent date: " + current_date + "." -}} + {%- elif strftime_now is defined -%} + {{- "\nCurrent date: " + strftime_now("%Y-%m-%d") + "." -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_reasoning() -}} + {%- if tools -%} + {{- "\n\n" -}} + {{- render_tool_defs(tools) -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_system_meta(tools) -}} + {{- "<|eot|>" -}} +{%- endif -%} +{%- for message in messages -%} + {%- set role = message["role"] -%} + {%- set end_token = "<|eom|>" if not loop.last and messages[loop.index0 + 1]["role"] == role else "<|eot|>" -%} + {%- if role == "system" -%} + {{- "<|start|>system<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "\n\n" -}} + {{- render_reasoning() -}} + {%- if tools -%} + {{- "\n\n" -}} + {{- render_tool_defs(tools) -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_system_meta(tools) -}} + {{- "<|eot|>" -}} + {%- elif role == "user" -%} + {{- "<|start|>user<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "<|eot|>" -}} + {%- elif role == "tool" -%} + {%- set tname = message.get("name") -%} + {%- if not tname -%} + {%- set tcid = message.get("tool_call_id") -%} + {%- set rns = namespace(name=tcid if tcid else "") -%} + {%- for m in messages -%} + {%- if m.get("tool_calls") -%} + {%- for tc in m["tool_calls"] -%} + {%- if tcid is not none and tc.id is defined and tc.id == tcid -%} + {%- set rns.name = tc.function.name -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endfor -%} + {%- set tname = rns.name -%} + {%- endif -%} + {{- "<|start|>tool " + tname + "<|message|>\n" -}} + {{- render_content(message["content"]) -}} + {{- "\n<|eot|>" -}} + {%- elif role == "assistant" -%} + {%- if message.get("reasoning_content") -%} + {{- "<|start|>assistant to=self<|message|>" + message["reasoning_content"] + "<|eom|>" -}} + {%- endif -%} + {%- if message.get("tool_calls") -%} + {%- for tc in message["tool_calls"] -%} + {{- "<|start|>assistant to=" + tc.function.name + "<|message|>" -}} + {{- render_atem(tc) -}} + {%- if loop.last -%} + {{- end_token -}} + {%- else -%} + {{- "<|eom|>" -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {%- set recipient = message.get("recipient") or "user" -%} + {%- set end_turn = message.get("end_turn") -%} + {%- if end_turn is none -%} + {%- set end_turn = not (recipient and recipient != "user") -%} + {%- endif -%} + {{- "<|start|>assistant" -}} + {%- if recipient -%} + {{- " to=" + recipient -}} + {%- endif -%} + {{- "<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "<|eot|>" if end_turn else "<|eom|>" -}} + {%- endif -%} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- "<|start|>assistant" -}} +{%- endif -%} +`.slice(1, -1).replaceAll("\\`", "`"); + +export const museGlimmerJinjaTemplate3 = String.raw` +{%- macro render_content(content) -%} + {%- if content is string -%} + {{- content -}} + {%- elif content is not none -%} + {%- for part in content -%} + {%- if part["type"] == "image" -%} + {{- "<|patch|>" -}} + {%- elif part["type"] == "video" -%} + {{- "<|video|>" -}} + {%- elif part["type"] == "text" -%} + {{- part["text"] -}} + {%- endif -%} + {%- endfor -%} + {%- endif -%} +{%- endmacro -%} +{%- macro render_atem(tc) -%} + {%- set args = tc.function.arguments -%} + {%- if args is not mapping -%} + {{- raise_exception("Onyx ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.") -}} + {%- endif -%} + {{- "\n\n" -}} + {%- for (k, v) in args.items() -%} + {{- "" -}} + {%- if v is boolean -%} + {%- if v -%} + {{- "true" -}} + {%- else -%} + {{- "false" -}} + {%- endif -%} + {%- elif v is none -%} + {{- "null" -}} + {%- elif v is mapping or v is iterable and v is not string -%} + {{- v | tojson -}} + {%- else -%} + {{- v -}} + {%- endif -%} + {{- "\n" -}} + {%- endfor -%} + {{- "\n" -}} +{%- endmacro -%} +{%- macro render_tool_defs(tools) -%} + {{- "In this environment you have access to a set of tools you can use to answer the user's question.\n\n" -}} + {{- "You can invoke a function by writing a \"\" block like the following:\n" -}} + {{- "\n\n$PARAMETER_VALUE\n...\n\n\n\n" -}} + {{- "String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\n" -}} + {{- "Here are the functions available in JSONSchema format:\n" -}} + {{- "// Tool metadata\n" -}} + {%- set nsns = namespace(seen=[]) -%} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set tns = fn.name.split(".")[0] -%} + {%- if tns not in nsns.seen -%} + {%- set nsns.seen = nsns.seen + [tns] -%} + {%- endif -%} + {%- endfor -%} + {%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%} + {%- for tns in nsns.seen -%} + {{- "{\"name\": " + tns | tojson + ", \"description\": " + (nd[tns] if tns in nd else "") | tojson + "}\n" -}} + {%- endfor -%} + {{- "// Function schemas" -}} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {{- "\n{\"name\": " + fn.name | tojson + ", \"description\": " + fn.description | tojson + ", \"parameters\": " + fn.parameters | tojson + "}" -}} + {%- endfor -%} + {{- "\n\nHere's an example of how to call a function in the tool set:\n" -}} + {{- "(If the tool namespace is not specified, invoke the function directly as \`example_function_name\` rather than \`example_tool_name.example_function_name\`)\n\n" -}} + {{- "to=example_tool_name.example_function_name\n\n" -}} + {{- "\n\n" -}} + {{- "value_1\n" -}} + {{- "This is the value for the second parameter\nthat can span\n\"multiple\" lines\n\n" -}} + {{- "\n" -}} +{%- endmacro -%} +{%- macro render_reasoning() -%} + {%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else "high" -%} + {{- "Reasoning strength: " + rs + "." -}} +{%- endmacro -%} +{%- macro render_system_meta(tools) -%} + {%- set rns = namespace(recipients=["\"self\""], nslist=[]) -%} + {%- if tools -%} + {%- for tool in tools -%} + {%- set fn = tool.function if tool.function is defined else tool -%} + {%- set tns = fn.name.split(".")[0] -%} + {%- if tns not in rns.nslist -%} + {%- set rns.nslist = rns.nslist + [tns] -%} + {%- endif -%} + {%- endfor -%} + {%- for tns in rns.nslist -%} + {%- set rns.recipients = rns.recipients + ["\"" + tns + ".*\""] -%} + {%- endfor -%} + {%- endif -%} + {%- set rns.recipients = rns.recipients + ["\"user\""] -%} + {{- "# Valid recipients: " + rns.recipients | join(", ") + "." -}} +{%- endmacro -%} +{{- bos_token -}} +{%- set ns = namespace(has_system=false) -%} +{%- for m in messages -%} + {%- if m["role"] == "system" -%} + {%- set ns.has_system = true -%} + {%- endif -%} +{%- endfor -%} +{%- if not ns.has_system -%} + {{- "<|start|>system<|message|>You are a helpful AI assistant." -}} + {%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else "2026-01-04" -%} + {{- "\nKnowledge cutoff: " + kc + "." -}} + {%- if current_date is defined and current_date -%} + {{- "\nCurrent date: " + current_date + "." -}} + {%- elif strftime_now is defined -%} + {{- "\nCurrent date: " + strftime_now("%Y-%m-%d") + "." -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_reasoning() -}} + {%- if tools -%} + {{- "\n\n" -}} + {{- render_tool_defs(tools) -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_system_meta(tools) -}} + {{- "<|eot|>" -}} +{%- endif -%} +{%- for message in messages -%} + {%- set role = message["role"] -%} + {%- set end_token = "<|eom|>" if not loop.last and messages[loop.index0 + 1]["role"] == role else "<|eot|>" -%} + {%- if role == "system" -%} + {{- "<|start|>system<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "\n\n" -}} + {{- render_reasoning() -}} + {%- if tools -%} + {{- "\n\n" -}} + {{- render_tool_defs(tools) -}} + {%- endif -%} + {{- "\n\n" -}} + {{- render_system_meta(tools) -}} + {{- "<|eot|>" -}} + {%- elif role == "user" -%} + {{- "<|start|>user<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "<|eot|>" -}} + {%- elif role == "tool" -%} + {%- set tname = message.get("name") -%} + {%- if not tname -%} + {%- set tcid = message.get("tool_call_id") -%} + {%- set rns = namespace(name=tcid if tcid else "") -%} + {%- for m in messages -%} + {%- if m.get("tool_calls") -%} + {%- for tc in m["tool_calls"] -%} + {%- if tcid is not none and tc.id is defined and tc.id == tcid -%} + {%- set rns.name = tc.function.name -%} + {%- endif -%} + {%- endfor -%} + {%- endif -%} + {%- endfor -%} + {%- set tname = rns.name -%} + {%- endif -%} + {{- "<|start|>tool " + tname + "<|message|>\n" -}} + {{- render_content(message["content"]) -}} + {{- "\n<|eot|>" -}} + {%- elif role == "assistant" -%} + {%- if message.get("reasoning_content") -%} + {{- "<|start|>assistant to=self<|message|>" + message["reasoning_content"] + "<|eom|>" -}} + {%- endif -%} + {%- if message.get("tool_calls") -%} + {%- for tc in message["tool_calls"] -%} + {{- "<|start|>assistant to=" + tc.function.name + "<|message|>" -}} + {{- render_atem(tc) -}} + {%- if loop.last -%} + {{- end_token -}} + {%- else -%} + {{- "<|eom|>" -}} + {%- endif -%} + {%- endfor -%} + {%- else -%} + {%- set recipient = message.get("recipient") or "user" -%} + {%- set end_turn = message.get("end_turn") -%} + {%- if end_turn is none -%} + {%- set end_turn = not (recipient and recipient != "user") -%} + {%- endif -%} + {{- "<|start|>assistant" -}} + {%- if recipient -%} + {{- " to=" + recipient -}} + {%- endif -%} + {{- "<|message|>" -}} + {{- render_content(message["content"]) -}} + {{- "<|eot|>" if end_turn else "<|eom|>" -}} + {%- endif -%} + {%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} + {{- "<|start|>assistant" -}} +{%- endif -%} +`.slice(1, -1).replaceAll("\\`", "`"); + +export const museGlimmerJinjaTemplate4 = ` +{%- macro render_content(content) -%} +\t{%- if content is string -%} +\t\t{{- content -}} +\t{%- elif content is not none -%} +\t\t{%- for part in content -%} +\t\t\t{%- if part["type"] == "image" -%} +\t\t\t\t{{- "<|patch|>" -}} +\t\t\t{%- elif part["type"] == "video" -%} +\t\t\t\t{{- "<|video|>" -}} +\t\t\t{%- elif part["type"] == "text" -%} +\t\t\t\t{{- part["text"] -}} +\t\t\t{%- endif -%} +\t\t{%- endfor -%} +\t{%- endif -%} +{%- endmacro -%} +{%- macro render_atem(tc) -%} +\t{%- set args = tc.function.arguments -%} +\t{%- if args is not mapping -%} +\t\t{{- raise_exception("Onyx ATEM chat template requires tool_call.function.arguments to be a dict (mapping); a JSON string cannot be parsed in the HF jinja sandbox.") -}} +\t{%- endif -%} +\t{{- "\\n\\n" -}} +\t{%- for (k, v) in args.items() -%} +\t\t{{- "" -}} +\t\t{%- if v is boolean -%} +\t\t\t{%- if v -%} +\t\t\t\t{{- "true" -}} +\t\t\t{%- else -%} +\t\t\t\t{{- "false" -}} +\t\t\t{%- endif -%} +\t\t{%- elif v is none -%} +\t\t\t{{- "null" -}} +\t\t{%- elif v is mapping or v is iterable and v is not string -%} +\t\t\t{{- v | tojson -}} +\t\t{%- else -%} +\t\t\t{{- v -}} +\t\t{%- endif -%} +\t\t{{- "\\n" -}} +\t{%- endfor -%} +\t{{- "\\n" -}} +{%- endmacro -%} +{%- macro render_tool_defs(tools) -%} +\t{{- "In this environment you have access to a set of tools you can use to answer the user's question.\\n\\n" -}} +\t{{- "You can invoke a function by writing a \\"\\" block like the following:\\n" -}} +\t{{- "\\n\\n$PARAMETER_VALUE\\n...\\n\\n\\n\\n" -}} +\t{{- "String and scalar parameters should be specified as is, while lists and objects should use JSON format. Note that spaces for string values are not stripped. The output is not expected to be valid XML and is parsed with regular expressions.\\n" -}} +\t{{- "Here are the functions available in JSONSchema format:\\n" -}} +\t{{- "// Tool metadata\\n" -}} +\t{%- set nsns = namespace(seen=[]) -%} +\t{%- for tool in tools -%} +\t\t{%- set fn = tool.function if tool.function is defined else tool -%} +\t\t{%- set tns = fn.name.split(".")[0] -%} +\t\t{%- if tns not in nsns.seen -%} +\t\t\t{%- set nsns.seen = nsns.seen + [tns] -%} +\t\t{%- endif -%} +\t{%- endfor -%} +\t{%- set nd = tool_namespace_descriptions if tool_namespace_descriptions is defined else {} -%} +\t{%- for tns in nsns.seen -%} +\t\t{{- "{\\"name\\": " + tns | tojson + ", \\"description\\": " + (nd[tns] if tns in nd else "") | tojson + "}\\n" -}} +\t{%- endfor -%} +\t{{- "// Function schemas" -}} +\t{%- for tool in tools -%} +\t\t{%- set fn = tool.function if tool.function is defined else tool -%} +\t\t{{- "\\n{\\"name\\": " + fn.name | tojson + ", \\"description\\": " + fn.description | tojson + ", \\"parameters\\": " + fn.parameters | tojson + "}" -}} +\t{%- endfor -%} +\t{{- "\\n\\nHere's an example of how to call a function in the tool set:\\n" -}} +\t{{- "(If the tool namespace is not specified, invoke the function directly as \`example_function_name\` rather than \`example_tool_name.example_function_name\`)\\n\\n" -}} +\t{{- "to=example_tool_name.example_function_name\\n\\n" -}} +\t{{- "\\n\\n" -}} +\t{{- "value_1\\n" -}} +\t{{- "This is the value for the second parameter\\nthat can span\\n\\"multiple\\" lines\\n\\n" -}} +\t{{- "\\n" -}} +{%- endmacro -%} +{%- macro render_reasoning() -%} +\t{%- set rs = reasoning_strength if reasoning_strength is defined and reasoning_strength else "high" -%} +\t{{- "Reasoning strength: " + rs + "." -}} +{%- endmacro -%} +{%- macro render_system_meta(tools) -%} +\t{%- set rns = namespace(recipients=["\\"self\\""], nslist=[]) -%} +\t{%- if tools -%} +\t\t{%- for tool in tools -%} +\t\t\t{%- set fn = tool.function if tool.function is defined else tool -%} +\t\t\t{%- set tns = fn.name.split(".")[0] -%} +\t\t\t{%- if tns not in rns.nslist -%} +\t\t\t\t{%- set rns.nslist = rns.nslist + [tns] -%} +\t\t\t{%- endif -%} +\t\t{%- endfor -%} +\t\t{%- for tns in rns.nslist -%} +\t\t\t{%- set rns.recipients = rns.recipients + ["\\"" + tns + ".*\\""] -%} +\t\t{%- endfor -%} +\t{%- endif -%} +\t{%- set rns.recipients = rns.recipients + ["\\"user\\""] -%} +\t{{- "# Valid recipients: " + rns.recipients | join(", ") + "." -}} +{%- endmacro -%} +{{- bos_token -}} +{%- set ns = namespace(has_system=false) -%} +{%- for m in messages -%} +\t{%- if m["role"] == "system" -%} +\t\t{%- set ns.has_system = true -%} +\t{%- endif -%} +{%- endfor -%} +{%- if not ns.has_system -%} +\t{{- "<|start|>system<|message|>You are a helpful AI assistant." -}} +\t{%- set kc = knowledge_cutoff if knowledge_cutoff is defined and knowledge_cutoff else "2026-01-04" -%} +\t{{- "\\nKnowledge cutoff: " + kc + "." -}} +\t{%- if current_date is defined and current_date -%} +\t\t{{- "\\nCurrent date: " + current_date + "." -}} +\t{%- elif strftime_now is defined -%} +\t\t{{- "\\nCurrent date: " + strftime_now("%Y-%m-%d") + "." -}} +\t{%- endif -%} +\t{{- "\\n\\n" -}} +\t{{- render_reasoning() -}} +\t{%- if tools -%} +\t\t{{- "\\n\\n" -}} +\t\t{{- render_tool_defs(tools) -}} +\t{%- endif -%} +\t{{- "\\n\\n" -}} +\t{{- render_system_meta(tools) -}} +\t{{- "<|eot|>" -}} +{%- endif -%} +{%- for message in messages -%} +\t{%- set role = message["role"] -%} +\t{%- set end_token = "<|eom|>" if not loop.last and messages[loop.index0 + 1]["role"] == role else "<|eot|>" -%} +\t{%- if role == "system" -%} +\t\t{{- "<|start|>system<|message|>" -}} +\t\t{{- render_content(message["content"]) -}} +\t\t{{- "\\n\\n" -}} +\t\t{{- render_reasoning() -}} +\t\t{%- if tools -%} +\t\t\t{{- "\\n\\n" -}} +\t\t\t{{- render_tool_defs(tools) -}} +\t\t{%- endif -%} +\t\t{{- "\\n\\n" -}} +\t\t{{- render_system_meta(tools) -}} +\t\t{{- "<|eot|>" -}} +\t{%- elif role == "user" -%} +\t\t{{- "<|start|>user<|message|>" -}} +\t\t{{- render_content(message["content"]) -}} +\t\t{{- "<|eot|>" -}} +\t{%- elif role == "tool" -%} +\t\t{%- set tname = message.get("name") -%} +\t\t{%- if not tname -%} +\t\t\t{%- set tcid = message.get("tool_call_id") -%} +\t\t\t{%- set rns = namespace(name=tcid if tcid else "") -%} +\t\t\t{%- for m in messages -%} +\t\t\t\t{%- if m.get("tool_calls") -%} +\t\t\t\t\t{%- for tc in m["tool_calls"] -%} +\t\t\t\t\t\t{%- if tcid is not none and tc.id is defined and tc.id == tcid -%} +\t\t\t\t\t\t\t{%- set rns.name = tc.function.name -%} +\t\t\t\t\t\t{%- endif -%} +\t\t\t\t\t{%- endfor -%} +\t\t\t\t{%- endif -%} +\t\t\t{%- endfor -%} +\t\t\t{%- set tname = rns.name -%} +\t\t{%- endif -%} +\t\t{{- "<|start|>tool " + tname + "<|message|>\\n" -}} +\t\t{{- render_content(message["content"]) -}} +\t\t{{- "\\n<|eot|>" -}} +\t{%- elif role == "assistant" -%} +\t\t{%- if message.get("reasoning_content") -%} +\t\t\t{{- "<|start|>assistant to=self<|message|>" + message["reasoning_content"] + "<|eom|>" -}} +\t\t{%- endif -%} +\t\t{%- if message.get("tool_calls") -%} +\t\t\t{%- for tc in message["tool_calls"] -%} +\t\t\t\t{{- "<|start|>assistant to=" + tc.function.name + "<|message|>" -}} +\t\t\t\t{{- render_atem(tc) -}} +\t\t\t\t{%- if loop.last -%} +\t\t\t\t\t{{- end_token -}} +\t\t\t\t{%- else -%} +\t\t\t\t\t{{- "<|eom|>" -}} +\t\t\t\t{%- endif -%} +\t\t\t{%- endfor -%} +\t\t{%- else -%} +\t\t\t{%- set recipient = message.get("recipient") or "user" -%} +\t\t\t{%- set end_turn = message.get("end_turn") -%} +\t\t\t{%- if end_turn is none -%} +\t\t\t\t{%- set end_turn = not (recipient and recipient != "user") -%} +\t\t\t{%- endif -%} +\t\t\t{{- "<|start|>assistant" -}} +\t\t\t{%- if recipient -%} +\t\t\t\t{{- " to=" + recipient -}} +\t\t\t{%- endif -%} +\t\t\t{{- "<|message|>" -}} +\t\t\t{{- render_content(message["content"]) -}} +\t\t\t{{- "<|eot|>" if end_turn else "<|eom|>" -}} +\t\t{%- endif -%} +\t{%- endif -%} +{%- endfor -%} +{%- if add_generation_prompt -%} +\t{{- "<|start|>assistant" -}} +{%- endif -%} +`.slice(1, -1); diff --git a/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts b/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts index 4704d89e..62b33c71 100644 --- a/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts +++ b/test/standalone/chatWrappers/utils/resolveChatWrapper.test.ts @@ -6,7 +6,8 @@ import { } from "../../../../src/index.js"; import { harmonyJinjaTemplate, harmonyJinjaTemplate2, harmonyJinjaTemplate3, harmonyJinjaTemplate4, harmonyJinjaTemplate5, - gemma4JinjaTemplate1, gemma4JinjaTemplate2, gemma4JinjaTemplate3, museGlimmerJinjaTemplate + gemma4JinjaTemplate1, gemma4JinjaTemplate2, gemma4JinjaTemplate3, + museGlimmerJinjaTemplate, museGlimmerJinjaTemplate2, museGlimmerJinjaTemplate3, museGlimmerJinjaTemplate4 } from "./jinjaTemplates.js"; @@ -1005,4 +1006,43 @@ describe("resolveChatWrapper", () => { expect(chatWrapper).to.be.instanceof(MuseChatWrapper); }); + + test("should resolve to specialized MuseChatWrapper 2", {timeout: 1000 * 60 * 60 * 2}, async () => { + const chatWrapper = resolveChatWrapper({ + customWrapperSettings: { + jinjaTemplate: { + template: museGlimmerJinjaTemplate2 + } + }, + fallbackToOtherWrappersOnJinjaError: false + }); + + expect(chatWrapper).to.be.instanceof(MuseChatWrapper); + }); + + test("should resolve to specialized MuseChatWrapper 3", {timeout: 1000 * 60 * 60 * 2}, async () => { + const chatWrapper = resolveChatWrapper({ + customWrapperSettings: { + jinjaTemplate: { + template: museGlimmerJinjaTemplate3 + } + }, + fallbackToOtherWrappersOnJinjaError: false + }); + + expect(chatWrapper).to.be.instanceof(MuseChatWrapper); + }); + + test("should resolve to specialized MuseChatWrapper 4", {timeout: 1000 * 60 * 60 * 2}, async () => { + const chatWrapper = resolveChatWrapper({ + customWrapperSettings: { + jinjaTemplate: { + template: museGlimmerJinjaTemplate4 + } + }, + fallbackToOtherWrappersOnJinjaError: false + }); + + expect(chatWrapper).to.be.instanceof(MuseChatWrapper); + }); }); diff --git a/test/standalone/llamaEvaluator/AddonJinjaRenderer.test.ts b/test/standalone/llamaEvaluator/AddonJinjaRenderer.test.ts new file mode 100644 index 00000000..d8658b10 --- /dev/null +++ b/test/standalone/llamaEvaluator/AddonJinjaRenderer.test.ts @@ -0,0 +1,164 @@ +import {beforeAll, describe, expect, test} from "vitest"; +import {Template} from "@huggingface/jinja"; +import {getTestLlama} from "../../utils/getTestLlama.js"; +import type {BindingModule} from "../../../src/bindings/AddonTypes.js"; + + +describe("AddonJinjaRenderer", () => { + let NativeRenderer: BindingModule["AddonJinjaRenderer"]; + + beforeAll(async () => { + const llama = await getTestLlama(); + NativeRenderer = llama._bindings.AddonJinjaRenderer; + }); + + test.each([ + ["{{ name }}", {name: "Hello 👋\0!"}], + ["{{ data | tojson }}", {data: {z: 1, a: [true, null, 1.5, "text"]}}], + ["{{ (value is defined) | tojson }}|{{ (missing is defined) | tojson }}|{{ (value is none) | tojson }}", {value: null}], + ["{% for message in messages %}{{ message.role }}: {{ message.content }}\n{% endfor %}", { + messages: [{role: "user", content: "Hello"}, {role: "assistant", content: "Hi"}] + }], + ["{% macro greet(name) %}Hello {{ name }}{% endmacro %}{{ greet(name) }}", {name: "Alice"}], + ["{{ number }}|{{ fraction }}|{{ negative }}", {number: Number.MAX_SAFE_INTEGER, fraction: 1.25, negative: -123}], + ["{{ (optional is defined) | tojson }}|{{ (sparse[0] is defined) | tojson }}", {optional: undefined, sparse: new Array(1)}] + ] as const)("renders %s", (template, variables) => { + const renderer = new NativeRenderer(template); + expect(renderer.render(variables)).toBe(new Template(template).render(variables)); + expect(renderer.render(variables)).toBe(new Template(template).render(variables)); + }); + + test("accepts empty templates and omitted variables", () => { + expect(new NativeRenderer("").render()).toBe(""); + expect(new NativeRenderer("hello").render(undefined)).toBe("hello"); + }); + + test("keeps variables and mutations local to each render", () => { + const renderer = new NativeRenderer("{{ previous is defined }}{% set previous = value %}{{ previous }}"); + expect(renderer.render({value: "first"})).toBe("Falsefirst"); + expect(renderer.render({value: "second"})).toBe("Falsesecond"); + const input = {values: ["a"]}; + const mutating = new NativeRenderer("{% set ignored = values.append('b') %}{{ values | join(',') }}"); + expect(mutating.render(input)).toBe("a,b"); + expect(mutating.render(input)).toBe("a,b"); + expect(input.values).toEqual(["a"]); + }); + + test("uses only own enumerable string properties and preserves their order", () => { + const data = Object.assign(Object.create({inherited: "ignored"}), {z: 1, a: 2}); + Object.defineProperty(data, "hidden", {get() { + throw new Error("must not read hidden"); + }}); + Object.defineProperty(data, Symbol("ignored"), {enumerable: true, get() { + throw new Error("must not read symbol"); + }}); + expect(new NativeRenderer("{{ data | tojson }}").render({data})).toBe('{"z": 1, "a": 2}'); + expect(new NativeRenderer("{{ data | tojson }}").render({data: Object.assign(Object.create(null), {a: 1})})) + .toBe('{"a": 1}'); + }); + + test("allows repeated references without treating them as cycles", () => { + const shared = {text: "same"}; + expect(new NativeRenderer("{{ a.text }} {{ b.text }}").render({a: shared, b: shared})).toBe("same same"); + }); + + test.each([NaN, Infinity, -Infinity, 1e100, 1n, Symbol("unsupported"), () => "value"])( + "rejects unsupported values safely: %s", (value) => { + expect(() => new NativeRenderer("ok").render({nested: {value}})).toThrow(/failed to convert input at items\["nested"\]\["value"\]/); + } + ); + + test("preserves getter and proxy errors as causes", () => { + const renderer = new NativeRenderer("{{ value }}"); + const cause = new Error("getter failed"); + const inputs = [ + {get value() { + throw cause; + }}, + new Proxy({}, {ownKeys() { + throw cause; + }}) + ]; + for (const input of inputs) { + try { + renderer.render(input); + expect.fail("Expected conversion to fail"); + } catch (error) { + expect(error).toBeInstanceOf(Error); + expect((error as Error).cause).toBe(cause); + expect(String(error)).toContain("getter failed"); + } + } + expect(renderer.render({value: "recovered"})).toBe("recovered"); + }); + + test("reports nested array paths without reading a throwing getter again", () => { + const renderer = new NativeRenderer("{{ messages[0].content }}"); + const cause = new Error("content getter failed"); + let reads = 0; + const messages = [{content: "valid"}, {get content() { + reads++; + throw cause; + }}]; + + try { + renderer.render({messages}); + expect.fail("Expected conversion to fail"); + } catch (error) { + expect((error as Error).cause).toBe(cause); + expect(String(error)).toContain('failed to convert input at items["messages"][1]["content"]'); + } + expect(reads).toBe(1); + expect(renderer.render({messages: [{content: "recovered"}]})).toBe("recovered"); + }); + + test("allows nested rendering from an input getter with independent variables", () => { + const renderer = new NativeRenderer("{{ prefix }}{{ value }}"); + expect(renderer.render({prefix: "outer:", get value() { + return renderer.render({prefix: "inner:", value: "nested"}); + }})).toBe("outer:inner:nested"); + expect(renderer.render({prefix: "next:", value: "value"})).toBe("next:value"); + }); + + test("can render after a nested render fails", () => { + const renderer = new NativeRenderer("{% filter indent(width) %}a\nb{% endfilter %}"); + expect(renderer.render({get width() { + expect(() => renderer.render({width: -1})).toThrow(/failed to render template/); + return 2; + }})).toBe("a\n b"); + expect(renderer.render({width: 4})).toBe("a\n b"); + }); + + test("reports compile errors", () => { + expect(() => new NativeRenderer("{% if %}")).toThrow(/AddonJinjaRenderer: failed to compile template/); + expect(() => new NativeRenderer("{{ 'unterminated")).toThrow(/lexer/); + }); + + test("reports runtime errors with source locations", () => { + const renderer = new NativeRenderer("hello\n{{ raise_exception('bad role') }}"); + expect(() => renderer.render()).toThrow(/failed to render template[\s\S]*line 2[\s\S]*bad role/); + }); + + test("can reuse a filter block after it throws", () => { + const renderer = new NativeRenderer("{% filter indent(width) %}a\nb{% endfilter %}"); + expect(() => renderer.render({width: -1})).toThrow(/failed to render template/); + expect(renderer.render({width: 2})).toBe("a\n b"); + expect(renderer.render({width: 4})).toBe("a\n b"); + }); + + test("renders recursive macros using the upstream runtime", () => { + const renderer = new NativeRenderer("{% macro recurse(n) %}{% if n > 0 %}{{ recurse(n - 1) }}{% else %}ok{% endif %}{% endmacro %}{{ recurse(n) }}"); + expect(renderer.render({n: 80})).toBe("ok"); + }); + + test("can render the same and another template after a macro throws on the same thread", () => { + const renderer = new NativeRenderer("{% macro recurse(n) %}{% if n > 0 %}{{ recurse(n - 1) }}{% elif fail %}{{ raise_exception('macro failed') }}{% else %}ok{% endif %}{% endmacro %}{{ recurse(n) }}"); + const other = new NativeRenderer("{% macro another(n) %}{% if n > 0 %}{{ another(n - 1) }}{% else %}other{% endif %}{% endmacro %}{{ another(n) }}"); + + for (let i = 0; i < 3; i++) { + expect(() => renderer.render({n: 8, fail: true})).toThrow(/macro failed/); + expect(other.render({n: 8})).toBe("other"); + expect(renderer.render({n: 8, fail: false})).toBe("ok"); + } + }); +});