A hands-on guide to training your own Small Language Model (SLM) in Python

Language models have revolutionized the field of natural language processing (NLP), powering applications like machine translation, text generation, and semantic search. While large language models like GPT-3 get much of the attention, small language models (SLMs) can be highly effective for specialized domains and applications. In this hands-on tutorial, we’ll walk through how to train your own small language model from scratch using Python.

Contents

• The benefits of thinking small
• Preparing your training data
• Defining the model architecture
• Training the model
• Evaluating and using the trained model
• A full example

The benefits of thinking small

You might be wondering, why bother with small language models when massive models like GPT-3 and PaLM already exist? While these large models are undeniably impressive, they have some significant drawbacks:

• Training them requires enormous computational resources that are out of reach for most organizations
• They can be overkill for narrow, domain-specific applications
• There are potential privacy and security concerns with sending sensitive data to third-party APIs

This is where small language models come in. By training a compact model on a curated dataset, you can create a highly specialized language model that is tailored to your exact use case. And best of all, you can do it with modest hardware and open-source tools.

Preparing your training data

The first step in training a language model is assembling your text corpus. This could be anything from a collection of technical papers to a database of customer support tickets. The key is to curate a dataset that is representative of the type of language you want your model to understand and generate.

Once you have your raw text data, you’ll need to preprocess it and convert it into a format suitable for training. This typically involves:

1. Cleaning the text by removing irrelevant characters, formatting, etc. 2. Tokenizing the text into individual words or subword units 3. Building a vocabulary of all unique tokens 4. Converting the tokens into integer IDs 5. Splitting the data into training and validation sets

Python libraries like NLTK and HuggingFace’s Tokenizers make this preprocessing pipeline straightforward. The end result will be a set of integer-encoded text sequences ready to feed into your model.

import nltk
from nltk.tokenize import word_tokenize

# Download the punkt tokenizer
nltk.download('punkt')

# Sample text corpus
corpus = [
    "This is the first sentence.",
    "Here's another sentence.",
    "And a third sentence for good measure.",
]

# Tokenize the corpus
tokenized_corpus = [word_tokenize(sentence) for sentence in corpus]

# Build the vocabulary
vocab = set(word for sentence in tokenized_corpus for word in sentence)

# Create a mapping from words to integer IDs
word_to_id = {word: i for i, word in enumerate(vocab)}

# Convert the tokenized corpus to integer sequences
int_sequences = [
    [word_to_id[word] for word in sentence] for sentence in tokenized_corpus
]

Defining the model architecture

The next step is to define the architecture of your small language model. A common choice is the transformer encoder, which uses attention mechanisms to capture long-range dependencies in the text. However, for a small model, a simple recurrent neural network (RNN) can work well.

Some key considerations when designing your model:

• Embedding size: This determines the dimensionality of the input token embeddings. A typical range is 128–512.
• Hidden size: The number of units in the RNN hidden state. Again, 128–512 is a common range.
• Number of layers: Stacking multiple RNN layers can increase the model’s capacity, but adds computational cost. 1–3 layers is usually sufficient for a small model.
• Dropout: This regularization technique randomly zeroes out activations during training, which can improve generalization. A dropout rate of 0.2–0.5 is typical.

In addition to the model hyper-parameters, you’ll need to choose an appropriate loss function (such as cross-entropy) and optimization algorithm (like Adam).

Training the model

With the data prepared and model defined, you’re ready to start training. The training loop will typically look something like this:

1- For each epoch:

• Iterate over the training data in batches
• Iterate over the training data in batches
• Feed each batch through the model to get the predictions
• Calculate the loss between the predictions and the true next-token targets
• Back-propagate the gradients and update the model parameters
• Calculate the validation loss on a held-out dataset

2- Stop training when the validation loss stops improving

It’s important to monitor both the training and validation loss during training. If the validation loss starts increasing while the training loss is still decreasing, that’s a sign that the model is starting to overfit and memorize the training data.

To combat overfitting, you can use techniques like early stopping (halting training when the validation loss stops improving) and model checkpointing (saving the model parameters at regular intervals during training).

Python libraries like PyTorch and TensorFlow make building and training custom models straightforward, even if you’re not an expert in machine learning.

import torch
import torch.nn as nn
import torch.optim as optim

