What instruction tuning is and why it differs from base pre-training

import json
from datasets import Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig

# Step 1: Create instruction tuning data (input-output pairs)
instruction_data = [
    {
        "instruction": "Classify the sentiment of this review.",
        "input": "The product broke after one week. Total waste of money.",
        "output": "Negative"
    },
    {
        "instruction": "Classify the sentiment of this review.",
        "input": "Arrived on time, works perfectly, great value!",
        "output": "Positive"
    },
    {
        "instruction": "Summarize this text in one sentence.",
        "input": "Machine learning models require large amounts of labeled data. The quality of this data directly impacts model performance. Preprocessing and validation are critical steps.",
        "output": "Data quality and preprocessing are essential for machine learning model performance."
    },
    {
        "instruction": "Translate to French.",
        "input": "Hello, how are you?",
        "output": "Bonjour, comment allez-vous?"
    }
]

# Step 2: Format data for instruction tuning (instruction + input → output)
def format_instruction_example(example):
    if example["input"]:
        return f"Instruction: {example['instruction']}\nInput: {example['input']}\nOutput: {example['output']}"
    else:
        return f"Instruction: {example['instruction']}\nOutput: {example['output']}"

formatted_data = [
    {"text": format_instruction_example(ex)}
    for ex in instruction_data
]

print("Example formatted instruction:")
print(formatted_data[0]["text"])
print("\n" + "="*60 + "\n")

# Step 3: Compare pre-training vs instruction tuning loss targets
print("PRE-TRAINING LOSS TARGET:")
print("Given: 'The quick brown fox jumps'")
print("Predict: 'over' (next token)")
print("Loss computed on: predicting the next word unconditionally\n")

print("INSTRUCTION TUNING LOSS TARGET:")
print("Given: 'Classify sentiment. Input: Great product! Output:'")
print("Predict: 'Positive' (conditional on instruction)")
print("Loss computed ONLY on: the output portion, not instruction/input\n")

# Step 4: Show the tokenization difference
tokenizer = AutoTokenizer.from_pretrained("gpt2")

inst_example = "Classify sentiment. Input: Good movie! Output: Positive"
tokens = tokenizer.encode(inst_example)
print(f"Full example tokenized ({len(tokens)} tokens):")
print(f"{inst_example}")
print(f"Token IDs: {tokens}\n")

# In instruction tuning, we only compute loss on token indices corresponding to "Positive"
output_start_idx = inst_example.find("Output: ") + len("Output: ")
output_tokens = tokenizer.encode("Positive")
print(f"Loss computed ONLY on these tokens: {output_tokens}")
print(f"(tokens for the answer 'Positive')")
print("\n" + "="*60 + "\n")

# Step 5: Demonstrate with a minimal trainer config (no actual training)
config = SFTConfig(
    output_dir="./instruction_tuned",
    max_steps=2,
    per_device_train_batch_size=1,
    learning_rate=2e-4,
    logging_steps=1,
)

print(f"Instruction Tuning Config Summary:")
print(f"- Loss function: Causal Language Modeling (still)")
print(f"- Data format: (instruction, input, output) tuples")
print(f"- Loss target: output tokens only (vs. all tokens in pre-training)")
print(f"- Learning rate: {config.learning_rate} (much lower than pre-training)")
print(f"- Typical duration: 1-3 epochs (vs. pre-training: 3-10 epochs)")