llama_cpp-rs

A LlamaSession represents an evaluation session for a LlamaModel. It maintains the model's context (tokens and KV cache) and can be used to advance the context or start token generation. You can create multiple independent sessions from a single model using LlamaModel::create_session.

Key capabilities:

• Advancing Context: Feed new tokens or raw bytes into the session to update the model's history.
• Completion: Start generating new tokens using a sampler.
• Context Manipulation: Truncate, remove specific ranges of tokens, or reset the context entirely.
• Deep Copying: Create an independent copy of a session (including its state) using deep_copy.

// Assuming a LlamaModel has been loaded
let mut session = model.create_session(params)?;

// Advance context with bytes
session.advance_context("Hello, how are you?")?;

// Start generating tokens
let mut handle = session.start_completing()?;

To use llama_cpp-rs, you need to create a LlamaModel to hold the weights and a LlamaSession to manage the context and generation. A LlamaModel is large and shared, while a LlamaSession is relatively small and used for individual interactions.

To generate text, you can use ctx.start_completing_with which returns an iterator of strings. Note that the completion process runs in a worker thread and will stop automatically when the completion handle is dropped.

use llama_cpp::{LlamaModel, LlamaParams, SessionParams};
use llama_cpp::standard_sampler::StandardSampler;
use std::io::{self, Write};

// 1. Load the model
let model = LlamaModel::load_from_file("path_to_model.gguf", LlamaParams::default()).expect("Could not load model");

// 2. Create a session
let mut ctx = model.create_session(SessionParams::default()).expect("Failed to create session");

// 3. Feed context
ctx.advance_context("This is the story of a man named Stanley.").unwrap();

// 4. Generate tokens using a sampler
let max_tokens = 1024;
let mut decoded_tokens = 0;
let mut completions = ctx.start_completing_with(StandardSampler::default(), 1024).into_strings();

for completion in completions {
    print!("{completion}");
    let _ = io::stdout().flush();

    decoded_tokens += 1;
    if decoded_tokens > max_tokens {
        break;
    }
}

You can run GGUF-based large language models by loading a model file into a LlamaModel, creating a Session (context), and using start_completing_with to generate tokens.

Key concepts:

• LlamaModel: Holds the heavy weights shared across multiple sessions.
• Session: A lightweight context (typically 10s to 100s of MBs) used for inference.
• start_completing_with: Spawns a worker thread for token generation. Dropping the completion handle stops generation.

// Create a model from anything that implements `AsRef<Path>`:
let model = LlamaModel::load_from_file("path_to_model.gguf", LlamaParams::default()).expect("Could not load model");

// A `LlamaModel` holds the weights shared across many _sessions_; while your model may be
// several gigabytes large, a session is typically a few dozen to a hundred megabytes!
let mut ctx = model.create_session(SessionParams::default()).expect("Failed to create session");

// You can feed anything that implements `AsRef<[u8]>` into the model's context.
ctx.advance_context("This is the story of a man named Stanley.").unwrap();

// LLMs are typically used to predict the next word in a sequence. Let's generate some tokens!
let max_tokens = 1024;
let mut decoded_tokens = 0;

// `ctx.start_completing_with` creates a worker thread that generates tokens. When the completion
// handle is dropped, tokens stop generating!
let mut completions = ctx.start_completing_with(StandardSampler::default(), 1024).into_strings();

for completion in completions {
    print!("{completion}");
    let _ = io::stdout().flush();
    
    decoded_tokens += 1;
    
    if decoded_tokens > max_tokens {
        break;
    }
}

The Batch struct is a safe wrapper around llama_batch used to prepare groups of tokens for model inference. It manages the memory for tokens, positions, sequence IDs, and logits required by the underlying llama.cpp engine.

Key Methods

• new(capacity: usize, embed: usize, max_sequences: usize) -> Self: Initializes a new batch.
  • capacity: The maximum number of tokens this batch can hold.
  • embed: The embedding dimension (used if the batch is for embedding tasks).
  • max_sequences: The maximum number of sequences that can be generated/processed in this batch.
  • Note: Panics if capacity or max_sequences is 0.
• add(&mut self, token: Token, position: usize, sequence_ids: &[i32], logits: bool) -> usize: Adds a token to the batch.
  • token: The Token to add.
  • position: The position of the token in the sequence.
  • sequence_ids: A slice of sequence IDs for the token.
  • logits: A boolean indicating whether to compute logits for this token.
  • Returns: The index of the added token, or usize::MAX if the batch is at capacity or exceeds max_sequences limits.
• clear(&mut self): Resets the number of tokens in the batch to 0, allowing the batch to be reused.
• set_logits(&self, idx: usize, value: bool): Manually sets whether logits should be computed for a specific token index.
• tokens(&self) -> usize: Returns the current number of tokens in the batch.
• handle(&self) -> llama_batch: Returns the underlying raw llama_batch pointer for low-level FFI operations.

// Example: Creating and adding tokens to a batch
let mut batch = Batch::new(512, 128, 1);

// Adding a token
let token_idx = batch.add(my_token, 0, &[0], true);
if token_idx == usize::MAX {
    // Handle capacity error
}

// Checking current token count
let count = batch.tokens();