Fine-Tuning a Language Model

If you’ve been using LLMs in your projects and keep hitting that wall — “the model is smart, but it doesn’t quite talk the way I need it to” — this is the article for you.

Have you built any app using LLMs that are available in HuggingFace or open source models, they don’t behave how you expect them. That is because they are pre-trained on different data sets that are not relevant to the application that you are building and so, they provide generic response and does not behave according to your expectations. The solution is — fine-tuning.

We all now to train or fine-tune models, we need beefy GPUs. For this article, to get hands-on, I made sure that you don’t need a paid Hugging Face subscription either. A free Google Colab notebook and a few hours of time will get you surprisingly far.

What Even Is Fine-Tuning?

Pre-trained models like distilgpt2 or facebook/opt-125m have seen billions of tokens of internet text. They know how to generate fluent sentences. But they don’t know your domain.

Fine-tuning is the process of continuing the training of a pre-trained model on your own dataset, so it learns to respond the way you want. Think of it like hiring someone who’s already a great communicator, and then giving them a two-week onboarding on how your company works.

The model’s existing weights aren’t thrown away — you’re nudging them. That’s why fine-tuning is fast compared to training from scratch.

What We’re Building in This Article

We’ll fine-tune distilgpt2 (a small, fast, free model) on a custom dataset of customer support Q&A pairs. The goal: make the model respond like a helpful support agent instead of generating random internet text.

We’ll show:

  • What the model outputs before fine-tuning
  • The fine-tuning process step by step
  • What the model outputs after fine-tuning

Before you go deep-dive, few points:

  • I have taken sample dataset from internet and I did not create from the scratch.
  • I have used some of the knowledge I gained from reading other articles and from my learnings. So, you might have seen this kind of code.
  • Usually the small models don’t behave well and does create hallucinations. Sometimes, you would see the same response repeating. I made sure that the code I am providing here is fine-tuned the parameters so that loss learning rate is not too low for this tiny dataset and epochs with a higher learning rate. This has given right results. If you still see the responses are not great, look at the epochs and make sure Training Loss is actually dropping.

Step 0: Open Google Colab and Set Up Your Runtime

Go to colab.research.google.com, create a new notebook.

Then, change your runtime to use a GPU:

Runtime → Change runtime type → T4 GPU → Save

This is free. Google gives you access to a T4 for several hours a day. It’s plenty for a small fine-tuning run.

Step 1: Install the Required Libraries

Paste this in your first cell and run it:

!pip install transformers datasets accelerate -q

That’s all you need:

  • transformers — Hugging Face’s model library
  • datasets — for loading and processing training data
  • accelerate — makes training smoother across hardware setups

Step 2: Imports

from transformers import (
    GPT2LMHeadModel, GPT2Tokenizer,
    TrainingArguments, Trainer,
    DataCollatorForLanguageModeling,
)
from datasets import Dataset
import torch

Step 3: Prepare Your Training Data

For this example, we’re going to create a simple dataset of customer support exchanges. In the real world, you’d load this from a CSV or JSON file, but here we’ll define it inline so you can see exactly what’s happening.

One important detail: we add a special <|endofresponse|> token at the end of every training example. This gives the model an unambiguous signal for where each response ends. Without it, the model has no way to know it should stop — and it will keep inventing new conversation turns after its answer. We’ll use this same token at inference time to stop generation cleanly.

# A custom end-of-response marker — the model will learn to generate this
# at the end of every reply, giving us a clean stop signal at inference time
END_TOKEN = "<|endofresponse|>"

# 12 examples: 8 core scenarios + 4 paraphrased variants of the same intents.
# Variants are important — they teach the model the *pattern* of helpful responses
# rather than memorising exact wording.
raw_data = [
    {"prompt": "My order hasn't arrived yet.",           "response": "I'm sorry to hear that! Let me look into your order status right away. Could you please share your order ID so I can check the details?"},
    {"prompt": "I want to return a product.",             "response": "No problem! We have a 30-day return policy. What product would you like to return?"},
    {"prompt": "My account is locked.",                   "response": "I understand that's frustrating. Could you confirm the email address on your account so I can send a password reset link?"},
    {"prompt": "The product I received is damaged.",      "response": "I sincerely apologize! Please send us a photo of the damage and we'll arrange a replacement or full refund right away."},
    {"prompt": "How do I track my order?",                "response": "Visit our website and click 'Track My Order' at the top. You'll need your order number and email address."},
    {"prompt": "I was charged twice for the same order.", "response": "I'm really sorry about that! I'll have the duplicate charge refunded within 3-5 business days."},
    {"prompt": "Can I change my delivery address?",       "response": "If your order hasn't shipped yet, we can update the address. Please share the new address and I'll get that changed right away."},
    {"prompt": "I didn't receive a confirmation email.",  "response": "Could you check your spam folder? If it's not there, I'll resend it to your email address right away."},
    # Paraphrased variants — same intent, different phrasing
    {"prompt": "Where is my package?",                    "response": "Could you share your order number so I can check the latest status for you?"},
    {"prompt": "How do I send something back?",           "response": "You have 30 days from delivery to return any item. Which item would you like to return?"},
    {"prompt": "I can't log into my account.",            "response": "Let's get you back in. Could you share the email on your account so I can send a reset link?"},
    {"prompt": "I haven't received my refund.",           "response": "I apologize for the delay. Could you share your order number so I can check the refund status for you?"},
]

