Skip to content

Commit 5068a80

Browse files
JamePengJamePeng
committed
Upload /docs/wiki/LlamaEmbedding.md for llama_embedding.py
Signed-off-by: JamePeng <jame_peng@sina.com>
1 parent 9d23199 commit 5068a80

3 files changed

Lines changed: 264 additions & 1 deletion

File tree

‎docs/wiki/SCHEMA.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
**Purpose**: Maintain a living, always-up-to-date, structured documentation wiki for the llama-cpp-python library using LLMs as the primary maintainer.
44

55
**Core Principles**:
6-
- The source of truth is the latest code in `llama_cpp/` (especially `llama.py`, `llama_chat_format.py`, `llama_cpp.py`, `llama_types.py`, `mtmd_cpp.py`, `_internals.py`, `_ggml.py`).
6+
- The source of truth is the latest code in `llama_cpp/` (especially `llama.py`, `llama_chat_format.py`, `llama_cpp.py`, `llama_types.py`, `llama_embedding.py`, `mtmd_cpp.py`, `_internals.py`, `_ggml.py`).
77
- Never invent parameters or behavior. Always read the current source code before writing/updating a page.
88
- All examples must be complete, runnable with the latest API, and include necessary imports.
99
- Clearly mark any deprecated/old usage with a warning and show the modern replacement.

‎docs/wiki/core/ChatHandler.md‎

Whitespace-only changes.

‎docs/wiki/core/LlamaEmbedding.md‎

