Fine-tuning: adapting a pretrained model

import os
os.environ["HF_HUB_OFFLINE"] = "1"
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import time
import numpy as np
import pandas as pd

messages = pd.read_csv(DATA / "mijoz_xabarlari.csv")
print(f"Messages: {len(messages)}")
print()
print(messages["maqsad"].value_counts().to_string())
print()
print("Samples:")
for _, q in messages.head(6).iterrows():
    label = "HARD" if q["qiyin"] else "easy"
    print(f"  [{q['maqsad']:<8}] [{label:<5}] {q['xabar']}")

from sklearn.dummy import DummyClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from sklearn.model_selection import train_test_split

classes = sorted(messages["maqsad"].unique())
class_index = {s: i for i, s in enumerate(classes)}
y = messages["maqsad"].map(class_index).values
X = messages["xabar"].values

tr, te = train_test_split(np.arange(len(messages)), test_size=0.25,
                          random_state=42, stratify=y)
print(f"train {len(tr)}, test {len(te)}, {len(classes)} classes: {classes}")
print()

baseline = DummyClassifier(strategy="most_frequent").fit(X[tr], y[tr])
tfidf = make_pipeline(TfidfVectorizer(ngram_range=(1, 2), min_df=1),
                      LogisticRegression(max_iter=2000)).fit(X[tr], y[tr])

t0 = time.perf_counter()
tfidf_time = time.perf_counter() - t0
print(f"Baseline (most frequent): {baseline.score(X[te], y[te]):.4f}")
print(f"TF-IDF + LogReg         : {tfidf.score(X[te], y[te]):.4f}")

hard_mask = messages["qiyin"].values[te].astype(bool)
print("We check TF-IDF part by part:")
print(f"  on EASY examples: {tfidf.score(X[te][~hard_mask], y[te][~hard_mask]):.4f}  ({(~hard_mask).sum()} examples)")
print(f"  on HARD examples: {tfidf.score(X[te][hard_mask], y[te][hard_mask]):.4f}  ({hard_mask.sum()} examples)")
print()
print("My hypothesis: TF-IDF falls over on the hard examples.")
print("The measurement: it didn't.")

show(pd.DataFrame([
    {"Approach": "Training from scratch", "Data": "millions", "Cost": "thousands of GPU-hours",
     "When": "Almost never: not your job"},
    {"Approach": "Fine-tuning (full)", "Data": "thousands", "Cost": "hours",
     "When": "Lots of data, task far from pretraining"},
    {"Approach": "LoRA / QLoRA", "Data": "hundreds to thousands", "Cost": "minutes",
     "When": "The usual choice: cheap and almost as good"},
    {"Approach": "Head layer only", "Data": "hundreds", "Cost": "minutes",
     "When": "Very little data"},
    {"Approach": "Prompt / RAG", "Data": "0", "Cost": "none",
     "When": "You need knowledge, not style: no training needed"},
]))

from transformers import AutoTokenizer

samples = ["Bugun to'lasam narx o'zgaradimi?",
           "Konditsioner buzilgan, tuzatib bering",
           "Bronni bekor qilmoqchiman"]
tok_table = []
for name in ["distilroberta-base", "distilbert-base-multilingual-cased"]:
    tok = AutoTokenizer.from_pretrained(name)
    tokens = [len(tok.tokenize(s)) for s in samples]
    chars = [len(s) for s in samples]
    tok_table.append({"model": name, "vocab size": f"{tok.vocab_size:,}",
                      "tokens": str(tokens),
                      "tokens/char": round(sum(tokens) / sum(chars), 3)})
show(pd.DataFrame(tok_table))

for name in ["distilroberta-base", "distilbert-base-multilingual-cased"]:
    tok = AutoTokenizer.from_pretrained(name)
    print(f"{name}:")
    print(f"  {tok.tokenize(samples[0])}")

"""Fine-tuning for text classification: full and LoRA."""
import time

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MODEL = "distilbert-base-multilingual-cased"
MAX_LENGTH = 48


def prepare(texts, model_name=MODEL, max_len=MAX_LENGTH):
    tok = AutoTokenizer.from_pretrained(model_name)
    return tok(list(texts), truncation=True, max_length=max_len,
               padding="max_length", return_tensors="pt")


def train(enc, y, train_idx, test_idx, n_classes, lora=False, epochs=3,
          lr=None, seed=0, model_name=MODEL):
    """Trains the model and returns (accuracy, trainable, total, time, model).

    With lora=True the base weights are FROZEN and only small
    adapter matrices are trained.
    """
    torch.manual_seed(seed)
    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=n_classes)
    if lora:
        from peft import LoraConfig, get_peft_model
        model = get_peft_model(model, LoraConfig(
            task_type="SEQ_CLS", r=8, lora_alpha=16, lora_dropout=0.05,
            target_modules=["q_lin", "v_lin"]))      # attention Q and V only

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())

    def subset(idx):
        return TensorDataset(enc["input_ids"][idx], enc["attention_mask"][idx],
                             torch.tensor(y[idx]))

    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad],
                            lr=lr or (5e-4 if lora else 3e-5))
    loader = DataLoader(subset(train_idx), batch_size=16, shuffle=True,
                        generator=torch.Generator().manual_seed(seed))
    t0 = time.perf_counter()
    model.train()
    for _ in range(epochs):
        for ids, mask, target in loader:
            opt.zero_grad()
            model(input_ids=ids, attention_mask=mask, labels=target).loss.backward()
            opt.step()
    elapsed = time.perf_counter() - t0

    model.eval()
    correct = 0
    with torch.no_grad():
        for ids, mask, target in DataLoader(subset(test_idx), batch_size=32):
            prediction = model(input_ids=ids, attention_mask=mask).logits.argmax(-1)
            correct += int((prediction == target).sum())
    return correct / len(test_idx), trainable, total, elapsed, model

