feat: expose ubatch param

mybigday · Jan 11, 2025 · c571221 · c571221
1 parent 00bb6a0
commit c571221
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/lib/binding.ts b/lib/binding.ts
@@ -12,6 +12,7 @@ export type LlamaModelOptions = {
   pooling_type?: number
   n_ctx?: number
   n_batch?: number
+  n_ubatch?: number
   n_threads?: number
   n_gpu_layers?: number
   use_mlock?: boolean

diff --git a/src/LlamaContext.cpp b/src/LlamaContext.cpp
@@ -83,6 +83,7 @@ LlamaContext::LlamaContext(const Napi::CallbackInfo &info)
 
   params.n_ctx = get_option<int32_t>(options, "n_ctx", 512);
   params.n_batch = get_option<int32_t>(options, "n_batch", 2048);
+  params.n_ubatch = get_option<int32_t>(options, "n_ubatch", 512);
   params.embedding = get_option<bool>(options, "embedding", false);
   if (params.embedding) {
     // For non-causal models, batch size must be equal to ubatch size