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:
  • 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
    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)]]

    26

    This website collects cookies to deliver better user experience

    Train a text labeler