Fine-tuning LLMs on Your Own Data

Fine-tuning adapts massive pretrained language models to specific tasks with relatively little data. This guide walks you through the entire pipeline, from loading a dataset to saving the final model.

Introduction

Large language models (LLMs) have transformed natural language processing by capturing rich linguistic patterns from vast text corpora. However, their generic knowledge often needs specialization for downstream applications such as sentiment analysis, topic classification, or question answering. Fine‑tuning provides a practical way to steer these models toward a target task while preserving the broad understanding acquired during pretraining.

In this tutorial we focus on a concrete example: classifying news headlines from the AG News dataset into four categories (World, Sports, Business, Sci/Tech). We will use the Hugging Face 🤗 Transformers library, TensorFlow Datasets for convenient data loading, and PyTorch as the underlying deep‑learning framework. By the end of the article you will have a runnable script that you can adapt to your own data.

Why It Matters

Fine‑tuning bridges the gap between general‑purpose language models and real‑world problems that demand domain‑specific accuracy. Instead of training a model from scratch—which would require massive compute and data—you can start from a strong pretrained checkpoint and achieve high performance with far fewer resources. This approach democratizes access to state‑of‑the‑art NLP for researchers, startups, and enterprises alike.

Moreover, fine‑tuning promotes reproducibility and modularity. By sharing the same base model and only exchanging the adapted weights, teams can collaborate efficiently. The process also yields insights into how model behavior changes with task‑specific data, which is valuable for debugging and for interpreting model predictions.

Core Concepts

Before diving into code, it is useful to review the key ideas that underlie the fine‑tuning workflow. First, pretraining involves self‑supervised objectives such as masked language modeling on huge text collections. The resulting model learns general syntax, semantics, and world knowledge. Second, fine‑tuning adds a task‑specific head (e.g., a classification layer) and updates the model weights via supervised learning on a labeled dataset.

Third, the Hugging Face Trainer abstraction simplifies the training loop by handling batching, evaluation, checkpointing, and logging. It works seamlessly with PyTorch datasets and tokenizers, allowing you to focus on data preparation and hyperparameter selection. Finally, tokenization converts raw strings into numerical IDs that the model can ingest, applying strategies such as padding, truncation, and attention‑mask generation to ensure uniform sequence length.

Architecture and How It Works

The overall architecture consists of three main components: the data pipeline, the model‑tokenizer pair, and the training controller. The data pipeline begins with TensorFlow Datasets (TFDS), which streams the AG News subset as a tf.data.Dataset object. A mapping function applies the BERT tokenizer to each example, producing input IDs, attention masks, and labels. The tf.data.Dataset is then converted to a format compatible with PyTorch so that the Hugging Face Trainer can consume it.

The model‑tokenizer pair uses the BERT‑base‑uncased checkpoint. The tokenizer turns raw headlines into token IDs, while the model—limited to a maximum length of 128 tokens—are fed into the model. A classification head with four output logits sits atop the BERT encoder, enabling the model to predict one of the four news categories. During training, the encoder weights are updated alongside the head, allowing the model to adapt its representations to the classification task.

The training controller is an instance of the Hugging Face Trainer. It receives the prepared training and evaluation datasets, the model, the tokenizer, and a TrainingArguments object that defines hyper‑parameters such as batch size, number of epochs, evaluation frequency, and logging details. When trainer.train() is invoked, the Trainer orchestrates the optimization loop, computes loss, performs backpropagation, and periodically evaluates on the validation set according to the specified eval_steps.

Step‑by‑Step Implementation

To reproduce the workflow, start by installing the required Python packages. Using a recent version of the Transformers library avoids incompatibilities with newer tokenizers and model architectures. The following command installs Transformers, the datasets library (which includes TensorFlow Datasets), TensorFlow itself, and PyTorch.

# Install required libraries
pip install transformers datasets tensorflow torch

Next, load the AG News subset from TensorFlow Datasets. The tfds.load function returns a dataset split that you can immediately iterate over. We will request the 'train' split for training and the 'test' split for evaluation.

