neuml--txtai/docs/pipeline/train/trainer.md

HFTrainer

Trains a new Hugging Face Transformer model using the Trainer framework.

Example

The following shows a simple example using this pipeline.

import pandas as pd

from datasets import load_dataset

from txtai.pipeline import HFTrainer

trainer = HFTrainer()

# Pandas DataFrame
df = pd.read_csv("training.csv")
model, tokenizer = trainer("bert-base-uncased", df)

# Hugging Face dataset
ds = load_dataset("glue", "sst2")
model, tokenizer = trainer("bert-base-uncased", ds["train"], columns=("sentence", "label"))

# List of dicts
dt = [{"text": "sentence 1", "label": 0}, {"text": "sentence 2", "label": 1}]]
model, tokenizer = trainer("bert-base-uncased", dt)

# Support additional TrainingArguments
model, tokenizer = trainer("bert-base-uncased", dt, 
                            learning_rate=3e-5, num_train_epochs=5)

All TrainingArguments are supported as function arguments to the trainer call.

See the links below for more detailed examples.

Notebook	Description
Train a text labeler	Build text sequence classification models
Train without labels	Use zero-shot classifiers to train new models
Train a QA model	Build and fine-tune question-answering models
Train a language model from scratch	Build new language models

Training tasks

The HFTrainer pipeline builds and/or fine-tunes models for following training tasks.

Task	Description
language-generation	Causal language model for text generation (e.g. GPT)
language-modeling	Masked language model for general tasks (e.g. BERT)
question-answering	Extractive question-answering model, typically with the SQuAD dataset
sequence-sequence	Sequence-Sequence model (e.g. T5)
text-classification	Classify text with a set of labels
token-detection	ELECTRA-style pre-training with replaced token detection

PEFT

Parameter-Efficient Fine-Tuning (PEFT) is supported through Hugging Face's PEFT library. Quantization is provided through bitsandbytes. See the examples below.

from txtai.pipeline import HFTrainer

trainer = HFTrainer()
trainer(..., quantize=True, lora=True)

When these parameters are set to True, they use default configuration. This can also be customized.

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

lora = {
    "r": 16,
    "lora_alpha": 8,
    "target_modules": "all-linear",
    "lora_dropout": 0.05,
    "bias": "none"
}

trainer(..., quantize=quantize, lora=lora)

The parameters also accept transformers.BitsAndBytesConfig and peft.LoraConfig instances.

See the following PEFT documentation links for more information.

• Quantization
• LoRA

Merge

An important parameter for language-generation and language-modeling tasks is merge or the packing of data into chunks.

It supports the following options.

• concat (default) - text is split into chunks up to maxlength, data can be split across multiple chunks
• pack - text is split into chunks up to maxlength, data guaranteed to be in same chunk, chunks can be smaller than maxlength
• None - disables merging

Merging helps reduce training time as data can be processed efficiently without padding. concat maximizes this as it guarantees each chunk will be up to maxlength size. pack is a middle ground where data is combined but records are preserved.

For general language modeling tasks like masked language modeling, concat is the best choice. For instruction/prompt fine-tuning, pack or None are the better choices as it guarantees complex logic is not split across chunks.

Methods

Python documentation for the pipeline.

::: txtai.pipeline.HFTrainer.call