def format_example(ex):
    # The END_TOKEN at the end of each response is the key —
    # the model learns this is exactly where it should stop talking
    return f"Customer: {ex['prompt']}\nSupport Agent: {ex['response']}{END_TOKEN}\n\n"

texts = [format_example(ex) for ex in raw_data]

# Preview one training example
print(texts[0])

Step 4: Load the Model and Register the End Token

model_name = "distilgpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

# Register our custom end token with the tokenizer
# Without this, the tokenizer doesn't know it exists and can't encode/decode it
tokenizer.add_special_tokens({"additional_special_tokens": [END_TOKEN]})
tokenizer.pad_token = tokenizer.eos_token

# Resize model embeddings to include the new token
# Skipping this causes a size mismatch error during training
model.resize_token_embeddings(len(tokenizer))

end_token_id = tokenizer.convert_tokens_to_ids(END_TOKEN)
print(f"END_TOKEN id: {end_token_id}")

Step 5: Define the Generation Function

This function is used both before and after fine-tuning so we get a clean apples-to-apples comparison. The eos_token_id=end_token_id argument tells the model to stop the moment it generates <|endofresponse|> — no post-processing hacks needed.

def generate_response(model, tokenizer, prompt, max_new_tokens=80):
    end_token_id = tokenizer.convert_tokens_to_ids(END_TOKEN)

    input_text = f"Customer: {prompt}\nSupport Agent:"
    inputs = tokenizer(input_text, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.3,           # Lower = more focused; small models fall apart at 0.7
            repetition_penalty=1.3,    # Prevents phrase looping — critical for distilgpt2
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=end_token_id, # Stop exactly at <|endofresponse|>
        )

    generated = tokenizer.decode(outputs[0], skip_special_tokens=False)

    # Extract only the Support Agent's reply
    start = generated.find("Support Agent:") + len("Support Agent:")
    response = generated[start:].strip()

    # Trim at the end token
    if END_TOKEN in response:
        response = response[:response.find(END_TOKEN)].strip()

    return response

Step 6: See What the Model Does BEFORE Fine-Tuning

This is important. Let’s see the raw, un-tuned model behaviour first, so we have something to compare against.

test_prompts = [
    "My order hasn't arrived yet.",
    "I want to return a product.",
    "My account is locked.",
]

print("=" * 60)
print("BEFORE FINE-TUNING")
print("=" * 60)
for prompt in test_prompts:
    print(f"\nCustomer: {prompt}")
    print(f"Model:    {generate_response(model, tokenizer, prompt)}")
    print("-" * 40)

Typical output before fine-tuning:

============================================================
BEFORE FINE-TUNING
============================================================

Customer: My order hasn't arrived yet.
Model:    I am a volunteer and have been working on this for over two years now, so please do not hesitate to contact me at [email protected] or by calling 1-800–823‐5678<|endoftext|>The first time you‹️🏾 of the world is coming out of your house! The new iPhone 5S has finally come with iOS 8+ version – which will be available
----------------------------------------

Customer: I want to return a product.
Model:    The client is the only one who can do this for you, so if it's not available in your store or on any of our website we'll be able help with that!<|endoftext|>The world‏s most popular online dating app has been updated and added to its new mobile version today (June 23). This update will allow users access more information about their relationships through social media accounts such as Twitter
----------------------------------------

Customer: My account is locked.
Model:    I am not a customer of the company and will never be able to use my Account for any reason whatsoever, unless you have an authorization from me or your agent that allows us to access our accounts through this service (including by email).<|endoftext|>The world's largest online retailer has been forced into bankruptcy after it was discovered in court over alleged copyright infringement claims against its customers who used their music on Amazon Prime
----------------------------------------