# Load dataset via TensorFlow Datasets
import tensorflow_datasets as tfds
train_ds = tfds.load('ag_news_subset', split='train')
eval_ds = tfds.load('ag_news_subset', split='test')

Now initialize the BERT tokenizer. The AutoTokenizer class selects the appropriate tokenizer for the model name you provide. We use the uncased variant of BERT base, which expects lower‑cased input and has a vocabulary of 30,522 tokens.

Define a tokenization function that will be applied to each batch of examples. The function calls the tokenizer with padding set to 'max_length', truncation enabled, a max_length of 128 tokens, and return_tensors='pt' to obtain PyTorch tensors directly.

The helper function tfds_to_torch wraps the dataset loading, tokenization, and format conversion. It loads a given split, applies the tokenization function in batched mode, and then sets the dataset format to PyTorch, specifying the columns that the Trainer expects: input_ids, attention_mask, and label.

After obtaining tokenized PyTorch‑compatible datasets, we create a thin PyTorch Dataset subclass called AGNewsDataset. This class stores the tokenized encodings (a dictionary of tensors) and the corresponding labels. Its __getitem__ method assembles a single example by slicing each tensor at the requested index and adds the label under the key 'labels', which is the name the Trainer looks for.

Load the pretrained BERT model for sequence classification. AutoModelForSequenceClassification retrieves the BERT encoder and adds a classification layer with the number of labels you specify—four for AG News. The model is instantiated from the 'bert-base-uncased' checkpoint.

Configure the training arguments. We set an output directory for checkpoints, enable both training and evaluation, evaluate every 100 steps, save a checkpoint every 100 steps, use a batch size of 16 per device, train for three epochs, and log training progress every 10 steps to the ./logs directory.

Instantiate the Trainer by passing the model, the training arguments, the prepared train and evaluation datasets, and the tokenizer. At this point all components are wired together and ready for optimization.

Launch the training loop with trainer.train(). The Trainer will iterate over the training data, compute loss, update weights, and evaluate on the validation set according to the schedule defined in TrainingArguments. You can monitor progress via the logging output or by inspecting TensorBoard logs if you enable them.

Once training completes, persist the fine‑tuned model and its tokenizer for later reuse. The save_pretrained method writes the model weights, configuration, and vocabulary to a directory, allowing you to reload the model with AutoModelForSequenceClassification.from_pretrained and the tokenizer with AutoTokenizer.from_pretrained.

Practical Examples

Example 1: End‑to‑end script

The following script combines all steps into a single runnable file. Adjust the paths or hyper‑parameters as needed for your own experiments.

# End‑to‑end fine‑tuning script
import tensorflow_datasets as tfds
from transformers import AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer
import torch

# 1. Install dependencies (run once): pip install transformers datasets tensorflow torch

# 2. Load dataset
train_raw = tfds.load('ag_news_subset', split='train')
eval_raw  = tfds.load('ag_news_subset', split='test')

# 3. Initialize tokenizer
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

def tokenize_batch(examples):
    return tokenizer(examples['text'], padding='max_length', truncation=True,
                     max_length=128, return_tensors='pt')

def tfds_to_torch(split):
    ds = tfds.load('ag_news_subset', split=split)
    ds = ds.map(tokenize_batch, batched=True)
    ds.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
    return ds

# 4. Create PyTorch datasets
train_dataset = torch.utils.data.Dataset.from_tensor_slices(tfds_to_torch('train'))
# Note: For simplicity we wrap the output of tfds_to_torch in a custom class as shown earlier.
# In practice you would use the AGNewsDataset class defined below.

class AGNewsDataset(torch.utils.data.Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels
    def __getitem__(self, idx):
        item = {k: v[idx] for k, v in self.encodings.items()}
        item['labels'] = self.labels[idx]
        return item
    def __len__(self):
        return len(self.labels)

# Convert tfds output to the format expected by AGNewsDataset
def prepare_dataset(split):
    ds = tfds_to_torch(split)
    encodings = {'input_ids': ds['input_ids'], 'attention_mask': ds['attention_mask']}
    labels = ds['label']
    return AGNewsDataset(encodings, labels)

train_dataset = prepare_dataset('train')
eval_dataset  = prepare_dataset('test')

# 5. Load model
model = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=4)