import sys
sys.path.insert(0, str(ISH))
from adapt.train import prepare, train

enc = prepare(X)
print(f"Tokenization ready: {tuple(enc['input_ids'].shape)}")
print()
results = []
models = {}
for name, lora in [("full fine-tuning", False), ("LoRA (r=8)", True)]:
    accuracy, trainable, total, elapsed, model = train(enc, y, tr, te, len(classes), lora=lora)
    models[name] = model
    results.append({"method": name, "accuracy": round(accuracy, 4),
                    "trainable": trainable, "total params": total,
                    "share": f"{trainable / total:.2%}",
                    "time (s)": round(elapsed, 1)})
show(pd.DataFrame(results))

import shutil
from pathlib import Path

print(f"{'method':<22} {'on disk':>12}")
print("-" * 36)
for name, model in models.items():
    folder = ISH / ("saved_" + ("lora" if "LoRA" in name else "full"))
    shutil.rmtree(folder, ignore_errors=True)
    model.save_pretrained(folder)
    size = sum(p.stat().st_size for p in Path(folder).rglob("*") if p.is_file())
    print(f"{name:<22} {size / 1e6:>9.2f} MB")
print()
print("With LoRA only the adapter is saved: the base model is unchanged and doesn't need saving again.")

KEY = "trainable"
comparison = pd.DataFrame([
    {"method": "baseline (most frequent)", "accuracy": round(baseline.score(X[te], y[te]), 4),
     "trained params": "0", "note": "learns nothing"},
    {"method": "TF-IDF + LogReg", "accuracy": round(tfidf.score(X[te], y[te]), 4),
     "trained params": "thousands", "note": "trains in a second"},
    {"method": results[0]["method"], "accuracy": results[0]["accuracy"],
     "trained params": f"{results[0][KEY]:,}", "note": "every weight changes"},
    {"method": results[1]["method"], "accuracy": results[1]["accuracy"],
     "trained params": f"{results[1][KEY]:,}", "note": "adapter: a few MB"},
])
show(comparison)

show(pd.DataFrame([
    {"Pitfall": "lr too large", "Symptom": "loss is NaN or accuracy at chance level",
     "Fix": "2e-5..5e-5 for full FT; 1e-4..1e-3 for LoRA"},
    {"Pitfall": "Catastrophic forgetting", "Symptom": "the model is good at the new task, bad at everything else",
     "Fix": "fewer epochs, a small lr, or use LoRA"},
    {"Pitfall": "Many epochs on little data", "Symptom": "train accuracy 1.0, test stays low",
     "Fix": "early stopping, watch the validation set"},
    {"Pitfall": "Tuning on the test set", "Symptom": "good on paper, bad in practice",
     "Fix": "tune on validation: Chapters 6 and 24"},
    {"Pitfall": "Imbalanced classes", "Symptom": "the small class is never found",
     "Fix": "class weights or the F1 metric: Chapters 8 and 9"},
    {"Pitfall": "Random seed", "Symptom": "a different result on every run",
     "Fix": "3 seeds, and report the worst one too: Chapter 35"},
]))

show(pd.DataFrame([
    {"Situation": "Keywords determine the class", "Right solution": "TF-IDF + logistic regression",
     "Reason": "measured in this chapter: cheaper and higher"},
    {"Situation": "You need knowledge, not style", "Right solution": "RAG (retrieval + LLM)",
     "Reason": "adding new facts doesn't require training"},
    {"Situation": "Fewer than 100 examples", "Right solution": "prompt or few-shot",
     "Reason": "fine-tuning fits the noise"},
    {"Situation": "Requirements change often", "Right solution": "prompt",
     "Reason": "retraining on every change is expensive"},
    {"Situation": "You need an output format", "Right solution": "fine-tuning OR enforced structure",
     "Reason": "style and format are fine-tuning's strong side"},
    {"Situation": "Specialized domain language", "Right solution": "fine-tuning",
     "Reason": "medical, legal text: far from pretraining"},
]))

show(pd.DataFrame([
    {"Step": "37. Data collected", "Status": "done"},
    {"Step": "38. Dataset built", "Status": "done: split by group"},
    {"Step": "39. Baseline measured", "Status": "FOUND that TF-IDF is enough"},
    {"Step": "39. Fine-tuning learned", "Status": "full and LoRA: the mechanism is known"},
    {"Step": "40. Saving the model", "Status": "next chapter"},
    {"Step": "41. Serving", "Status": "Django API"},
]))

show(pd.DataFrame([
    {"Concept": "Baseline first", "In one sentence": "Measured: TF-IDF came out above the transformer"},
    {"Concept": "Fine-tuning", "In one sentence": "Start from pretrained weights and keep training on your data"},
    {"Concept": "LoRA", "In one sentence": "Base weights frozen, a small adapter is trained"},
    {"Concept": "LoRA's advantage", "In one sentence": "Memory, storage, swapping: not accuracy"},
    {"Concept": "Tokenization", "In one sentence": "Measured: both models chop Uzbek words into pieces"},
    {"Concept": "Multilingual model", "In one sentence": "Better than an English model for Uzbek text"},
    {"Concept": "Catastrophic forgetting", "In one sentence": "Many epochs break the old knowledge"},
    {"Concept": "When not to", "In one sentence": "When keywords are enough, or RAG or a prompt is"},
]))