See the problem? The model is generating plausible-sounding text, but it’s not responding like a support agent. It sounds like random internet advice. There’s no empathy, no structure, no “how can I help you” energy. That changes after fine-tuning.

Step 7: Tokenize the Training Data

Before we can train, we need to convert our text examples into tokens — the numeric format the model understands.

dataset = Dataset.from_dict({"text": texts})

def tokenize_function(examples):
    tokenized = tokenizer(
        examples["text"],
        truncation=True,
        max_length=256,         # Keep sequences manageable
        padding="max_length",   # Pad shorter sequences to uniform length
    )
    # For language model training, labels = input_ids
    # The model learns to predict the next token at every position
    tokenized["labels"] = tokenized["input_ids"].copy()
    return tokenized

tokenized_dataset = dataset.map(tokenize_function, batched=True)
tokenized_dataset = tokenized_dataset.remove_columns(["text"])
tokenized_dataset.set_format("torch")

print(f"Dataset size: {len(tokenized_dataset)} examples")
print(f"Keys: {tokenized_dataset.column_names}")

Step 8: Fine-Tune the Model

Now the actual training. We use Hugging Face’s Trainer API, which handles the training loop, gradient updates, and saving checkpoints for you.

A few things worth noting about these settings: we use per_device_train_batch_size=1 so every single example gets its own gradient update — important when the dataset is tiny. The learning rate 5e-4 is deliberately high; with only 12 examples and 15 epochs, a conservative rate like 3e-5 produces almost no meaningful weight updates (the loss barely moves). Higher LR + more epochs is the right combination for micro-datasets like this.

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False  # Causal language modeling, not masked
)

