Implementation:Neuml Txtai Pooling Base
| Knowledge Sources | |
|---|---|
| Domains | Embeddings, Transformer Models, Vector Pooling |
| Last Updated | 2026-02-10 01:00 GMT |
Overview
Concrete tool for building pooled embedding vectors from transformer model outputs provided by txtai.
Description
The Pooling class is a PyTorch nn.Module that serves as the base pooling model for generating vector embeddings from transformer model outputs. It handles the full encoding pipeline: loading a transformer model and tokenizer from Hugging Face Hub or local paths, managing device placement (CPU/GPU), batching documents for efficient processing, tokenizing input text, running the forward pass through the transformer, and collecting outputs. The class implements performance optimizations including sorting documents by length for efficient batching (matching sentence-transformers behavior) and restoring original order after encoding. It supports configurable maximum sequence length, instruction prompt prepending (loaded from sentence-transformers configuration files), and pre/post encoding hooks for subclass customization. The base forward method returns the raw first output tensor from the transformer model.
Usage
Use Pooling as the base class for all pooling strategies in txtai, or directly when the raw hidden state output from a transformer is desired without further pooling operations. It is typically instantiated through the PoolingFactory rather than directly. The encode method is the primary entry point, accepting a list of documents and returning a NumPy array of embeddings.
Code Reference
Source Location
- Repository: Neuml_Txtai
- File:
src/python/txtai/models/pooling/base.py
Signature
class Pooling(nn.Module):
def __init__(self, path, device, tokenizer=None, maxlength=None, loadprompts=None, modelargs=None)
def encode(self, documents, batch=32, category=None)
def chunk(self, texts, size)
def forward(self, **inputs)
def preencode(self, documents, category)
def postencode(self, results, category)
def load(self, path, name)
def loadprompts(self, path)
Import
from txtai.models.pooling.base import Pooling
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| path | str or bytes | Yes | Path to the model. Accepts a Hugging Face model hub id (e.g., "sentence-transformers/all-MiniLM-L6-v2") or a local filesystem path. |
| device | int or str | Yes | Tensor device id for model placement (e.g., 0 for GPU 0, "cpu" for CPU).
|
| tokenizer | str | No | Optional path to a custom tokenizer. Defaults to using the same path as the model. |
| maxlength | int | No | Maximum sequence length for tokenization. Defaults to the tokenizer's model_max_length if bounded, otherwise None.
|
| loadprompts | bool | No | Whether to load instruction prompts from config_sentence_transformers.json.
|
| modelargs | dict | No | Additional keyword arguments passed to the model loader. |
| documents | list of str | Yes (for encode) | List of text documents to encode into embeddings. |
| batch | int | No | Batch size for processing documents. Default is 32. |
| category | str | No | Embeddings category, either "query" or "data". Used for prompt selection and subclass-specific behavior. |
Outputs
| Name | Type | Description |
|---|---|---|
| encode() | numpy.ndarray | 2D array of shape (num_documents, embedding_dim) containing pooled embeddings, in the original document order. |
| forward() | torch.Tensor | Raw transformer output tensor (first element of model output tuple). |
| preencode() | list | Documents with optional prompt prepended based on category. |
| postencode() | list | Results after post-encoding transformation (identity in base class). |
| load() | dict or None | Parsed JSON configuration from a Hugging Face Hub config file, or None if not found. |
| loadprompts() | dict or None | Dictionary of instruction prompts keyed by category ("query", "data", "document"), or None. |
Usage Examples
from txtai.models.pooling.base import Pooling
# Create a pooling model from a Hugging Face model
model = Pooling(
path="sentence-transformers/all-MiniLM-L6-v2",
device="cpu",
maxlength=512,
loadprompts=True
)
# Encode documents
documents = [
"Machine learning is a branch of artificial intelligence.",
"Natural language processing enables text understanding.",
"Deep learning uses multi-layered neural networks."
]
embeddings = model.encode(documents, batch=16, category="data")
# embeddings.shape: (3, 384)
# Encode a query with prompt
query_embedding = model.encode(["What is machine learning?"], category="query")