Implementation:Neuml Txtai Data Init

Overview

This page documents the Data base class and the Labels subclass, which provide the foundation for tokenizing and preparing datasets for transformer model training in txtai. The Data class defines the shared tokenization interface, while Labels implements text-classification-specific tokenization logic.

API

Data.__init__

def __init__(self, tokenizer, columns, maxlength)

Creates a new base instance for tokenizing data. This constructor stores the tokenizer, column definitions, and maximum sequence length that will be used throughout the tokenization pipeline.

Parameters:

Name	Type	Description
tokenizer	PreTrainedTokenizer	Model tokenizer used to convert text into token IDs
columns	tuple or None	Column names defining the structure of input data (task-specific)
maxlength	int	Maximum sequence length for truncation and padding

Returns: None

Example:

from txtai.data import Data
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
data = Data(tokenizer, ("text", None, "label"), maxlength=512)

Data.__call__

def __call__(self, train, validation, workers)

Tokenizes training and validation data and returns processed datasets. This is the main entry point for the data tokenization pipeline. It delegates to prepare() for the actual tokenization work.

Parameters:

Name	Type	Description
train	Dataset, DataFrame, or iterable of dicts	Training data to tokenize
validation	Dataset, DataFrame, iterable of dicts, or None	Validation data to tokenize; returns None if not provided
workers	int or None	Number of concurrent tokenizers when processing Hugging Face datasets; only main process used when set to None

Returns: (train, validation) -- A tuple of tokenized training and validation datasets. Validation is None if no validation data was provided.

Example:

from txtai.data import Labels
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
process = Labels(tokenizer, ("text", None, "label"), maxlength=512)

train_data = [{"text": "great movie", "label": 1}, {"text": "terrible film", "label": 0}]
train_tokens, val_tokens = process(train_data, None, None)

Data.prepare

def prepare(self, data, fn, workers)

Prepares and tokenizes data for training. Handles three input formats: Hugging Face Datasets (via .map()), Polars/pandas DataFrames (via .columns), and iterable dicts. For non-Dataset inputs, data is re-oriented from row format to column format for efficient batch tokenization, then wrapped in a Tokens dataset.

Parameters:

Name	Type	Description
data	Dataset, DataFrame, or iterable of dicts	Input data to tokenize
fn	callable	Tokenize processing function to apply (typically self.process)
workers	int or None	Number of concurrent workers for multiprocessing with Hugging Face datasets

Returns: Tokenized data as a Hugging Face Dataset or Tokens instance.

Data.labels

def labels(self, data)

Extracts the number of unique labels from the training data. Used to configure the classification head of the model.

Parameters:

Name	Type	Description
data	Dataset, DataFrame, or iterable of dicts	Input data containing a label column

Returns: int -- Number of unique labels for classification tasks, or 1 for regression tasks (when at least one label is a non-integer float). If labels are arrays (multi-label), returns the array length.

Behavior:

• If a label value has a __len__ attribute (e.g., a list), its length is returned directly (multi-label classification).
• If all label values are integers, the count of unique labels is returned (classification).
• If any label has a fractional component, 1 is returned (regression).

Data.process

def process(self, data)

Base tokenization method. In the Data base class, this is a no-op that returns data unchanged. Subclasses override this method with task-specific tokenization logic.

Parameters:

Name	Type	Description
data	dict	Input data batch as a dictionary of column lists

Returns: Tokenized data (identity in base class).

Data.length

def length(self, value)

Utility method that returns the length of a value if it supports len(), otherwise returns None. Used internally by labels() to detect multi-label classification.

Parameters:

Name	Type	Description
value	any	Value to check for length

Returns: int or None

Labels API

Labels.__init__

def __init__(self, tokenizer, columns, maxlength)

Creates a new instance for tokenizing text-classification training data. Extends Data.__init__ by standardizing the column definitions for classification tasks.

Parameters:

Name	Type	Description
tokenizer	PreTrainedTokenizer	Model tokenizer
columns	tuple or None	Tuple of columns for text/label mapping. Defaults to ("text", None, "label") when None.
maxlength	int	Maximum sequence length

Column standardization rules:

• If columns is None or falsy: defaults to ("text", None, "label").
• If columns has fewer than 3 elements: maps to (columns[0], None, columns[-1]), treating the first element as the text column and the last as the label column.
• If columns has 3 elements: used as-is as (text1, text2, label) for sentence-pair tasks.

Example:

from txtai.data import Labels
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

# Single text column with default column names
labels = Labels(tokenizer, None, maxlength=512)
# columns becomes ("text", None, "label")

# Custom column names
labels = Labels(tokenizer, ("sentence", "label"), maxlength=512)
# columns becomes ("sentence", None, "label")

# Sentence pair task
labels = Labels(tokenizer, ("sentence1", "sentence2", "label"), maxlength=512)
# columns remains ("sentence1", "sentence2", "label")

Labels.process

def process(self, data)

Tokenizes a batch of text-classification input data. Handles both single-text and sentence-pair inputs depending on whether the second column is set.

Parameters:

Name	Type	Description
data	dict	Batch of input data as column-oriented dictionary

Returns: dict -- Tokenized data with input_ids, attention_mask, token_type_ids, and the label column.

Behavior:

1. Unpacks columns as (text1, text2, label).
2. If text2 is set, passes a sentence pair to the tokenizer; otherwise passes a single text.
3. Calls self.tokenizer(*text, max_length=self.maxlength, padding=True, truncation=True).
4. Attaches the label data to the tokenized output.

Source

• src/python/txtai/data/base.py (lines 8-138)
• src/python/txtai/data/labels.py (lines 8-42)

Import

from txtai.data import Data, Labels

See Also

• Neuml_Txtai_Training_Data_Preparation -- Principle: tokenization theory and task-specific data formatting
• Neuml_Txtai_HFTrainer_Call -- The training pipeline that consumes prepared data