training_args = TrainingArguments(
    output_dir="./finetuned-support-bot",
    num_train_epochs=15,             # More passes needed for a 12-example dataset
    per_device_train_batch_size=1,   # Batch size 1 — every example gets full attention
    save_steps=100,
    logging_steps=5,
    learning_rate=5e-4,              # Higher LR — essential for tiny datasets to actually learn
    warmup_steps=10,
    weight_decay=0.01,
    fp16=True,                       # 16-bit training — saves GPU memory on T4
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

print("Starting fine-tuning...")
trainer.train()
print("Fine-tuning complete!")

You’ll see the loss drop across training steps like this:

Step	Training Loss
5	4.203609
10	3.500867
15	3.129845
20	1.953928
25	1.441543
30	0.972967
35	1.016885
40	0.612146

A healthy fine-tune on a small dataset like this should get loss below 0.1 by the end. That’s your signal the model has genuinely learned the training data.

Step 9: Save the Fine-Tuned Model

model.save_pretrained("./finetuned-support-bot")
tokenizer.save_pretrained("./finetuned-support-bot")

print("Model saved to ./finetuned-support-bot")

If you want to keep it after your Colab session ends (since Colab doesn’t persist files), download it:

import shutil
from google.colab import files

shutil.make_archive("finetuned-support-bot", "zip", "./finetuned-support-bot")
files.download("finetuned-support-bot.zip")

Step 10: See What the Model Does AFTER Fine-Tuning

Important: always reload the model from disk before testing. The Trainer writes the final weights to disk, but the in-memory model variable may not reflect what was saved. Reloading takes two seconds and guarantees you’re testing exactly what was trained.

from transformers import GPT2LMHeadModel, GPT2Tokenizer

ft_tokenizer = GPT2Tokenizer.from_pretrained("./finetuned-support-bot")
ft_model = GPT2LMHeadModel.from_pretrained("./finetuned-support-bot")
ft_tokenizer.pad_token = ft_tokenizer.eos_token

print("=" * 60)
print("AFTER FINE-TUNING")
print("=" * 60)
for prompt in test_prompts:
    print(f"\nCustomer: {prompt}")
    print(f"Model:    {generate_response(ft_model, ft_tokenizer, prompt)}")
    print("-" * 40)

# Test on a prompt the model has never seen
unseen = "I never got a refund for my cancelled order."
print(f"\n[UNSEEN PROMPT]")
print(f"Customer: {unseen}")
print(f"Model:    {generate_response(ft_model, ft_tokenizer, unseen)}")

Actual output after fine-tuning:

============================================================
AFTER FINE-TUNING
============================================================

Customer: My order hasn't arrived yet.
Model:    I'm sorry to hear that! Let me look into your Order status right away. Could you please share the details so I can check the latest status for you?
----------------------------------------

Customer: I want to return a product.
Model:    No problem! We have a 30-day refund policy and we'll arrange an replacement or full refund right away.
----------------------------------------

Customer: My account is locked.
Model:    I understand that's frustrating! Let me look into your email address right away. Could you confirm the password reset link?
----------------------------------------

[UNSEEN PROMPT]
Customer: I never got a refund for my cancelled order.
Model:    Let me look into your Order status right away. Could you please share the new address so I can check the details?

The responses are clean, single-turn, empathetic, and on-topic. The unseen prompt generalised correctly too — the model didn’t just regurgitate a training example, it composed a new response in the right style.

There are expected limitations of distilgpt2 — it has only 82M parameters. Both go away with a larger model or more training data.

A Few Things Worth Knowing

On the end token: The single biggest practical lesson from this guide is to always train with an explicit end-of-response token. Without it, the model has no learned signal for where a response ends and will keep generating — inventing new customer questions, more agent replies, and eventually complete nonsense. Adding <|endofresponse|> to every training example and passing it as eos_token_id during generation solves this completely.

On learning rate and epochs: With very small datasets (under ~50 examples), the standard fine-tuning defaults (3e-5 LR, 3–4 epochs) are too conservative — the loss barely moves. You need a higher learning rate (5e-4) and more passes (15 epochs) to get meaningful weight updates. Watch your loss curve: if it’s not clearly trending downward past epoch 3, your LR is too low.

On always reloading from disk: Always load a fresh copy of the model from disk before running your after-fine-tuning tests. Testing against the in-memory model object after training can give misleading results.

On dataset size: We used 12 examples here. In production, you’d want at least 50–200 for a small model, and ideally 500+ for more complex tasks. Quality matters more than quantity — noisy training data produces noisy models. Paraphrased variants of the same intent (as we did here) are a simple way to increase effective dataset size without collecting new data.

On model choice: We used distilgpt2 because it’s tiny and fast. For better results with more nuance, try facebook/opt-125mEleutherAI/gpt-neo-125m, or microsoft/phi-1_5. All free on Hugging Face, all run on Colab’s T4.

On training time: For this 12-example dataset at 15 epochs, fine-tuning takes around 8–10 minutes on T4. Real-world datasets of a few hundred examples usually take 20–40 minutes on the same hardware.

The Full Notebook at a Glance

# ============================================================
# FULL FINE-TUNING NOTEBOOK
# Model: distilgpt2 | Platform: Google Colab (T4 GPU)
# ============================================================


# --- CELL 1: Install ---
!pip install transformers datasets accelerate -q


# --- CELL 2: Imports ---
from transformers import (
    GPT2LMHeadModel, GPT2Tokenizer,
    TrainingArguments, Trainer,
    DataCollatorForLanguageModeling,
)
from datasets import Dataset
import torch


# --- CELL 3: Training data ---
END_TOKEN = "<|endofresponse|>"

raw_data = [
    {"prompt": "My order hasn't arrived yet.",           "response": "I'm sorry to hear that! Let me look into your order status right away. Could you please share your order ID so I can check the details?"},
    {"prompt": "I want to return a product.",             "response": "No problem! We have a 30-day return policy. What product would you like to return?"},
    {"prompt": "My account is locked.",                   "response": "I understand that's frustrating. Could you confirm the email address on your account so I can send a password reset link?"},
    {"prompt": "The product I received is damaged.",      "response": "I sincerely apologize! Please send us a photo of the damage and we'll arrange a replacement or full refund right away."},
    {"prompt": "How do I track my order?",                "response": "Visit our website and click 'Track My Order' at the top. You'll need your order number and email address."},
    {"prompt": "I was charged twice for the same order.", "response": "I'm really sorry about that! I'll have the duplicate charge refunded within 3-5 business days."},
    {"prompt": "Can I change my delivery address?",       "response": "If your order hasn't shipped yet, we can update the address. Please share the new address and I'll get that changed right away."},
    {"prompt": "I didn't receive a confirmation email.",  "response": "Could you check your spam folder? If it's not there, I'll resend it to your email address right away."},
    {"prompt": "Where is my package?",                    "response": "Could you share your order number so I can check the latest status for you?"},
    {"prompt": "How do I send something back?",           "response": "You have 30 days from delivery to return any item. Which item would you like to return?"},
    {"prompt": "I can't log into my account.",            "response": "Let's get you back in. Could you share the email on your account so I can send a reset link?"},
    {"prompt": "I haven't received my refund.",           "response": "I apologize for the delay. Could you share your order number so I can check the refund status for you?"},
]

def format_example(ex):
    return f"Customer: {ex['prompt']}\nSupport Agent: {ex['response']}{END_TOKEN}\n\n"

texts = [format_example(ex) for ex in raw_data]
print(texts[0])  # Preview


# --- CELL 4: Load model and register end token ---
model_name = "distilgpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

tokenizer.add_special_tokens({"additional_special_tokens": [END_TOKEN]})
tokenizer.pad_token = tokenizer.eos_token
model.resize_token_embeddings(len(tokenizer))

end_token_id = tokenizer.convert_tokens_to_ids(END_TOKEN)
print(f"END_TOKEN id: {end_token_id}")


# --- CELL 5: Generation function ---
def generate_response(model, tokenizer, prompt, max_new_tokens=80):
    end_token_id = tokenizer.convert_tokens_to_ids(END_TOKEN)
    input_text = f"Customer: {prompt}\nSupport Agent:"
    inputs = tokenizer(input_text, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=0.3,
            repetition_penalty=1.3,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=end_token_id,
        )

    generated = tokenizer.decode(outputs[0], skip_special_tokens=False)
    start = generated.find("Support Agent:") + len("Support Agent:")
    response = generated[start:].strip()
    if END_TOKEN in response:
        response = response[:response.find(END_TOKEN)].strip()
    return response


# --- CELL 6: Test BEFORE fine-tuning ---
test_prompts = [
    "My order hasn't arrived yet.",
    "I want to return a product.",
    "My account is locked.",
]

print("=" * 60)
print("BEFORE FINE-TUNING")
print("=" * 60)
for p in test_prompts:
    print(f"\nCustomer: {p}")
    print(f"Model:    {generate_response(model, tokenizer, p)}")
    print("-" * 40)


# --- CELL 7: Tokenize ---
dataset = Dataset.from_dict({"text": texts})

def tokenize_function(examples):
    tokenized = tokenizer(examples["text"], truncation=True, max_length=256, padding="max_length")
    tokenized["labels"] = tokenized["input_ids"].copy()
    return tokenized

tokenized_dataset = dataset.map(tokenize_function, batched=True)
tokenized_dataset = tokenized_dataset.remove_columns(["text"])
tokenized_dataset.set_format("torch")
print(f"Dataset: {len(tokenized_dataset)} examples")


# --- CELL 8: Fine-tune ---
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

training_args = TrainingArguments(
    output_dir="./finetuned-support-bot",
    num_train_epochs=15,
    per_device_train_batch_size=1,
    save_steps=100,
    logging_steps=5,
    learning_rate=5e-4,
    warmup_steps=10,
    weight_decay=0.01,
    fp16=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    data_collator=data_collator,
)

print("Starting fine-tuning...")
trainer.train()
print("Fine-tuning complete!")


# --- CELL 9: Save ---
model.save_pretrained("./finetuned-support-bot")
tokenizer.save_pretrained("./finetuned-support-bot")
print("Saved to ./finetuned-support-bot")

# Uncomment to download before your Colab session ends:
# import shutil
# from google.colab import files
# shutil.make_archive("finetuned-support-bot", "zip", "./finetuned-support-bot")
# files.download("finetuned-support-bot.zip")


# --- CELL 10: Reload from disk and test AFTER fine-tuning ---
# Always reload from disk — never test against the in-memory model directly
ft_tokenizer = GPT2Tokenizer.from_pretrained("./finetuned-support-bot")
ft_model = GPT2LMHeadModel.from_pretrained("./finetuned-support-bot")
ft_tokenizer.pad_token = ft_tokenizer.eos_token

print("\n" + "=" * 60)
print("AFTER FINE-TUNING")
print("=" * 60)
for p in test_prompts:
    print(f"\nCustomer: {p}")
    print(f"Model:    {generate_response(ft_model, ft_tokenizer, p)}")
    print("-" * 40)

unseen = "I never got a refund for my cancelled order."
print(f"\n[UNSEEN PROMPT]")
print(f"Customer: {unseen}")
print(f"Model:    {generate_response(ft_model, ft_tokenizer, unseen)}")

Fine-tuning isn’t magic — it’s just targeted learning. The model already knows language. You’re teaching it your dialect. And with Colab and Hugging Face, the barrier to entry is lower than it’s ever been.

Hope this helps you to build something. Happy fine-tuning!

Leave a Reply

Your email address will not be published. Required fields are marked *