Train a text labeler

The Hugging Face Model Hub has a wide range of models that can handle many tasks. While these models perform well, the best performance often is found when fine-tuning a model with task-specific data.

Hugging Face provides a number of full-featured examples available to assist with training task-specific models. When building models from the command line, these scripts are a great way to get started.

txtai provides a training pipeline that can be used to train new models programatically using the Transformers Trainer framework. The training pipeline supports the following:

• Building transient models without requiring an output directory
• Load training data from Hugging Face datasets, pandas DataFrames and list of dicts
• Text sequence classification tasks (single/multi label classification and regression) including all GLUE tasks
• All training arguments

This article shows examples of how to use txtai to train/fine-tune new models.

Install dependencies

Install txtai and all dependencies.

pip install txtai[pipeline-train] datasets pandas

Train a model

Let's get right to it! The following example fine-tunes a tiny Bert model with the sst2 dataset.

The trainer pipeline is basically a one-liner that fine-tunes any text classification/regression model available (locally and/or from the HF Hub).

from datasets import load_dataset

from txtai.pipeline import HFTrainer

trainer = HFTrainer()

# Hugging Face dataset
ds = load_dataset("glue", "sst2")
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", ds["train"], columns=("sentence", "label"))

The default trainer pipeline functionality will not store any logs, checkpoints or models to disk. The trainer can take any of the standard TrainingArguments to enable persistent models.

The next section creates a Labels pipeline using the newly built model and runs the model against the sst2 validation set.

from txtai.pipeline import Labels

labels = Labels((model, tokenizer), dynamic=False)

# Determine accuracy on validation set
results = [row["label"] == labels(row["sentence"])[0][0] for row in ds["validation"]]
sum(results) / len(ds["validation"])

0.8268348623853211

82.68% accuracy - not bad for a tiny Bert model.

Train a model with Lists

As mentioned earlier, the trainer pipeline supports Hugging Face datasets, pandas DataFrames and lists of dicts. The example below trains a model using lists.

data = [{"text": "This is a test sentence", "label": 0}, {"text": "This is not a test", "label": 1}]

model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)

Train a model with DataFrames

The next section builds a new model using data stored in a pandas DataFrame.

import pandas as pd

df = pd.DataFrame(data)

model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", data)

Train a regression model

The previous models were classification tasks. The following model trains a sentence similarity model with a regression output per sentence pair between 0 (dissimilar) and 1 (similar).

ds = load_dataset("glue", "stsb")
model, tokenizer = trainer("google/bert_uncased_L-2_H-128_A-2", ds["train"], columns=("sentence1", "sentence2", "label"))

labels = Labels((model, tokenizer), dynamic=False)
labels([("Sailing to the arctic", "Dogs and cats don't get along"), 
        ("Walking down the road", "Walking down the street")])

[[(0, 0.5648878216743469)], [(0, 0.97544926404953)]]