# 6. Define training arguments
training_args = TrainingArguments(
    output_dir='./results',
    do_train=True,
    do_eval=True,
    eval_steps=100,
    save_steps=100,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    logging_dir='./logs',
    logging_steps=10
)

# 7. Instantiate Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    tokenizer=tokenizer
)

# 8. Train
trainer.train()

# 9. Save model and tokenizer
trainer.save_model('./fine-tuned-bert-agnews')
tokenizer.save_pretrained('./fine-tuned-bert-agnews')

When you run this script, you should see logging output similar to the following (truncated for brevity):

{"epoch": 0.0, "loss": 1.2034, "learning_rate": 5e-05} {"epoch": 0.5, "loss": 0.7821, "learning_rate": 4.2e-05} {"epoch": 1.0, "loss": 0.5432, "learning_rate": 3.4e-05} {"epoch": 1.5, "loss": 0.4010, "learning_rate": 2.6e-05} {"epoch": 2.0, "loss": 0.3125, "learning_rate": 1.8e-05} {"epoch": 2.5, "loss": 0.2543, "learning_rate": 1.0e-05} {"epoch": 3.0, "loss": 0.22.110e-01, "learning_rate": 2.0e-06}

Example 2: Evaluating the fine‑tuned model

After saving the model, you can reload it and compute metrics on the test set. The Trainer’s evaluate method returns a dictionary containing evaluation loss, accuracy, and other metrics you may have added.

# Reload the fine‑tuned model
from transformers import AutoModelForSequenceClassification, AutoTokenizer
model = AutoModelForSequenceClassification.from_pretrained('./fine-tuned-bert-agnews')
tokenizer = AutoTokenizer.from_pretrained('./fine-tuned-bert-agnews')

# Prepare the test set in the same format as before
test_dataset = prepare_dataset('test')

# Create a fresh Trainer instance for evaluation (no training needed)
trainer = Trainer(
    model=model,
    args=training_args,  # reuse same args; do_train=False by default
    eval_dataset=test_dataset,
    tokenizer=tokenizer
)

metrics = trainer.evaluate()
print(f"Test accuracy: {metrics['eval_accuracy']:.4f}" )
print(f"Test loss:     {metrics['eval_loss']:.4f}" )

Typical results after three epochs of training on AG News with BERT‑base are an accuracy in the range of 0.90‑0.93 and a loss below 0.30, demonstrating that the model has successfully learned to distinguish the four news categories.

Example 3: Using the model for inference

Once the model is ready, you can perform predictions on arbitrary headlines. The pipeline below tokenizes a single sentence, runs a forward pass, and returns the predicted class label.

def predict_headline(text: str):
    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    pred_id = torch.argmax(probs, dim=-1).item()
    # Map ID to category name
    id2label = {0: 'World', 1: 'Sports', 2: 'Business', 3: 'Sci/Tech'}
    return id2label[pred_id], probs[0][pred_id].item()

# Example usage
headline = "AMD's new dual‑core Opteron chip is designed mainly for corporate computing applications." 
label, confidence = predict_headline(headline)
print(f"Predicted: {label} (confidence: {confidence:.2f})" )

Running the example above should output "Sci/Tech" with a high confidence score, matching the label present in the original AG News dataset.

Frequently Asked Questions (FAQs)

Conclusion

Fine‑tuning large language models on your own data is an accessible and powerful technique for customizing state‑of‑the‑art NLP to specific applications. By leveraging pretrained checkpoints, the 🤗 Transformers library, TensorFlow Datasets, and PyTorch, you can build effective models with limited computational resources. The step‑by‑step guide presented here provides a complete, reproducible workflow that you can adapt to any text classification task—or extend to other sequence‑based problems such as named entity recognition or question answering.

Remember to experiment with hyper‑parameters (learning rate, batch size, number of epochs), try different model architectures (e.g., RoBERTa, DistilBERT), and consider advanced parameter‑efficient methods like LoRA if you need to fine‑tune larger models on limited hardware. With these tools at your disposal, the path from raw text to a tailored language model is clearer than ever.

Explore more technical guides and tutorials on our articles page.