# Define the model architecture
class SLM(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super(SLM, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        
    def forward(self, x):
        x = self.embedding(x)
        _, (h, _) = self.lstm(x)
        output = self.fc(h.squeeze(0))
        return output

# Hyper-parameters
vocab_size = len(vocab)
embedding_dim = 128
hidden_dim = 256
learning_rate = 0.001
num_epochs = 10

# Create the model and optimizer
model = SLM(vocab_size, embedding_dim, hidden_dim)
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
criterion = nn.CrossEntropyLoss()

# Training loop
for epoch in range(num_epochs):
    for sequence in int_sequences:
        optimizer.zero_grad()
        input_seq = torch.tensor(sequence[:-1]).unsqueeze(0)
        target_seq = torch.tensor(sequence[1:])
        
        output = model(input_seq)
        loss = criterion(output.view(-1, vocab_size), target_seq)
        
        loss.backward()
        optimizer.step()
    
    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")

Evaluating and using the trained model

Once the model is trained, it’s important to evaluate its performance on a held-out test set. A common metric for language models is perplexity, which measures how “surprised” the model is by the test data. Lower perplexity indicates better performance.

With a trained model in hand, you can use it for a variety of downstream tasks, such as:

• Text generation: Given a prompt, use the model to generate novel text in the style of the training data.
• Text classification: Fine-tune the model on a labeled dataset to create a custom text classifier.
• Semantic search: Use the model’s learned embeddings to find similar documents in a database.

The compact size of small language models makes them ideal for deploying in resource-constrained environments like mobile devices and web browsers.

# Generate text using the trained model
def generate_text(prompt, max_length=20):
    model.eval()
    words = word_tokenize(prompt)
    input_seq = torch.tensor([word_to_id[word] for word in words]).unsqueeze(0)
    
    generated_text = prompt
    for _ in range(max_length):
        output = model(input_seq)
        _, next_word_id = torch.max(output, dim=1)
        next_word = list(vocab)[next_word_id.item()]
        
        generated_text += " " + next_word
        input_seq = torch.cat((input_seq[:, 1:], next_word_id.unsqueeze(0)), dim=1)
    
    return generated_text

# Example usage
prompt = "The quick brown"
generated_text = generate_text(prompt)
print(generated_text)

A full example

from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    Seq2SeqTrainingArguments,
    Seq2SeqTrainer,
)
from datasets import load_dataset

# Load the pre-trained tokenizer and model
model_name = "google/flan-t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Load and preprocess the dataset
dataset = load_dataset(
    "squad_v2", split="train[:1000]"
)  # Use a smaller subset for demonstration purposes


def preprocess_function(examples):
    inputs = [
        f"question: {q}  context: {c}"
        for q, c in zip(examples["question"], examples["context"])
    ]
    targets = [
        a["text"][0] if len(a["text"]) > 0 else "" for a in examples["answers"]
    ]  # Use the first answer as the target, if available
    model_inputs = tokenizer(inputs, max_length=1024, truncation=True, padding=True)
    labels = tokenizer(targets, max_length=128, truncation=True, padding=True)
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs


tokenized_dataset = dataset.map(
    preprocess_function, batched=True, remove_columns=dataset.column_names
)

# Set up the training arguments and trainer
training_args = Seq2SeqTrainingArguments(
    output_dir="output",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=1,
    weight_decay=0.01,
    save_total_limit=3,
    fp16=True,
)

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    eval_dataset=tokenized_dataset,
    tokenizer=tokenizer,
)

# Train the model
trainer.train()

# Save the trained model
trainer.save_model("trained_model")


# Ask a question to the trained model
def ask_question(question, context):
    input_text = f"question: {question}  context: {context}"
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    output = model.generate(input_ids, max_length=128, num_return_sequences=1)
    return tokenizer.decode(output[0], skip_special_tokens=True)


question = "What is the capital of France?"
context = "The capital of France is Paris. It is the largest city in France and one of the most iconic cities in the world, known for its art, fashion, and cuisine."

answer = ask_question(question, context)
print(f"Question: {question}")
print(f"Answer: {answer}")

In this project:

1. We load the pre-trained flan-t5-base tokenizer and model using the Hugging Face Transformers library.
2. We load a subset of the SQuAD v2.0 dataset and preprocess it by formatting the questions and contexts as input sequences and the answers as target sequences.
3. We set up the training arguments and create a Seq2SeqTrainer to train the model on the preprocessed dataset.
4. After training, we save the trained model for later use.
5. We define a function ask_question that takes a question and a context as input, encodes them using the tokenizer, generates an answer using the trained model, and decodes the generated output.
6. Finally, we ask a sample question to the trained model and print the generated answer.

Note: This project assumes you have the necessary dependencies installed, including the Hugging Face Transformers library and the datasets library. I used the following:

transformers==4.39.2
datasets==2.18.0
torch==2.2.2
nltk==3.8.1
accelerate==0.28.0

The output

python slm.py
{'eval_loss': 52.156742095947266, 'eval_runtime': 55.9558, 'eval_samples_per_second': 17.871, 'eval_steps_per_second': 4.468, 'epoch': 1.0}  
{'train_runtime': 255.4264, 'train_samples_per_second': 3.915, 'train_steps_per_second': 0.979, 'train_loss': 13.0256796875, 'epoch': 1.0}   
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████| 250/250 [04:15<00:00,  1.02s/it]
Question: What is the capital of France?
Answer: Paris