Usage

// Compile a vocabulary and index
std::string schema = "{\"type\":\"object\",\"properties\":{\"x\":{\"type\":\"integer\"}}}";
std::string model = "NousResearch/Hermes-3-Llama-3.1-8B";
dotjson::Vocabulary vocabulary(model);
dotjson::Index index(schema, vocabulary);

// Save the index
index.serialize_to_disk(path);

// Load the index
std::filesystem::path path = "path/to/file";
dotjson::Index index(path);

// Create a guide from the index with the batch size
std::size_t batch_size = 1;
dotjson::Guide guide(index, batch_size);

// Check if a specific token is allowed in each batch
u_int32_t token_to_check = 42;
std::vector<bool> is_allowed = token_set.contains(token_to_check);
// is_allowed[i] is true if token 42 is allowed in batch i

// Get the number of allowed tokens in each batch
std::vector<std::size_t> allowed_count = token_set.num_allowed();
// allowed_count[i] contains the number of allowed tokens in batch i

// Get the vocabulary size
u_int32_t vocab_size = vocabulary.max_token_id();

// Allocate a single vector for all batches
std::vector<u_int16_t> all_logits(n_batches * vocab_size, 1);
std::vector<std::span<u_int16_t>> logits(n_batches);

// Create non-overlapping spans pointing to sections of the vector
for (int i = 0; i < n_batches; ++i) {
    logits[i] = std::span<u_int16_t>(all_logits.data() + i * vocab_size, vocab_size);
}

// Store the vectors in a container that persists outside the loop
std::vector<std::vector<u_int16_t>> batch_logits_storage(n_batches);
std::vector<std::span<u_int16_t>> logits(n_batches);

// Initialize each batch's logits:
for (int i = 0; i < n_batches; ++i) {
    batch_logits_storage[i] = std::vector<u_int16_t>(vocab_size, 1);
    logits[i] = std::span<u_int16_t>(batch_logits_storage[i].data(), batch_logits_storage[i].size());
}

// Set the batch size
std::size_t batch_size = 1;

// The value to use for masking (aka disabling tokens)
// This will be determined by your quantization scheme,
// but in this example we use 0
u_int16_t mask_value = 0;

// Create the processor using the guide object
dotjson::LogitsProcessor processor(guide, mask_value);

processor(logits, token_set);

// Single logits vector, with a three token vocabulary
logits_vec = { {1, 2, 3} };

// Import dotjson
// Note: The include path may vary depending on your installation method
// For system installation, use: #include "dotjson.hpp"
// For local installation, use: #include "vendor/dotjson/include/dotjson.hpp"
#include "../src/dotjson.hpp"
#include <iostream>
#include <vector>
#include <limits>
#include <span>

// A function that returns logits in the appropriate form,
// taking the full sequence history as input for the model's forward pass
std::vector<std::span<u_int16_t>> get_logits(std::string model, const std::vector<std::vector<u_int32_t>>& sequences) {
    // This is your language model's forward pass.
    // It should use the sequences (token history) to compute the next token logits
    // ...

    // Return the computed logits
    std::vector<std::span<u_int16_t>> logits;
    // ... populate logits from your forward pass ...
    return logits;
}

std::vector<u_int32_t> sample_tokens(std::vector<std::span<u_int16_t>> &logits) {
    // A function that returns the next tokens generated from a given set of logits, such as multinomial or greedy sampling.
    // ...
    return sampled_tokens;
}

int main() {
    // Specify the model to use -- this downloads the tokenizer
    // from HuggingFace.
    std::string model = "gpt2";

    // Specify a JSON schema to use.
    std::string schema = "{\"type\":\"object\",\"properties\":{\"x\":{\"type\":\"integer\"}}}";

    // Compile the index and vocabulary for the schema.
    dotjson::Vocabulary vocabulary(model);
    dotjson::Index index(schema, vocabulary);

    // Specify the mask value to use for disabling tokens.
    // Tokens that don't match the schema will have their logit values set to this.
    u_int16_t mask_value = 0;
    std::size_t batch_size = 1;

    // Create a guide to generate sets of allowed tokens
    dotjson::Guide guide(index, batch_size);

    // Create a logit processor using the guide
    dotjson::LogitsProcessor processor(guide, mask_value);

    // Get the initial set of allowed tokens
    dotjson::BatchedTokenSet token_set = guide.get_start_tokensets();

    // Initialize sequence tracking for each batch
    std::vector<std::vector<u_int32_t>> sequences(batch_size);

    // Maximum allowed sequence length
    const size_t max_sequence_length = 1024;

    // Initialize the first logits based on empty sequences
    // (this is your language model's forward pass)
    std::vector<std::span<u_int16_t>> logits = get_logits(model, sequences);

    // Process the initial logits with the token set
    processor(logits, token_set);

    // Sample the first set of tokens
    std::vector<u_int32_t> sampled_tokens = sample_tokens(logits);

    // Add sampled tokens to the sequences history
    for (size_t i = 0; i < batch_size; i++) {
        sequences[i].push_back(sampled_tokens[i]);
    }

    // Tracking boolean to know when to end generation
    bool is_completed = false;

    // A complete inference loop would look like:
    while (!is_completed) {
        // Get the next set of allowed tokens based on the
        // most recently sampled tokens
        token_set = guide.get_next_tokensets(sampled_tokens);

        // Get the next set of logits from the model
        // The full sequences history is used for the model's context
        logits = get_logits(model, sequences);

        // Process logits with the current token set
        processor(logits, token_set);

        // After processing, only tokens that are valid according to the schema
        // will retain their original values. Others will be set to mask_value.

        // Sample the next tokens
        sampled_tokens = sample_tokens(logits);

        // Add the new tokens to the sequence history
        for (size_t i = 0; i < batch_size; i++) {
            sequences[i].push_back(sampled_tokens[i]);

            // Check for completion (e.g., EOS token)
            // In this case there is only one EOS token, but in general
            // you need to check the entire vector.
            if (sampled_tokens[i] == vocabulary.eos_token_ids()[0]) {
                is_completed = true;
            }
        }
    }

    // Process the generated sequences to return the final
    // completion to the user
    for (size_t i = 0; i < batch_size; i++) {
        // Convert token IDs back to text
        // ...
    }
}