Implementation:Neuml Txtai HFTrainer Model

Knowledge Sources	txtai txtai Documentation
Domains	Training, NLP
Last Updated	2026-02-09 00:00 GMT

Overview

Concrete tool for loading and configuring pretrained transformer models for fine-tuning provided by the txtai library. This method wraps HuggingFace transformers AutoModel classes, adding task-based dispatch, optional BitsAndBytes quantization, and LoRA adapter injection.

Description

HFTrainer.model() is the central method responsible for instantiating the correct pretrained model architecture based on the training task. It performs the following steps:

1. If a label count is provided (classification tasks), it updates the model configuration with num_labels.
2. It formats the quantization configuration -- accepting True (for sensible 4-bit NF4 defaults), a custom dict, or a BitsAndBytesConfig object. Quantization is automatically disabled when no GPU is available.
3. It dispatches to the appropriate AutoModel class based on the task string.
4. If the base argument is already a (model, tokenizer) tuple (from a prior training run), the existing model is returned directly.

The task-to-model mapping is:

• "language-generation" -- AutoModelForCausalLM
• "language-modeling" -- AutoModelForMaskedLM
• "question-answering" -- AutoModelForQuestionAnswering
• "sequence-sequence" -- AutoModelForSeq2SeqLM
• "token-detection" -- TokenDetection (paired MaskedLM + PreTraining)
• Default ("text-classification") -- AutoModelForSequenceClassification

Usage

This method is called internally by HFTrainer.__call__() during the training pipeline. It can also be used directly when only model instantiation (without full training) is needed, for example to inspect a model's architecture or to prepare a model for manual training loops.

Code Reference

Source Location

• Repository: txtai
• File: src/python/txtai/pipeline/train/hftrainer.py (Lines 230-303)

Signature

def model(self, task, base, config, labels, tokenizer, quantize):
    """
    Loads the base model to train.

    Args:
        task: model task or category, determines the model type
        base: base model path or (model, tokenizer) tuple
        config: AutoConfig instance
        labels: number of labels (classification tasks)
        tokenizer: model tokenizer
        quantize: quantization config (bool, dict, or BitsAndBytesConfig)

    Returns:
        Pretrained HF model instance
    """

Import

from txtai.pipeline import HFTrainer

I/O Contract

Inputs

Name	Type	Required	Description
task	str	Yes	Task identifier that selects the model architecture. Supported values: "language-generation", "language-modeling", "question-answering", "sequence-sequence", "token-detection", "text-classification" (default).
base	str or tuple	Yes	Path to pretrained model (HuggingFace hub ID or local path), or a (model, tokenizer) tuple from a prior training run.
config	AutoConfig	Yes	Model configuration loaded from the base model. Updated with num_labels when labels are provided.
labels	int or None	No	Number of output labels for classification tasks. When None, the config is used as-is.
tokenizer	AutoTokenizer	Yes	Tokenizer instance, used only for the "token-detection" task.
quantize	bool / dict / BitsAndBytesConfig / None	No	Quantization configuration. True applies 4-bit NF4 defaults (load_in_4bit=True, double quantization, bfloat16 compute dtype). A dict is converted to BitsAndBytesConfig. Ignored when no GPU is available.

Outputs

Name	Type	Description
model	PreTrainedModel	A pretrained HuggingFace model instance, optionally quantized. The specific class depends on the task: AutoModelForCausalLM, AutoModelForMaskedLM, AutoModelForQuestionAnswering, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification, or TokenDetection.

Usage Examples

Basic Example: Load a Classification Model

from transformers import AutoConfig, AutoTokenizer
from txtai.pipeline import HFTrainer

trainer = HFTrainer()

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

model = trainer.model(
    task="text-classification",
    base="bert-base-uncased",
    config=config,
    labels=2,
    tokenizer=tokenizer,
    quantize=None
)

Load a Quantized Causal LM

from transformers import AutoConfig, AutoTokenizer
from txtai.pipeline import HFTrainer

trainer = HFTrainer()

config = AutoConfig.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

# True enables 4-bit NF4 defaults
model = trainer.model(
    task="language-generation",
    base="meta-llama/Llama-2-7b-hf",
    config=config,
    labels=None,
    tokenizer=tokenizer,
    quantize=True
)

Custom Quantization Dict

from txtai.pipeline import HFTrainer

trainer = HFTrainer()

quantize_config = {
    "load_in_4bit": True,
    "bnb_4bit_use_double_quant": True,
    "bnb_4bit_quant_type": "nf4",
    "bnb_4bit_compute_dtype": "bfloat16",
}

model = trainer.model(
    task="language-generation",
    base="mistralai/Mistral-7B-v0.1",
    config=config,
    labels=None,
    tokenizer=tokenizer,
    quantize=quantize_config
)

Related Pages

Implements Principle

• Principle:Neuml_Txtai_Base_Model_Configuration

Requires Environment

• Environment:Neuml_Txtai_GPU_Accelerator_Environment