From fe3dbbffbf1a1c00311b94ab190741034b757e68 Mon Sep 17 00:00:00 2001 From: Claude Date: Thu, 25 Dec 2025 15:16:41 +0000 Subject: [PATCH] Expand FFI bindings with grammar, LoRA, embeddings, and more New FFI bindings: - Grammar sampling (llama_sampler_init_grammar) for structured output - LoRA adapter support (load, apply, remove adapters) - State/session save and load functions - Additional samplers (XTC, top-n-sigma, infill, DRY) - FIM (Fill-in-Middle) tokens for code completion - Extended model metadata access - Chat template formatting - Embeddings extraction helpers - Control vectors support - Context settings (causal attention, warmup mode) - Additional vocab functions (sep, mask, score, attr) - NUMA and RPC support detection New high-level API functions: - sampler_create_with_grammar() - Grammar-constrained sampling - lora_load/apply/clear() - LoRA adapter management - state_save/load() - Session persistence - chat_format() - Format messages with model template - extract_embeddings() - Get text embeddings as array - fim_prompt() - Create FIM prompts for code models - model_info_extended() - Extended model information - model_meta/model_meta_all() - Access GGUF metadata New examples: - json_output.hml - Generate structured JSON - lora.hml - LoRA adapter usage - embeddings.hml - Text embeddings and similarity - fim.hml - Fill-in-middle code completion - session.hml - Save/restore conversation state Constants added: - Token types and attributes - ROPE_SCALING_LONGROPE - GRAMMAR_JSON (built-in JSON grammar) --- README.md | 184 ++++++++++- examples/embeddings.hml | 145 +++++++++ examples/fim.hml | 120 +++++++ examples/json_output.hml | 87 +++++ examples/lora.hml | 107 +++++++ examples/session.hml | 150 +++++++++ llama.hml | 672 +++++++++++++++++++++++++++++++++++++++ 7 files changed, 1464 insertions(+), 1 deletion(-) create mode 100644 examples/embeddings.hml create mode 100644 examples/fim.hml create mode 100644 examples/json_output.hml create mode 100644 examples/lora.hml create mode 100644 examples/session.hml diff --git a/README.md b/README.md index d036865..a489290 100644 --- a/README.md +++ b/README.md @@ -2,6 +2,17 @@ Hemlock bindings for [llama.cpp](https://github.com/ggml-org/llama.cpp), enabling LLM inference in the Hemlock programming language. +## Features + +- **Full llama.cpp API** - Access to model loading, tokenization, inference, and sampling +- **Grammar Constraints** - Generate structured output (JSON, etc.) with GBNF grammars +- **LoRA Adapters** - Load and apply LoRA fine-tuned adapters +- **Session Management** - Save and restore conversation state +- **Embeddings** - Extract text embeddings for semantic similarity +- **FIM Support** - Fill-in-Middle code completion +- **Chat Templates** - Format conversations using model chat templates +- **GPU Acceleration** - CUDA and Metal support + ## Prerequisites - [Hemlock](https://github.com/yourusername/hemlock) programming language installed @@ -123,6 +134,104 @@ llama.model_free(model); llama.backend_free(); ``` +## Grammar-Constrained Generation + +Generate valid JSON or other structured output: + +```hemlock +// Use built-in JSON grammar +let sampler = llama.sampler_create_with_grammar(model, { + grammar: llama.GRAMMAR_JSON, + grammar_root: "root", + temp: 0.7 +}); + +// Or define custom grammar +let custom_grammar = "root ::= \"yes\" | \"no\""; +let sampler = llama.sampler_create_with_grammar(model, { + grammar: custom_grammar, + temp: 0.1 +}); +``` + +## LoRA Adapters + +Apply fine-tuned LoRA adapters: + +```hemlock +// Load adapter +let lora = llama.lora_load(model, "adapter.gguf"); + +// Apply with full strength +llama.lora_apply(ctx, lora, 1.0); + +// Or partial strength +llama.lora_apply(ctx, lora, 0.5); + +// Remove adapter +llama.lora_clear(ctx); +llama.lora_free(lora); +``` + +## Session Management + +Save and restore conversation state: + +```hemlock +// Save session +let tokens = [...]; // collected tokens +llama.state_save(ctx, "session.bin", tokens); + +// Load session later +let loaded = llama.state_load(ctx, "session.bin", 4096); +if (loaded.success) { + print("Loaded " + loaded.tokens.length + " tokens"); +} +``` + +## Chat Templates + +Format conversations using model templates: + +```hemlock +let messages = [ + { role: "system", content: "You are a helpful assistant." }, + { role: "user", content: "Hello!" } +]; + +let formatted = llama.chat_format(model, messages, true); +print(formatted); // Formatted prompt with model's template +``` + +## Fill-in-Middle (Code Completion) + +```hemlock +if (llama.supports_fim(model)) { + let prompt = llama.fim_prompt(model, + "def add(a, b):\n ", // prefix + "\n return result" // suffix + ); + // Generate completion for the middle +} +``` + +## Embeddings + +Extract text embeddings: + +```hemlock +let ctx = llama.context_create(model, { + embeddings: true, + pooling_type: llama.POOLING_TYPE_MEAN +}); + +let tokens = llama.tokenize(model, "Hello world", { add_special: true }); +llama.decode(ctx, tokens, 0); + +let embeddings = llama.extract_embeddings(ctx, model, 0); +print("Embedding dimension: " + embeddings.length); +``` + ## API Reference ### Backend Functions @@ -131,6 +240,8 @@ llama.backend_free(); - `backend_free()` - Free the backend (call once at end) - `system_info()` - Get system info string - `supports_gpu()` - Check if GPU offloading is supported +- `supports_rpc()` - Check if RPC is supported +- `numa_init(numa)` - Initialize NUMA optimizations ### Model Functions @@ -138,15 +249,23 @@ llama.backend_free(); - Options: `n_gpu_layers`, `use_mmap`, `use_mlock`, `split_mode`, `main_gpu` - `model_free(model)` - Free a loaded model - `model_info(model)` - Get model information +- `model_info_extended(model)` - Get extended model info (hybrid, diffusion, RoPE, etc.) - `model_vocab(model)` - Get the model's vocabulary pointer +- `model_meta(model, key)` - Get metadata value by key +- `model_meta_all(model)` - Get all metadata as object +- `chat_template(model)` - Get model's chat template ### Context Functions - `context_create(model, options)` - Create an inference context - - Options: `n_ctx`, `n_batch`, `n_ubatch`, `n_threads`, `embeddings` + - Options: `n_ctx`, `n_batch`, `n_ubatch`, `n_threads`, `embeddings`, `n_seq_max` - `context_free(ctx)` - Free a context - `context_info(ctx)` - Get context information - `set_threads(ctx, n, n_batch)` - Set thread counts +- `set_embeddings(ctx, enabled)` - Enable/disable embeddings +- `set_causal_attn(ctx, causal)` - Set causal attention mode +- `set_warmup(ctx, warmup)` - Set warmup mode +- `n_ctx_seq(ctx)` - Get sequence context size ### Tokenization @@ -157,16 +276,60 @@ llama.backend_free(); - `token_to_piece(model, token, special)` - Convert single token to text - `vocab_info(model)` - Get vocabulary information - `is_eog(model, token)` - Check if token is end-of-generation +- `token_score(model, token)` - Get token score +- `token_attr(model, token)` - Get token attributes +- `vocab_sep(model)` - Get separator token +- `vocab_mask(model)` - Get mask token + +### FIM (Fill-in-Middle) + +- `supports_fim(model)` - Check if model supports FIM +- `fim_tokens(model)` - Get FIM special tokens +- `fim_prompt(model, prefix, suffix)` - Create FIM prompt ### Sampling - `sampler_create(options)` - Create a sampler chain - Options: `top_k`, `top_p`, `min_p`, `temp`, `repeat_penalty`, `repeat_last_n`, `seed` +- `sampler_create_with_grammar(model, options)` - Create sampler with grammar + - Additional options: `grammar`, `grammar_root` +- `sampler_grammar(model, grammar_str, grammar_root)` - Create grammar-only sampler - `sampler_greedy()` - Create a greedy sampler +- `sampler_xtc(p, t, min_keep, seed)` - Create XTC sampler +- `sampler_top_n_sigma(n)` - Create Top-N-Sigma sampler +- `sampler_infill(model)` - Create infill sampler for FIM - `sampler_free(sampler)` - Free a sampler - `sampler_reset(sampler)` - Reset sampler state - `sample(sampler, ctx, idx)` - Sample a token - `sampler_accept(sampler, token)` - Accept a token (for penalties) +- `GRAMMAR_JSON` - Built-in JSON grammar string + +### LoRA Adapters + +- `lora_load(model, path)` - Load a LoRA adapter +- `lora_free(adapter)` - Free a LoRA adapter +- `lora_apply(ctx, adapter, scale)` - Apply adapter to context +- `lora_remove(ctx, adapter)` - Remove specific adapter +- `lora_clear(ctx)` - Remove all adapters + +### State/Sessions + +- `state_size(ctx)` - Get state size in bytes +- `state_save(ctx, path, tokens)` - Save state to file +- `state_load(ctx, path, max_tokens)` - Load state from file +- `state_seq_save(ctx, path, seq_id, tokens)` - Save single sequence +- `state_seq_load(ctx, path, seq_id, max_tokens)` - Load single sequence + +### Embeddings + +- `get_embeddings(ctx, seq_id)` - Get embeddings pointer +- `extract_embeddings(ctx, model, seq_id)` - Extract embeddings as array +- `get_pooling_type(ctx)` - Get pooling type + +### Chat Templates + +- `chat_template(model)` - Get model's chat template string +- `chat_format(model, messages, add_assistant)` - Format messages array ### Inference @@ -183,8 +346,10 @@ llama.backend_free(); - `kv_cache_seq_cp(ctx, src, dst, p0, p1)` - Copy a sequence - `kv_cache_seq_keep(ctx, seq_id)` - Keep only one sequence - `kv_cache_seq_add(ctx, seq_id, p0, p1, delta)` - Shift positions +- `kv_cache_seq_div(ctx, seq_id, p0, p1, d)` - Divide positions - `kv_cache_seq_pos_min(ctx, seq_id)` - Get min position - `kv_cache_seq_pos_max(ctx, seq_id)` - Get max position +- `memory_can_shift(ctx)` - Check if memory supports shifting ### High-Level @@ -196,8 +361,18 @@ llama.backend_free(); - `perf_print(ctx)` - Print context performance stats - `perf_reset(ctx)` - Reset performance stats +- `perf_sampler_print(sampler)` - Print sampler stats - `time_us()` - Get current time in microseconds +### Constants + +- Token types: `TOKEN_TYPE_NORMAL`, `TOKEN_TYPE_CONTROL`, etc. +- Token attributes: `TOKEN_ATTR_NORMAL`, `TOKEN_ATTR_CONTROL`, etc. +- Vocab types: `VOCAB_TYPE_SPM`, `VOCAB_TYPE_BPE`, `VOCAB_TYPE_WPM`, etc. +- Split modes: `SPLIT_MODE_NONE`, `SPLIT_MODE_LAYER`, `SPLIT_MODE_ROW` +- Pooling types: `POOLING_TYPE_NONE`, `POOLING_TYPE_MEAN`, `POOLING_TYPE_CLS`, `POOLING_TYPE_LAST` +- RoPE scaling: `ROPE_SCALING_NONE`, `ROPE_SCALING_LINEAR`, `ROPE_SCALING_YARN`, `ROPE_SCALING_LONGROPE` + ## Examples See the `examples/` directory: @@ -205,11 +380,18 @@ See the `examples/` directory: - `simple.hml` - Basic text generation - `chat.hml` - Interactive chat interface - `oneshot.hml` - One-shot completion +- `json_output.hml` - Grammar-constrained JSON generation +- `lora.hml` - LoRA adapter usage +- `embeddings.hml` - Text embeddings extraction +- `fim.hml` - Fill-in-Middle code completion +- `session.hml` - Save/restore conversation sessions Run examples: ```bash export LD_LIBRARY_PATH="$(pwd)/lib:$LD_LIBRARY_PATH" hemlock examples/simple.hml ~/models/llama-7b.gguf "What is the meaning of life?" +hemlock examples/json_output.hml ~/models/llama.gguf "List 3 colors" +hemlock examples/chat.hml ~/models/llama.gguf ``` ## License diff --git a/examples/embeddings.hml b/examples/embeddings.hml new file mode 100644 index 0000000..67d5417 --- /dev/null +++ b/examples/embeddings.hml @@ -0,0 +1,145 @@ +// embeddings.hml - Extract text embeddings for semantic similarity +// +// Usage: +// hemlock examples/embeddings.hml "" "" +// +// Example: +// hemlock examples/embeddings.hml ~/models/nomic-embed.gguf "Hello world" "Hi there" +// +// Note: This requires an embedding model (e.g., nomic-embed, bge, e5) + +import * as llama from "../llama.hml"; + +// Compute cosine similarity between two embedding vectors +fn cosine_similarity(a: array, b: array): f32 { + if (a.length != b.length || a.length == 0) { + return 0.0; + } + + let dot = 0.0; + let norm_a = 0.0; + let norm_b = 0.0; + + for (let i = 0; i < a.length; i = i + 1) { + dot = dot + a[i] * b[i]; + norm_a = norm_a + a[i] * a[i]; + norm_b = norm_b + b[i] * b[i]; + } + + if (norm_a == 0.0 || norm_b == 0.0) { + return 0.0; + } + + return dot / (sqrt(norm_a) * sqrt(norm_b)); +} + +fn sqrt(x: f32): f32 { + // Newton's method for square root + if (x <= 0.0) { + return 0.0; + } + let guess = x / 2.0; + for (let i = 0; i < 10; i = i + 1) { + guess = (guess + x / guess) / 2.0; + } + return guess; +} + +fn main() { + let args = __args(); + if (args.length < 4) { + print("Usage: hemlock examples/embeddings.hml \"\" \"\""); + print(""); + print("This example extracts embeddings and computes semantic similarity."); + print("Use an embedding model like nomic-embed, bge, or e5."); + return; + } + + let model_path = args[1]; + let text1 = args[2]; + let text2 = args[3]; + + print("=== Text Embeddings Example ===\n"); + + llama.backend_init(); + + // Load embedding model + print("Loading model: " + model_path); + let model = llama.model_load(model_path, { + n_gpu_layers: llama.supports_gpu() ? 99 : 0 + }); + + let info = llama.model_info(model); + print("Model: " + info.description); + print("Embedding dimension: " + info.n_embd); + + // Create context with embeddings enabled + let ctx = llama.context_create(model, { + n_ctx: 512, + n_batch: 512, + embeddings: true + }); + + // Check pooling type + let pooling = llama.get_pooling_type(ctx); + print("Pooling type: " + pooling); + + // Get embeddings for text1 + print("\nText 1: \"" + text1 + "\""); + let tokens1 = llama.tokenize(model, text1, { add_special: true }); + print(" Tokens: " + tokens1.length); + + llama.decode(ctx, tokens1, 0); + let emb1 = llama.extract_embeddings(ctx, model, 0); + + if (emb1.length == 0) { + print(" Warning: Could not extract embeddings. Model may not support pooling."); + print(" Try using llama.get_logits() for token-level representations instead."); + } else { + print(" Embedding dimension: " + emb1.length); + print(" First 5 values: [" + emb1[0] + ", " + emb1[1] + ", " + emb1[2] + ", " + emb1[3] + ", " + emb1[4] + ", ...]"); + } + + // Clear cache and get embeddings for text2 + llama.kv_cache_clear(ctx); + + print("\nText 2: \"" + text2 + "\""); + let tokens2 = llama.tokenize(model, text2, { add_special: true }); + print(" Tokens: " + tokens2.length); + + llama.decode(ctx, tokens2, 0); + let emb2 = llama.extract_embeddings(ctx, model, 0); + + if (emb2.length > 0) { + print(" Embedding dimension: " + emb2.length); + print(" First 5 values: [" + emb2[0] + ", " + emb2[1] + ", " + emb2[2] + ", " + emb2[3] + ", " + emb2[4] + ", ...]"); + } + + // Compute similarity + if (emb1.length > 0 && emb2.length > 0) { + let similarity = cosine_similarity(emb1, emb2); + print("\n=== Similarity Score ==="); + print("Cosine similarity: " + similarity); + + if (similarity > 0.9) { + print("Interpretation: Very similar (nearly identical meaning)"); + } else if (similarity > 0.7) { + print("Interpretation: Similar (related meaning)"); + } else if (similarity > 0.5) { + print("Interpretation: Somewhat similar"); + } else if (similarity > 0.3) { + print("Interpretation: Weakly related"); + } else { + print("Interpretation: Different meaning"); + } + } + + // Cleanup + llama.context_free(ctx); + llama.model_free(model); + llama.backend_free(); + + print("\nDone!"); +} + +main(); diff --git a/examples/fim.hml b/examples/fim.hml new file mode 100644 index 0000000..6aa252c --- /dev/null +++ b/examples/fim.hml @@ -0,0 +1,120 @@ +// fim.hml - Fill-in-Middle code completion +// +// Usage: +// hemlock examples/fim.hml "" "" +// +// Example: +// hemlock examples/fim.hml ~/models/codellama.gguf "def add(a, b):\n " "\n return result" +// +// Note: This requires a model that supports FIM (e.g., CodeLlama, StarCoder, DeepSeek-Coder) + +import * as llama from "../llama.hml"; + +fn main() { + let args = __args(); + if (args.length < 4) { + print("Usage: hemlock examples/fim.hml \"\" \"\""); + print(""); + print("This example demonstrates Fill-in-Middle (FIM) code completion."); + print("Use a code model that supports FIM like CodeLlama, StarCoder, or DeepSeek-Coder."); + print(""); + print("Example:"); + print(" hemlock examples/fim.hml codellama.gguf \"def add(a, b):\\n \" \"\\n return result\""); + return; + } + + let model_path = args[1]; + let prefix = args[2]; + let suffix = args[3]; + + print("=== Fill-in-Middle Code Completion ===\n"); + + llama.backend_init(); + + // Load model + print("Loading model: " + model_path); + let model = llama.model_load(model_path, { + n_gpu_layers: llama.supports_gpu() ? 99 : 0 + }); + print("Model: " + llama.model_info(model).description); + + // Check FIM support + if (!llama.supports_fim(model)) { + print("\nError: This model does not support FIM (Fill-in-Middle)."); + print("FIM requires special tokens (FIM_PREFIX, FIM_SUFFIX, FIM_MIDDLE)."); + print("Try using CodeLlama, StarCoder, DeepSeek-Coder, or similar code models."); + llama.model_free(model); + llama.backend_free(); + return; + } + + // Get FIM tokens + let fim = llama.fim_tokens(model); + print("\nFIM tokens found:"); + print(" PREFIX: " + fim.pre + " -> \"" + llama.token_to_piece(model, fim.pre, true) + "\""); + print(" SUFFIX: " + fim.suf + " -> \"" + llama.token_to_piece(model, fim.suf, true) + "\""); + print(" MIDDLE: " + fim.mid + " -> \"" + llama.token_to_piece(model, fim.mid, true) + "\""); + + // Create context + let ctx = llama.context_create(model, { + n_ctx: 2048, + n_batch: 512 + }); + + // Create infill-optimized sampler chain + let sampler = llama.sampler_create({ + temp: 0.2, // Lower temperature for more deterministic code + top_k: 40, + top_p: 0.95, + repeat_penalty: 1.0 // No repeat penalty for code + }); + + // Add infill sampler for better FIM behavior + let infill_sampler = llama.sampler_infill(model); + + // Build FIM prompt + let fim_prompt = llama.fim_prompt(model, prefix, suffix); + + print("\n--- Input ---"); + print("Prefix: " + prefix); + print("Suffix: " + suffix); + print("\n--- Completion ---"); + + // Tokenize and decode + let tokens = llama.tokenize(model, fim_prompt, { add_special: false, parse_special: true }); + llama.decode(ctx, tokens, 0); + + // Generate completion + let n_cur = tokens.length; + let output = ""; + + for (let i = 0; i < 256; i = i + 1) { + let token = llama.sample(sampler, ctx, -1); + + // Stop on end-of-generation or suffix token + if (llama.is_eog(model, token) || token == fim.suf) { + break; + } + + let piece = llama.token_to_piece(model, token, false); + output = output + piece; + print_inline(piece); + + llama.decode(ctx, [token], n_cur); + n_cur = n_cur + 1; + } + + print("\n\n--- Full Result ---"); + print(prefix + output + suffix); + + // Cleanup + llama.sampler_free(infill_sampler); + llama.sampler_free(sampler); + llama.context_free(ctx); + llama.model_free(model); + llama.backend_free(); + + print("\nDone!"); +} + +main(); diff --git a/examples/json_output.hml b/examples/json_output.hml new file mode 100644 index 0000000..af586a4 --- /dev/null +++ b/examples/json_output.hml @@ -0,0 +1,87 @@ +// json_output.hml - Generate structured JSON output using grammar constraints +// +// Usage: +// hemlock examples/json_output.hml "" +// +// Example: +// hemlock examples/json_output.hml ~/models/llama.gguf "List 3 programming languages with name and year" + +import * as llama from "../llama.hml"; + +fn main() { + let args = __args(); + if (args.length < 3) { + print("Usage: hemlock examples/json_output.hml \"\""); + return; + } + + let model_path = args[1]; + let prompt = args[2]; + + print("Loading model: " + model_path); + llama.backend_init(); + + let model = llama.model_load(model_path, { + n_gpu_layers: llama.supports_gpu() ? 99 : 0 + }); + + print("Model loaded: " + llama.model_info(model).description); + + let ctx = llama.context_create(model, { + n_ctx: 2048, + n_batch: 512 + }); + + // Create a sampler with JSON grammar constraint + let sampler = llama.sampler_create_with_grammar(model, { + grammar: llama.GRAMMAR_JSON, + grammar_root: "root", + temp: 0.7, + top_k: 40, + top_p: 0.9 + }); + + // Format the prompt to request JSON output + let full_prompt = "You are a helpful assistant that responds only in valid JSON.\n\n" + + "User: " + prompt + "\n\n" + + "Respond with a JSON object or array:\n"; + + print("\nPrompt: " + prompt); + print("\nGenerating JSON response...\n"); + + // Tokenize and decode prompt + let tokens = llama.tokenize(model, full_prompt, { add_special: true }); + llama.decode(ctx, tokens, 0); + + // Generate tokens + let output_tokens = []; + let n_cur = tokens.length; + let max_tokens = 512; + + for (let i = 0; i < max_tokens; i = i + 1) { + let token = llama.sample(sampler, ctx, -1); + + if (llama.is_eog(model, token)) { + break; + } + + output_tokens.push(token); + let piece = llama.token_to_piece(model, token, false); + print_inline(piece); + + llama.decode(ctx, [token], n_cur); + n_cur = n_cur + 1; + } + + print("\n"); + + // Cleanup + llama.sampler_free(sampler); + llama.context_free(ctx); + llama.model_free(model); + llama.backend_free(); + + print("\nDone!"); +} + +main(); diff --git a/examples/lora.hml b/examples/lora.hml new file mode 100644 index 0000000..7a4dc09 --- /dev/null +++ b/examples/lora.hml @@ -0,0 +1,107 @@ +// lora.hml - Demonstrate LoRA adapter loading and application +// +// Usage: +// hemlock examples/lora.hml "" +// +// Example: +// hemlock examples/lora.hml ~/models/llama.gguf ~/models/my-lora.gguf "Hello!" + +import * as llama from "../llama.hml"; + +fn main() { + let args = __args(); + if (args.length < 4) { + print("Usage: hemlock examples/lora.hml \"\""); + print(""); + print("This example demonstrates loading and applying a LoRA adapter."); + print("LoRA adapters allow fine-tuning models without modifying base weights."); + return; + } + + let model_path = args[1]; + let lora_path = args[2]; + let prompt = args[3]; + + print("=== LoRA Adapter Example ===\n"); + + llama.backend_init(); + + // Load base model + print("Loading base model: " + model_path); + let model = llama.model_load(model_path, { + n_gpu_layers: llama.supports_gpu() ? 99 : 0 + }); + print("Model: " + llama.model_info(model).description); + + // Load LoRA adapter + print("\nLoading LoRA adapter: " + lora_path); + let lora = llama.lora_load(model, lora_path); + print("LoRA adapter loaded successfully!"); + + // Create context + let ctx = llama.context_create(model, { + n_ctx: 2048, + n_batch: 512 + }); + + // Apply LoRA adapter with full strength (scale = 1.0) + print("\nApplying LoRA adapter with scale=1.0..."); + llama.lora_apply(ctx, lora, 1.0); + + // Create sampler + let sampler = llama.sampler_create({ + temp: 0.7, + top_k: 40, + top_p: 0.9 + }); + + // Generate with LoRA + print("\n--- Generation with LoRA (scale=1.0) ---"); + print("Prompt: " + prompt + "\n"); + + let tokens = llama.tokenize(model, prompt, { add_special: true }); + llama.decode(ctx, tokens, 0); + + let n_cur = tokens.length; + for (let i = 0; i < 128; i = i + 1) { + let token = llama.sample(sampler, ctx, -1); + if (llama.is_eog(model, token)) { + break; + } + print_inline(llama.token_to_piece(model, token, false)); + llama.decode(ctx, [token], n_cur); + n_cur = n_cur + 1; + } + print("\n"); + + // Clear and try with different scale + print("\n--- Now with reduced LoRA scale (0.5) ---"); + llama.lora_clear(ctx); + llama.lora_apply(ctx, lora, 0.5); + llama.kv_cache_clear(ctx); + llama.sampler_reset(sampler); + + llama.decode(ctx, tokens, 0); + n_cur = tokens.length; + for (let i = 0; i < 128; i = i + 1) { + let token = llama.sample(sampler, ctx, -1); + if (llama.is_eog(model, token)) { + break; + } + print_inline(llama.token_to_piece(model, token, false)); + llama.decode(ctx, [token], n_cur); + n_cur = n_cur + 1; + } + print("\n"); + + // Cleanup + llama.sampler_free(sampler); + llama.context_free(ctx); + llama.lora_free(lora); + llama.model_free(model); + llama.backend_free(); + + print("\nDone!"); +} + +main(); diff --git a/examples/session.hml b/examples/session.hml new file mode 100644 index 0000000..cfc29ff --- /dev/null +++ b/examples/session.hml @@ -0,0 +1,150 @@ +// session.hml - Save and restore conversation sessions +// +// Usage: +// hemlock examples/session.hml save "" +// hemlock examples/session.hml load "" +// +// Example: +// # Start a new session +// hemlock examples/session.hml llama.gguf save session.bin "Tell me about cats." +// +// # Continue from saved session +// hemlock examples/session.hml llama.gguf load session.bin "What else can you tell me?" + +import * as llama from "../llama.hml"; + +fn main() { + let args = __args(); + if (args.length < 5) { + print("Usage:"); + print(" hemlock examples/session.hml save \"\""); + print(" hemlock examples/session.hml load \"\""); + print(""); + print("This example demonstrates saving and restoring KV cache state."); + print("Useful for long conversations or expensive prompt processing."); + return; + } + + let model_path = args[1]; + let mode = args[2]; // "save" or "load" + let session_path = args[3]; + let prompt = args[4]; + + if (mode != "save" && mode != "load") { + print("Error: Mode must be 'save' or 'load'"); + return; + } + + print("=== Session Management Example ===\n"); + + llama.backend_init(); + + // Load model + print("Loading model: " + model_path); + let model = llama.model_load(model_path, { + n_gpu_layers: llama.supports_gpu() ? 99 : 0 + }); + print("Model: " + llama.model_info(model).description); + + // Create context + let ctx = llama.context_create(model, { + n_ctx: 4096, // Larger context for sessions + n_batch: 512 + }); + + let sampler = llama.sampler_create({ + temp: 0.7, + top_k: 40, + top_p: 0.9 + }); + + let all_tokens = []; + + if (mode == "load") { + // Load existing session + print("\nLoading session from: " + session_path); + let loaded = llama.state_load(ctx, session_path, 4096); + + if (!loaded.success) { + print("Error: Failed to load session file"); + llama.sampler_free(sampler); + llama.context_free(ctx); + llama.model_free(model); + llama.backend_free(); + return; + } + + all_tokens = loaded.tokens; + print("Loaded " + all_tokens.length + " tokens from session"); + + // Show what was in the session + let prev_text = llama.detokenize(model, all_tokens, { remove_special: false }); + print("\n--- Previous conversation ---"); + print(prev_text); + print("\n--- Continuing ---\n"); + } + + // Add new prompt + print("Prompt: " + prompt + "\n"); + let new_tokens = llama.tokenize(model, prompt, { + add_special: mode == "save" // Only add BOS on new session + }); + + // Decode all tokens + if (mode == "load") { + // Already have KV cache from loaded state, just decode new tokens + llama.decode(ctx, new_tokens, all_tokens.length); + } else { + // New session, decode from scratch + llama.decode(ctx, new_tokens, 0); + } + + for (let i = 0; i < new_tokens.length; i = i + 1) { + all_tokens.push(new_tokens[i]); + } + + // Generate response + print("Response: "); + let n_cur = all_tokens.length; + + for (let i = 0; i < 256; i = i + 1) { + let token = llama.sample(sampler, ctx, -1); + + if (llama.is_eog(model, token)) { + break; + } + + all_tokens.push(token); + let piece = llama.token_to_piece(model, token, false); + print_inline(piece); + + llama.decode(ctx, [token], n_cur); + n_cur = n_cur + 1; + } + print("\n"); + + // Save session + print("\nSaving session to: " + session_path); + print("Total tokens: " + all_tokens.length); + + let saved = llama.state_save(ctx, session_path, all_tokens); + if (saved) { + print("Session saved successfully!"); + + // Show file size + let state_size = llama.state_size(ctx); + print("State size: " + (state_size / 1024 / 1024) + " MB"); + } else { + print("Error: Failed to save session"); + } + + // Cleanup + llama.sampler_free(sampler); + llama.context_free(ctx); + llama.model_free(model); + llama.backend_free(); + + print("\nDone!"); +} + +main(); diff --git a/llama.hml b/llama.hml index 9d69451..8442f19 100644 --- a/llama.hml +++ b/llama.hml @@ -45,6 +45,29 @@ export let ROPE_SCALING_UNSPECIFIED = -1; export let ROPE_SCALING_NONE = 0; export let ROPE_SCALING_LINEAR = 1; export let ROPE_SCALING_YARN = 2; +export let ROPE_SCALING_LONGROPE = 3; + +// Token types +export let TOKEN_TYPE_UNDEFINED = 0; +export let TOKEN_TYPE_NORMAL = 1; +export let TOKEN_TYPE_UNKNOWN = 2; +export let TOKEN_TYPE_CONTROL = 3; +export let TOKEN_TYPE_USER_DEFINED = 4; +export let TOKEN_TYPE_UNUSED = 5; +export let TOKEN_TYPE_BYTE = 6; + +// Token attributes (bitflags) +export let TOKEN_ATTR_UNDEFINED = 0; +export let TOKEN_ATTR_UNKNOWN = 1; +export let TOKEN_ATTR_UNUSED = 2; +export let TOKEN_ATTR_NORMAL = 4; +export let TOKEN_ATTR_CONTROL = 8; +export let TOKEN_ATTR_USER_DEFINED = 16; +export let TOKEN_ATTR_BYTE = 32; +export let TOKEN_ATTR_NORMALIZED = 64; +export let TOKEN_ATTR_LSTRIP = 128; +export let TOKEN_ATTR_RSTRIP = 256; +export let TOKEN_ATTR_SINGLE_WORD = 512; // ============================================================================ // Low-level FFI declarations @@ -185,6 +208,89 @@ extern fn llama_perf_context_reset(ctx: ptr): void; extern fn llama_perf_sampler_print(chain: ptr): void; extern fn llama_perf_sampler_reset(chain: ptr): void; +// Grammar sampling +extern fn llama_sampler_init_grammar(vocab: ptr, grammar_str: string, grammar_root: string): ptr; + +// Additional samplers +extern fn llama_sampler_init_xtc(p: f32, t: f32, min_keep: u64, seed: u32): ptr; +extern fn llama_sampler_init_top_n_sigma(n: f32): ptr; +extern fn llama_sampler_init_dry(vocab: ptr, n_ctx_train: i32, dry_multiplier: f32, dry_base: f32, dry_allowed_length: i32, dry_penalty_last_n: i32, seq_breakers: ptr, num_breakers: u64): ptr; +extern fn llama_sampler_init_logit_bias(n_vocab: i32, n_logit_bias: i32, logit_bias: ptr): ptr; +extern fn llama_sampler_init_infill(vocab: ptr): ptr; + +// LoRA adapters +extern fn llama_adapter_lora_init(model: ptr, path_lora: string): ptr; +extern fn llama_adapter_lora_free(adapter: ptr): void; +extern fn llama_set_adapter_lora(ctx: ptr, adapter: ptr, scale: f32): i32; +extern fn llama_rm_adapter_lora(ctx: ptr, adapter: ptr): i32; +extern fn llama_clear_adapter_lora(ctx: ptr): void; + +// Control vectors +extern fn llama_apply_adapter_cvec(ctx: ptr, data: ptr, len: u64, n_embd: i32, il_start: i32, il_end: i32): i32; + +// State/session save/load +extern fn llama_state_get_size(ctx: ptr): u64; +extern fn llama_state_get_data(ctx: ptr, dst: ptr, size: u64): u64; +extern fn llama_state_set_data(ctx: ptr, src: ptr, size: u64): u64; +extern fn llama_state_save_file(ctx: ptr, path: string, tokens: ptr, n_token_count: u64): i32; +extern fn llama_state_load_file(ctx: ptr, path: string, tokens_out: ptr, n_token_capacity: u64, n_token_count_out: ptr): i32; +extern fn llama_state_seq_get_size(ctx: ptr, seq_id: i32): u64; +extern fn llama_state_seq_get_data(ctx: ptr, dst: ptr, size: u64, seq_id: i32): u64; +extern fn llama_state_seq_set_data(ctx: ptr, src: ptr, size: u64, seq_id: i32): u64; +extern fn llama_state_seq_save_file(ctx: ptr, filepath: string, seq_id: i32, tokens: ptr, n_token_count: u64): u64; +extern fn llama_state_seq_load_file(ctx: ptr, filepath: string, seq_id: i32, tokens_out: ptr, n_token_capacity: u64, n_token_count_out: ptr): u64; + +// Additional vocab functions +extern fn llama_vocab_sep(vocab: ptr): i32; +extern fn llama_vocab_mask(vocab: ptr): i32; +extern fn llama_vocab_get_score(vocab: ptr, token: i32): f32; +extern fn llama_vocab_get_attr(vocab: ptr, token: i32): i32; +extern fn llama_vocab_get_add_sep(vocab: ptr): i32; + +// FIM (Fill-in-Middle) tokens +extern fn llama_vocab_fim_pre(vocab: ptr): i32; +extern fn llama_vocab_fim_suf(vocab: ptr): i32; +extern fn llama_vocab_fim_mid(vocab: ptr): i32; +extern fn llama_vocab_fim_pad(vocab: ptr): i32; +extern fn llama_vocab_fim_rep(vocab: ptr): i32; +extern fn llama_vocab_fim_sep(vocab: ptr): i32; + +// Additional model info +extern fn llama_model_n_head_kv(model: ptr): i32; +extern fn llama_model_n_swa(model: ptr): i32; +extern fn llama_model_rope_freq_scale_train(model: ptr): f32; +extern fn llama_model_decoder_start_token(model: ptr): i32; +extern fn llama_model_is_hybrid(model: ptr): i32; +extern fn llama_model_is_diffusion(model: ptr): i32; +extern fn llama_model_rope_type(model: ptr): i32; + +// Model metadata +extern fn llama_model_meta_val_str(model: ptr, key: string, buf: ptr, buf_size: u64): i32; +extern fn llama_model_meta_count(model: ptr): i32; +extern fn llama_model_meta_key_by_index(model: ptr, i: i32, buf: ptr, buf_size: u64): i32; +extern fn llama_model_meta_val_str_by_index(model: ptr, i: i32, buf: ptr, buf_size: u64): i32; + +// Chat templates +extern fn llama_chat_builtin_templates(output: ptr, len: u64): i32; + +// Context settings +extern fn llama_set_embeddings(ctx: ptr, embeddings: i32): void; +extern fn llama_set_causal_attn(ctx: ptr, causal_attn: i32): void; +extern fn llama_set_warmup(ctx: ptr, warmup: i32): void; +extern fn llama_pooling_type(ctx: ptr): i32; + +// Embeddings +extern fn llama_get_embeddings_seq(ctx: ptr, seq_id: i32): ptr; + +// Memory helpers +extern fn llama_memory_can_shift(mem: ptr): i32; +extern fn llama_memory_seq_div(mem: ptr, seq_id: i32, p0: i32, p1: i32, d: i32): void; + +// Additional utilities +extern fn llama_supports_rpc(): i32; +extern fn llama_numa_init(numa: i32): void; +extern fn llama_n_ctx_seq(ctx: ptr): u32; + // ============================================================================ // High-level API // ============================================================================ @@ -854,6 +960,572 @@ export fn complete(model_path: string, prompt: string, options): string { return result; } +// ============================================================================ +// Grammar Sampling +// ============================================================================ + +// Create a grammar-constrained sampler +// grammar_str: GBNF grammar string (see llama.cpp/grammars/README.md) +// grammar_root: name of the root rule (typically "root") +export fn sampler_grammar(model: ptr, grammar_str: string, grammar_root: string): ptr { + let vocab = llama_model_get_vocab(model); + let sampler = llama_sampler_init_grammar(vocab, grammar_str, grammar_root); + if (sampler == null) { + throw "Failed to create grammar sampler - check grammar syntax"; + } + return sampler; +} + +// Create a sampler chain with grammar constraint +// Options: +// grammar: GBNF grammar string +// grammar_root: root rule name (default: "root") +// top_k, top_p, temp, etc.: standard sampling parameters +export fn sampler_create_with_grammar(model: ptr, options): ptr { + let opts = options || {}; + + let chain_params = alloc(8); + memset(chain_params, 0, 8); + let chain = llama_sampler_chain_init(chain_params); + free(chain_params); + + if (chain == null) { + throw "Failed to create sampler chain"; + } + + // Add penalties if specified + let repeat_penalty = opts.repeat_penalty || 1.1; + let repeat_last_n = opts.repeat_last_n || 64; + if (repeat_penalty != 1.0) { + let freq_penalty = opts.freq_penalty || 0.0; + let presence_penalty = opts.presence_penalty || 0.0; + let penalties = llama_sampler_init_penalties(repeat_last_n, repeat_penalty, freq_penalty, presence_penalty); + llama_sampler_chain_add(chain, penalties); + } + + // Add top-k + let top_k = opts.top_k || 40; + if (top_k > 0) { + let top_k_sampler = llama_sampler_init_top_k(top_k); + llama_sampler_chain_add(chain, top_k_sampler); + } + + // Add top-p + let top_p = opts.top_p || 0.95; + if (top_p < 1.0) { + let top_p_sampler = llama_sampler_init_top_p(top_p, 1); + llama_sampler_chain_add(chain, top_p_sampler); + } + + // Add temperature + let temp = opts.temp || 0.8; + if (temp > 0.0) { + let temp_sampler = llama_sampler_init_temp(temp); + llama_sampler_chain_add(chain, temp_sampler); + } + + // Add grammar constraint + if (opts.grammar) { + let grammar_root = opts.grammar_root || "root"; + let vocab = llama_model_get_vocab(model); + let grammar_sampler = llama_sampler_init_grammar(vocab, opts.grammar, grammar_root); + if (grammar_sampler != null) { + llama_sampler_chain_add(chain, grammar_sampler); + } + } + + // Add distribution sampler (final selection) + let seed = opts.seed || DEFAULT_SEED; + let dist = llama_sampler_init_dist(seed); + llama_sampler_chain_add(chain, dist); + + return chain; +} + +// Common JSON grammar for structured output +export let GRAMMAR_JSON = "root ::= object | array\n" + + "object ::= \"{\" ws ( pair ( \",\" ws pair )* )? ws \"}\"\n" + + "pair ::= string ws \":\" ws value\n" + + "array ::= \"[\" ws ( value ( \",\" ws value )* )? ws \"]\"\n" + + "value ::= string | number | object | array | \"true\" | \"false\" | \"null\"\n" + + "string ::= \"\\\"\" ( [^\"\\\\] | \"\\\\\" [\"\\\\nrt] )* \"\\\"\"\n" + + "number ::= \"-\"? ( \"0\" | [1-9] [0-9]* ) ( \".\" [0-9]+ )? ( [eE] [\"+\\-\"]? [0-9]+ )?\n" + + "ws ::= [ \\t\\n]*"; + +// ============================================================================ +// LoRA Adapters +// ============================================================================ + +// Load a LoRA adapter from file +export fn lora_load(model: ptr, path: string): ptr { + let adapter = llama_adapter_lora_init(model, path); + if (adapter == null) { + throw "Failed to load LoRA adapter: " + path; + } + return adapter; +} + +// Free a LoRA adapter +export fn lora_free(adapter: ptr) { + if (adapter != null) { + llama_adapter_lora_free(adapter); + } +} + +// Apply a LoRA adapter to a context +// scale: strength of the adapter (1.0 = full strength) +export fn lora_apply(ctx: ptr, adapter: ptr, scale: f32): i32 { + return llama_set_adapter_lora(ctx, adapter, scale); +} + +// Remove a specific LoRA adapter from context +export fn lora_remove(ctx: ptr, adapter: ptr): i32 { + return llama_rm_adapter_lora(ctx, adapter); +} + +// Remove all LoRA adapters from context +export fn lora_clear(ctx: ptr) { + llama_clear_adapter_lora(ctx); +} + +// ============================================================================ +// State/Session Management +// ============================================================================ + +// Get the size needed to save the context state +export fn state_size(ctx: ptr): u64 { + return llama_state_get_size(ctx); +} + +// Save context state to a file +// Returns true on success +export fn state_save(ctx: ptr, path: string, tokens: array): bool { + let n_tokens = tokens.length; + let tokens_buf = alloc(n_tokens * 4); + for (let i = 0; i < n_tokens; i = i + 1) { + ptr_write_i32(ptr_offset(tokens_buf, i, 4), tokens[i]); + } + + let result = llama_state_save_file(ctx, path, tokens_buf, n_tokens); + free(tokens_buf); + + return result != 0; +} + +// Load context state from a file +// Returns: { tokens: array, success: bool } +export fn state_load(ctx: ptr, path: string, max_tokens: i32): object { + let tokens_buf = alloc(max_tokens * 4); + let n_tokens_ptr = alloc(8); + + let result = llama_state_load_file(ctx, path, tokens_buf, max_tokens, n_tokens_ptr); + + if (result == 0) { + free(tokens_buf); + free(n_tokens_ptr); + return { tokens: [], success: false }; + } + + let n_tokens = ptr_deref_u64(n_tokens_ptr); + free(n_tokens_ptr); + + let tokens = []; + for (let i = 0; i < n_tokens; i = i + 1) { + tokens.push(ptr_deref_i32(ptr_offset(tokens_buf, i, 4))); + } + free(tokens_buf); + + return { tokens: tokens, success: true }; +} + +// Save a single sequence to file +export fn state_seq_save(ctx: ptr, path: string, seq_id: i32, tokens: array): bool { + let n_tokens = tokens.length; + let tokens_buf = alloc(n_tokens * 4); + for (let i = 0; i < n_tokens; i = i + 1) { + ptr_write_i32(ptr_offset(tokens_buf, i, 4), tokens[i]); + } + + let result = llama_state_seq_save_file(ctx, path, seq_id, tokens_buf, n_tokens); + free(tokens_buf); + + return result > 0; +} + +// Load a single sequence from file +export fn state_seq_load(ctx: ptr, path: string, seq_id: i32, max_tokens: i32): object { + let tokens_buf = alloc(max_tokens * 4); + let n_tokens_ptr = alloc(8); + + let result = llama_state_seq_load_file(ctx, path, seq_id, tokens_buf, max_tokens, n_tokens_ptr); + + if (result == 0) { + free(tokens_buf); + free(n_tokens_ptr); + return { tokens: [], success: false }; + } + + let n_tokens = ptr_deref_u64(n_tokens_ptr); + free(n_tokens_ptr); + + let tokens = []; + for (let i = 0; i < n_tokens; i = i + 1) { + tokens.push(ptr_deref_i32(ptr_offset(tokens_buf, i, 4))); + } + free(tokens_buf); + + return { tokens: tokens, success: true }; +} + +// ============================================================================ +// Embeddings +// ============================================================================ + +// Get embeddings for a sequence +// Returns null if pooling is not enabled +export fn get_embeddings(ctx: ptr, seq_id: i32): ptr { + return llama_get_embeddings_seq(ctx, seq_id); +} + +// Extract embeddings as an array of floats +export fn extract_embeddings(ctx: ptr, model: ptr, seq_id: i32): array { + let embd_ptr = llama_get_embeddings_seq(ctx, seq_id); + if (embd_ptr == null) { + return []; + } + + let n_embd = llama_model_n_embd(model); + let embeddings = []; + for (let i = 0; i < n_embd; i = i + 1) { + // Read f32 values from the embeddings pointer + let bytes = []; + for (let j = 0; j < 4; j = j + 1) { + bytes.push(ptr_deref_u8(ptr_offset(embd_ptr, i * 4 + j, 1))); + } + // Convert bytes to float (little-endian IEEE 754) + embeddings.push(__bytes_to_f32(bytes)); + } + + return embeddings; +} + +// Enable/disable embeddings extraction +export fn set_embeddings(ctx: ptr, enabled: bool) { + llama_set_embeddings(ctx, enabled ? 1 : 0); +} + +// Get the pooling type of the context +export fn get_pooling_type(ctx: ptr): i32 { + return llama_pooling_type(ctx); +} + +// ============================================================================ +// FIM (Fill-in-Middle) for Code Completion +// ============================================================================ + +// Get FIM special tokens +export fn fim_tokens(model: ptr): object { + let vocab = llama_model_get_vocab(model); + return { + pre: llama_vocab_fim_pre(vocab), // Prefix token + suf: llama_vocab_fim_suf(vocab), // Suffix token + mid: llama_vocab_fim_mid(vocab), // Middle token + pad: llama_vocab_fim_pad(vocab), // Padding token + rep: llama_vocab_fim_rep(vocab), // Repository token + sep: llama_vocab_fim_sep(vocab) // Separator token + }; +} + +// Check if model supports FIM +export fn supports_fim(model: ptr): bool { + let vocab = llama_model_get_vocab(model); + let pre = llama_vocab_fim_pre(vocab); + let suf = llama_vocab_fim_suf(vocab); + let mid = llama_vocab_fim_mid(vocab); + return pre != TOKEN_NULL && suf != TOKEN_NULL && mid != TOKEN_NULL; +} + +// Create FIM prompt for code infilling +// prefix: code before cursor +// suffix: code after cursor +export fn fim_prompt(model: ptr, prefix: string, suffix: string): string { + let fim = fim_tokens(model); + if (fim.pre == TOKEN_NULL) { + throw "Model does not support FIM"; + } + + let pre_piece = token_to_piece(model, fim.pre, true); + let suf_piece = token_to_piece(model, fim.suf, true); + let mid_piece = token_to_piece(model, fim.mid, true); + + return pre_piece + prefix + suf_piece + suffix + mid_piece; +} + +// ============================================================================ +// Chat Templates +// ============================================================================ + +// Get the model's default chat template +export fn chat_template(model: ptr): string { + let tmpl_ptr = llama_model_chat_template(model, null); + if (tmpl_ptr == null) { + return ""; + } + return ptr_to_string(tmpl_ptr); +} + +// Apply chat template to format messages +// messages: array of { role: string, content: string } +// add_assistant: whether to add assistant prompt at end +export fn chat_format(model: ptr, messages: array, add_assistant: bool): string { + let tmpl_ptr = llama_model_chat_template(model, null); + if (tmpl_ptr == null) { + throw "Model does not have a chat template"; + } + let tmpl = ptr_to_string(tmpl_ptr); + + // Allocate chat messages array + // Each llama_chat_message is 16 bytes (2 pointers) + let n_msg = messages.length; + let chat_buf = alloc(n_msg * 16); + let string_bufs = []; + + for (let i = 0; i < n_msg; i = i + 1) { + let msg = messages[i]; + + // Allocate and write role string + let role_bytes = msg.role.bytes; + let role_buf = alloc(role_bytes.length + 1); + for (let j = 0; j < role_bytes.length; j = j + 1) { + ptr_write_u8(ptr_offset(role_buf, j, 1), role_bytes[j]); + } + ptr_write_u8(ptr_offset(role_buf, role_bytes.length, 1), 0); + string_bufs.push(role_buf); + + // Allocate and write content string + let content_bytes = msg.content.bytes; + let content_buf = alloc(content_bytes.length + 1); + for (let j = 0; j < content_bytes.length; j = j + 1) { + ptr_write_u8(ptr_offset(content_buf, j, 1), content_bytes[j]); + } + ptr_write_u8(ptr_offset(content_buf, content_bytes.length, 1), 0); + string_bufs.push(content_buf); + + // Write pointers to chat struct + ptr_write_ptr(ptr_offset(chat_buf, i * 16, 1), role_buf); + ptr_write_ptr(ptr_offset(chat_buf, i * 16 + 8, 1), content_buf); + } + + // Allocate output buffer + let buf_size = 4096; + let out_buf = alloc(buf_size); + + let result_len = llama_chat_apply_template(tmpl, chat_buf, n_msg, add_assistant ? 1 : 0, out_buf, buf_size); + + if (result_len < 0) { + // Buffer too small, reallocate + free(out_buf); + buf_size = -result_len + 64; + out_buf = alloc(buf_size); + result_len = llama_chat_apply_template(tmpl, chat_buf, n_msg, add_assistant ? 1 : 0, out_buf, buf_size); + } + + // Read result + let result = ""; + if (result_len > 0) { + result = ptr_to_string_n(out_buf, result_len); + } + + // Cleanup + free(out_buf); + free(chat_buf); + for (let i = 0; i < string_bufs.length; i = i + 1) { + free(string_bufs[i]); + } + + return result; +} + +// ============================================================================ +// Model Metadata +// ============================================================================ + +// Get model metadata value by key +export fn model_meta(model: ptr, key: string): string { + let buf = alloc(512); + let len = llama_model_meta_val_str(model, key, buf, 512); + + if (len < 0) { + free(buf); + return ""; + } + + let result = ptr_to_string_n(buf, len); + free(buf); + return result; +} + +// Get number of metadata key-value pairs +export fn model_meta_count(model: ptr): i32 { + return llama_model_meta_count(model); +} + +// Get all model metadata as an object +export fn model_meta_all(model: ptr): object { + let count = llama_model_meta_count(model); + let result = {}; + + let key_buf = alloc(256); + let val_buf = alloc(512); + + for (let i = 0; i < count; i = i + 1) { + let key_len = llama_model_meta_key_by_index(model, i, key_buf, 256); + if (key_len > 0) { + let key = ptr_to_string_n(key_buf, key_len); + let val_len = llama_model_meta_val_str_by_index(model, i, val_buf, 512); + if (val_len > 0) { + result[key] = ptr_to_string_n(val_buf, val_len); + } + } + } + + free(key_buf); + free(val_buf); + return result; +} + +// Extended model info +export fn model_info_extended(model: ptr): object { + let desc_buf = alloc(256); + llama_model_desc(model, desc_buf, 256); + let desc = ptr_to_string(desc_buf); + free(desc_buf); + + return { + description: desc, + n_ctx_train: llama_model_n_ctx_train(model), + n_embd: llama_model_n_embd(model), + n_layer: llama_model_n_layer(model), + n_head: llama_model_n_head(model), + n_head_kv: llama_model_n_head_kv(model), + n_swa: llama_model_n_swa(model), + size: llama_model_size(model), + n_params: llama_model_n_params(model), + has_encoder: llama_model_has_encoder(model) != 0, + has_decoder: llama_model_has_decoder(model) != 0, + is_recurrent: llama_model_is_recurrent(model) != 0, + is_hybrid: llama_model_is_hybrid(model) != 0, + is_diffusion: llama_model_is_diffusion(model) != 0, + rope_type: llama_model_rope_type(model), + rope_freq_scale: llama_model_rope_freq_scale_train(model), + decoder_start_token: llama_model_decoder_start_token(model) + }; +} + +// ============================================================================ +// Additional Samplers +// ============================================================================ + +// Create XTC sampler +export fn sampler_xtc(p: f32, t: f32, min_keep: u64, seed: u32): ptr { + return llama_sampler_init_xtc(p, t, min_keep, seed); +} + +// Create Top-N-Sigma sampler +export fn sampler_top_n_sigma(n: f32): ptr { + return llama_sampler_init_top_n_sigma(n); +} + +// Create infill sampler for FIM +export fn sampler_infill(model: ptr): ptr { + let vocab = llama_model_get_vocab(model); + return llama_sampler_init_infill(vocab); +} + +// ============================================================================ +// Additional Vocab Functions +// ============================================================================ + +// Get separator token +export fn vocab_sep(model: ptr): i32 { + let vocab = llama_model_get_vocab(model); + return llama_vocab_sep(vocab); +} + +// Get mask token (for BERT-style models) +export fn vocab_mask(model: ptr): i32 { + let vocab = llama_model_get_vocab(model); + return llama_vocab_mask(vocab); +} + +// Get token score +export fn token_score(model: ptr, token: i32): f32 { + let vocab = llama_model_get_vocab(model); + return llama_vocab_get_score(vocab, token); +} + +// Get token attributes +export fn token_attr(model: ptr, token: i32): i32 { + let vocab = llama_model_get_vocab(model); + return llama_vocab_get_attr(vocab, token); +} + +// Check if model adds separator token +export fn vocab_add_sep(model: ptr): bool { + let vocab = llama_model_get_vocab(model); + return llama_vocab_get_add_sep(vocab) != 0; +} + +// ============================================================================ +// Context Settings +// ============================================================================ + +// Set causal attention mode +export fn set_causal_attn(ctx: ptr, causal: bool) { + llama_set_causal_attn(ctx, causal ? 1 : 0); +} + +// Set warmup mode (for pre-loading weights) +export fn set_warmup(ctx: ptr, warmup: bool) { + llama_set_warmup(ctx, warmup ? 1 : 0); +} + +// ============================================================================ +// Additional KV Cache Functions +// ============================================================================ + +// Divide positions in a sequence +export fn kv_cache_seq_div(ctx: ptr, seq_id: i32, p0: i32, p1: i32, d: i32) { + let mem = llama_get_memory(ctx); + llama_memory_seq_div(mem, seq_id, p0, p1, d); +} + +// Check if memory supports shifting +export fn memory_can_shift(ctx: ptr): bool { + let mem = llama_get_memory(ctx); + return llama_memory_can_shift(mem) != 0; +} + +// ============================================================================ +// System Capabilities +// ============================================================================ + +// Check RPC support +export fn supports_rpc(): bool { + return llama_supports_rpc() != 0; +} + +// Initialize NUMA (call before loading models) +// numa: 0 = disabled, 1 = distribute, 2 = isolate, 3 = numactl, 4 = mirror +export fn numa_init(numa: i32) { + llama_numa_init(numa); +} + +// Get context sequence size +export fn n_ctx_seq(ctx: ptr): u32 { + return llama_n_ctx_seq(ctx); +} + // ============================================================================ // Helper functions for pointer operations // ============================================================================