WordLlama NLP Utility

WordLlama provides a suite of NLP utilities optimized for speed and low resource usage:

• Fast Embeddings: Efficient generation via token lookup and average pooling.
• Similarity & Ranking: Compute cosine similarity and rank documents or retrieve Top-K results.
• Deduplication & Clustering: Perform fuzzy deduplication and KMeans clustering.
• Semantic Text Splitting: Split text into semantically coherent chunks.
• Advanced Representations: Support for Binary Embeddings (Hamming similarity) and Matryoshka Representations (dimension truncation).

WordLlama is a lightweight NLP toolkit optimized for CPU inference. It creates compact word representations by extracting token embedding codebooks from large language models (like LLaMA 2 or 3) and training a small context-less model.

Key technical characteristics:

• Matryoshka Representations: Supports truncating embedding dimensions for flexibility.
• Low Resource Requirements: Uses simple token lookup with average pooling, requiring only NumPy for inference.
• Binary Embeddings: Supports packing embeddings into small integer arrays for fast Hamming distance calculations.
• Performance: Outperforms traditional models like GloVe 300d on MTEB benchmarks while remaining significantly smaller (e.g., 16MB default model).

Install the WordLlama toolkit using pip to get started with fast, lightweight NLP tasks like similarity computation and ranking.

pip install wordllama

To perform model surgery or extract embeddings for training purposes, you must install the wordllama package with the [train] extra to ensure all necessary training dependencies are present.

pip install wordllama[train]

To train a model, install the training extras and use the train.py script with a configuration file.

1. Install: pip install wordllama[train]
2. Start training: python train.py train --config your_new_config
3. Save model: python train.py save --config your_new_config --checkpoint ... --outdir /path/to/weights/ (This saves one model per Matryoshka dimension).

Semantic splitting is a three-step process used to prepare text for Retrieval-Augmented Generation (RAG) applications. This method aims to maximize information continuity and produce consistent chunk sizes by using embeddings to identify semantic boundaries.

The Semantic Splitting Workflow

1. Split: Divide the original text into small, grammatically intact units (typically at the paragraph or sentence level).
2. Embed: Generate vector embeddings for each small chunk using WordLlama to represent its semantic meaning.
3. Reconstruct: Combine adjacent chunks based on their embedding similarity. This ensures that semantically related content stays together in a single chunk while maintaining a target size (e.g., 256-2048 tokens).

To use a specific model (like Gemma2 2B) for embedding extraction, you must create a configuration file in wordllama/config. You can base your new config on an existing .toml file.

Key fields include:

• dim: The size of the embedding vector.
• n_vocab: The vocabulary size.
• hf_model_id: The Hugging Face model identifier.
• pad_token: A special token from the model's tokenizer (WordLlama's average pooling ignores these tokens).

[model]
dim = 2304
n_vocab = 256000
hf_model_id = "google/gemma-2-2b-it"
pad_token = ""

Use wl.vector_similarity(embeddings_a, embeddings_b) to compute the similarity between two sets of embeddings. This is useful for identifying semantic relationships or blocks of similar text in a dataset.

# calculate the cross-similarity
embeddings = wl.embed(lines, norm=True)
xsim = wl.vector_similarity(embeddings, embeddings)

Use WordLlama.load() to initialize the model and the .embed() method to convert a list of strings into embeddings. You can use the trunc_dim parameter during loading to reduce the embedding dimensionality (e.g., to 64).

from wordllama import WordLlama

# Load pre-trained embeddings (truncate dimension to 64)
wl = WordLlama.load(trunc_dim=64)

# Embed text
embeddings = wl.embed(["The quick brown fox jumps over the lazy dog", "And all that jazz"])
print(embeddings.shape)  # Output: (2, 64)

The .key(query) method returns a Callable[[str], float]. This function can be passed directly to standard Python functions like sorted() or max() to rank or find the most similar candidate to the original query based on semantic similarity.

query = "Machine learning methods"
candidates = [
    "Foundations of neural science",
    "Introduction to neural networks",
    "Cooking delicious pasta at home",
    "Introduction to philosophy: logic",
]

# Returns a Callable[[str], float] function
sim_key = wl.key(query)

# Sort candidates, most similar first
sorted_candidates = sorted(candidates, key=sim_key, reverse=True)

# Most similar candidate
best_candidate = max(candidates, key=sim_key)