Lines changed: 263 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,263 @@
1+
---
2+
title: LlamaEmbedding
3+
class_name: LlamaEmbedding
4+
last_updated: 2026-04-27
5+
version_target: "latest"
6+
---
7+
8+
# LlamaEmbedding
9+
10+
## Overview
11+
12+
`LlamaEmbedding` is a specialized class for high-performance Text Embedding and Reranking. It inherits from the base `Llama` class but is optimized for vector operations.
13+
14+
### Support Embeddings & Rerank Model:
15+
16+
17+
| Model | Type | Link | Status |
18+
|--------------------|-----------|--------------------------------------------------------|--------------|
19+
| `bge-m3` | Embedding |[bge-m3-GGUF](https://huggingface.co/gpustack/bge-m3-GGUF) | Useful ✅ |
20+
|`bge-reranker-v2-m3`| Rerank |[bge-reranker-v2-m3-GGUF](https://huggingface.co/gpustack/bge-reranker-v2-m3-GGUF) | Useful ✅ |
21+
|`qwen3-reranker`| Rerank |[Qwen3-Reranker-GGUF](https://huggingface.co/JamePeng2023/Qwen3-Reranker-GGUF) | Useful ✅ |
22+
23+
**Core Features:**
24+
1. **Auto-configuration**: Automatically sets `embeddings=True`.
25+
2. **Streaming Batch**: Handles massive datasets without OOM (Out Of Memory).
26+
3. **Native Reranking Support**: Specifically handles `LLAMA_POOLING_TYPE_RANK` models (like BGE-Reranker, Qwen3-Reranker). It correctly identifies classification heads to output scalar relevance scores instead of high-dimensional vectors.
27+
4. **Advanced Normalization**: Implements MaxInt16, Taxicab (L1), and Euclidean (L2) normalization strategies using NumPy for optimal performance and compatibility with various vector databases.
28+
29+
## Constructor `__init__`
30+
31+
| Parameter | Type | Default | Description |
32+
|-----------|------|---------|-------------|
33+
| `model_path` | str | Required | Path to the GGUF model file. |
34+
| `n_ctx` | int | 0 | Text context window size (0 = model default). |
35+
| `n_batch` | int | 512 | Maximum prompt processing batch size. |
36+
| `n_ubatch` | int | 512 | Physical batch size. |
37+
| `pooling_type` | int | `LLAMA_POOLING_TYPE_UNSPECIFIED` (-1) | Pooling strategy used by the model: `LLAMA_POOLING_TYPE_RANK` (4) for rerankers, `LLAMA_POOLING_TYPE_UNSPECIFIED` (-1) for embeddings. |
38+
| `n_gpu_layers` | int | 0 | Number of layers offloaded to GPU (0 = CPU only, -1 = all layers). |
39+
| `verbose` | bool | True | Whether to print debug information. |
40+
| `**kwargs` | Any | — | Extra arguments passed to the `Llama` base class (e.g., `n_batch`, `n_ctx`, `verbose`). |
41+
42+
### Initialization Logic
43+
44+
1. Forces `embeddings=True` to enable embedding support.
45+
2. Sets `kv_unified=True` to enable unified KV Cache, allowing arbitrary sequence IDs in a batch without "invalid seq_id" errors.
46+
3. Passes `pooling_type` to the parent class constructor.
47+
48+
## Core Methods
49+
50+
### `embed(input, normalize=NORM_MODE_EUCLIDEAN, truncate=True, separator=None, return_count=False)`
51+
52+
**Description**: Computes embedding vectors for input text (standard embeddings or reranking scores).
53+
54+
| Parameter | Type | Default | Description |
55+
|-----------|------|---------|-------------|
56+
| `input` | `Union[str, List[str], List[List[int]]]` | — | Input format: string (can be split), list of strings, or list of integer lists (token IDs). |
57+
| `normalize` | int | `NORM_MODE_EUCLIDEAN` (2) | Vector normalization mode (see below). |
58+
| `truncate` | bool | True | Whether to truncate input. |
59+
| `separator` | str | None | Separator for splitting string input into multiple documents. |
60+
| `return_count` | bool | False | If True, returns `(embeddings, token_count)`. |
61+
62+
**Normalization Modes:**
63+
- `NORM_MODE_NONE` (-1): No normalization.
64+
- `NORM_MODE_MAX_INT16` (0): Max absolute value normalization (scaled to 32760).
65+
- `NORM_MODE_TAXICAB` (1): L1 Taxicab norm.
66+
- `NORM_MODE_EUCLIDEAN` (2): L2 Euclidean norm.
67+
- `NORM_MODE_PNORM` (>2): p-norm normalization.
68+
69+
**Returns:**
70+
- `return_count=False`: List of embedding vectors.
71+
- `return_count=True`: Tuple `(embeddings, token_count)`.
72+
73+
**Internal Logic:**
74+
1. Determines mode based on `pooling_type`: `LLAMA_POOLING_TYPE_NONE` (token-level), `LLAMA_POOLING_TYPE_RANK` (rerank), or other (sequence-level).
75+
2. Uses streaming batch decoding to process embeddings in chunks.
76+
3. For token-level mode, extracts and normalizes per-token vectors.
77+
4. For sequence-level mode, extracts sequence vectors and normalizes.
78+
5. Supports `separator` for splitting input into multiple documents.
79+
80+
### `rank(query, documents)`
81+
82+
**Description**: Calculates relevance scores for a list of documents against a query using a Reranking model.
83+
84+
| Parameter | Type | Description |
85+
|-----------|------|-------------|
86+
| `query` | str | Search query string. |
87+
| `documents` | `List[str]` | List of candidate document strings to be scored. |
88+
89+
**Returns**: List of float scores, where higher values indicate greater relevance.
90+
91+
**Internal Logic:**
92+
1. Checks if model is a reranker (`pooling_type == LLAMA_POOLING_TYPE_RANK`).
93+
2. Attempts to retrieve the built-in 'rerank' chat template.
94+
3. If template exists: dynamically replaces `{query}` and `{document}` and tokenizes; otherwise, manually constructs `[BOS] Query [SEP] Doc [EOS]` sequence.
95+
4. Executes embedding inference (`embed`), returning raw logits/scores.
96+
5. For generative rerankers (e.g., Qwen3-Reranker, output dim = 2), uses `yes_logit` as relevance score.
97+
98+
### `create_embedding(input, model=None, normalize=NORM_MODE_EUCLIDEAN, output_format="json")`
99+
100+
**Description**: High-level API compatible with OpenAI format.
101+
102+
| Parameter | Type | Default | Description |
103+
|-----------|------|---------|-------------|
104+
| `input` | `Union[str, List[str]]` | — | Input text or list of texts. |
105+
| `model` | str | None | Model name (optional, uses `self.model_path` if None). |
106+
| `normalize` | int | `NORM_MODE_EUCLIDEAN` (2) | Normalization mode. |
107+
| `output_format` | str | "json" | Output format: 'json', 'json+', or 'array'. |
108+
109+
**Output Formats:**
110+
- `'json'`: OpenAI-style dictionary list.
111+
- `'json+'`: OpenAI dictionary list + cosine similarity matrix.
112+
- `'array'`: Raw Python list (`List[float]` or `List[List[float]]`).
113+
114+
**Returns**: Data structure according to `output_format`.
115+
116+
## Deprecated / Changed APIs
117+
118+
- **Note**: The TODO comments `# TODO(JamePeng): Needs more extensive testing with various embedding and reranking models.` indicate that support for various embedding and reranking models may be incomplete. Further testing is recommended.
119+
120+
## Best Practices & Common Patterns
121+
122+
1. **Select Correct `pooling_type`**:
123+
- Standard embeddings: `LLAMA_POOLING_TYPE_UNSPECIFIED (-1)`.
124+
- Reranker models: `LLAMA_POOLING_TYPE_RANK (4)`.
125+
- Token-level embeddings: `LLAMA_POOLING_TYPE_NONE (0)`.
126+
127+
2. **Batch Optimization for Large Datasets**:
128+
- Adjust `n_batch` and `n_ubatch` to balance performance and memory.
129+
- Streaming processing avoids OOM for large datasets.
130+
131+
3. **Normalization Selection**:
132+
- Vector databases typically prefer L2 normalization (Euclidean), but other norms may be needed in specific scenarios.
133+
134+
4. **Reranker Models**:
135+
- Ensure `pooling_type` is set to `LLAMA_POOLING_TYPE_RANK`.
136+
- Note that output is scalar scores, not vectors.
137+
138+
5. **Performance Tuning**:
139+
- For GPU acceleration, set `n_gpu_layers` to -1 (recommended).
140+
- Use `verbose=True` for debugging configuration.
141+
142+
## Example Code
143+
144+
### 1. Text Embeddings (Vector Search)
145+
146+
To generate embeddings, use the `LlamaEmbedding` class. It automatically configures the model for vector generation.
147+
148+
```python
149+
from llama_cpp.llama_embedding import LlamaEmbedding, LLAMA_POOLING_TYPE_NONE
150+
151+
# Initialize the model (automatically sets embeddings=True)
152+
llm = LlamaEmbedding(model_path="path/to/bge-m3.gguf", n_gpu_layers=-1, pooling_type=LLAMA_POOLING_TYPE_NONE)
153+
154+
# 1. Simple usage (OpenAI-compatible format)
155+
response = llm.create_embedding("Hello, world!")
156+
print(response['data'][0]['embedding'])
157+
158+
# 2. Batch processing (High Performance)
159+
# You can pass a large list of strings; the streaming batcher handles memory automatically.
160+
documents = ["Hello, world!", "Goodbye, world!", "Llama is cute."] * 100
161+
embeddings = llm.embed(documents) # Returns a list of lists (vectors)
162+
163+
print(f"Generated {len(embeddings)} vectors.")
164+
```
165+
166+
**Advanced Output Formats:**
167+
You can request raw arrays or cosine similarity matrices directly:
168+
169+
```python
170+
from llama_cpp.llama_embedding import LlamaEmbedding, LLAMA_POOLING_TYPE_NONE
171+
172+
# Initialize the model (automatically sets embeddings=True)
173+
llm = LlamaEmbedding(model_path="path/to/bge-m3.gguf", n_gpu_layers=-1, pooling_type=LLAMA_POOLING_TYPE_NONE)
174+
175+
# Returns raw List[float] instead of a dictionary wrapper
176+
vector = llm.create_embedding("Text", output_format="array")
177+
178+
# Returns a similarity matrix (A @ A.T) in the response
179+
# Note: Requires numpy installed
180+
response = llm.create_embedding(
181+
["apple", "fruit", "car"],
182+
output_format="json+"
183+
)
184+
print(response["cosineSimilarity"])
185+
```
186+
187+
### 2. Reranking (Cross-Encoder Scoring)
188+
189+
Reranking models (like `bge-reranker`) take a **Query** and a list of **Documents** as input and output a relevance score (scalar) for each document.
190+
191+
> **Important:** You must explicitly set `pooling_type` to `LLAMA_POOLING_TYPE_RANK` (4) when initializing the model.
192+
193+
```python
194+
import llama_cpp
195+
from llama_cpp.llama_embedding import LlamaEmbedding
196+
197+
# Initialize a Reranking model
198+
ranker = LlamaEmbedding(
199+
model_path="path/to/qwen3-reranker-0.6b-q8_0.gguf",
200+
pooling_type=llama_cpp.LLAMA_POOLING_TYPE_RANK, # Crucial for Rerankers!
201+
n_gpu_layers=-1,
202+
n_ctx=0
203+
)
204+
205+
query = "What causes Rain?"
206+
docs = [
207+
"Clouds are made of water droplets...", # Relevant
208+
"To bake a cake you need flour...", # Irrelevant
209+
"Rain is liquid water in the form of droplets..." # Highly Relevant
210+
]
211+
212+
# Calculate relevance scores
213+
# Logic: Constructs inputs like "[BOS] query [SEP] doc [EOS]" automatically
214+
scores = ranker.rank(query, docs)
215+
216+
# Result: List of floats (higher means more relevant)
217+
print(scores)
218+
# e.g., [0.0011407170677557588, 5.614783731289208e-05, 0.7173627614974976] -> The 3rd doc is the best match
219+
```
220+
221+
### 3. Normalization
222+
223+
The `embed` method supports various mathematical normalization strategies via the `normalize` parameter.
224+
225+
| Normalization modes | $Integer$ | Description | Formula |
226+
|---------------------|-----------|---------------------|---------|
227+
| NORM_MODE_NONE | $-1$ | none |
228+
| NORM_MODE_MAX_INT16 | $0$ | max absolute int16 | $\Large{{32760 * x_i} \over\max \lvert x_i\rvert}$
229+
| NORM_MODE_TAXICAB | $1$ | taxicab | $\Large{x_i \over\sum \lvert x_i\rvert}$
230+
| NORM_MODE_EUCLIDEAN | $2$ | euclidean (default) | $\Large{x_i \over\sqrt{\sum x_i^2}}$
231+
| NORM_MODE_PNORM | $>2$ | p-norm | $\Large{x_i \over\sqrt[p]{\sum \lvert x_i\rvert^p}}$
232+
233+
This is useful for optimizing storage or preparing vectors for cosine similarity search (which requires L2 normalization).
234+
235+
```python
236+
from llama_cpp.llama_embedding import (
237+
LLAMA_POOLING_TYPE_NONE,
238+
NORM_MODE_MAX_INT16,
239+
NORM_MODE_TAXICAB,
240+
NORM_MODE_EUCLIDEAN
241+
)
242+
243+
# Initialize the model (automatically sets embeddings=True)
244+
llm = LlamaEmbedding(model_path="path/to/bge-m3.gguf", n_gpu_layers=-1, pooling_type=LLAMA_POOLING_TYPE_NONE)
245+
246+
# Taxicab (L1)
247+
vec_l1 = llm.embed("text", normalize=NORM_MODE_TAXICAB)
248+
249+
# Default is Euclidean (L2) - Standard for vector databases
250+
vec_l2 = llm.embed("text", normalize=NORM_MODE_EUCLIDEAN)
251+
252+
# Max Absolute Int16 - Useful for quantization/compression
253+
vec_int16 = llm.embed("text", normalize=NORM_MODE_MAX_INT16)
254+
255+
# Raw Output (No Normalization) - Get the raw floating point values from the model
256+
embeddings_raw = llm.embed(["search query", "document text"], normalize=NORM_MODE_NONE)
257+
```
258+
259+
## Notes
260+
261+
- This class is in development; some features may be unstable, especially reranking model support.
262+
- Performance issues can be addressed by adjusting `n_batch`, `n_ubatch`, and `n_gpu_layers`.
263+
- For custom models, manual `pooling_type` configuration may be required to match model behavior.

0 commit comments

Comments
 (0)