26
loading...
Train a text labeler
This article is part of a tutorial series on txtai, an AI-powered semantic search platform.
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:
This article shows examples of how to use txtai to train/fine-tune new models.
Install
txtai and all dependencies.pip install txtaiLet'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.826834862385321182.68% accuracy - not bad for a tiny Bert model.
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)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)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)]]26