LoRA, PEFT & Quantization

# original: output = x @ W          (W is frozen, huge)
# LoRA:     output = x @ W  +  x @ (B @ A) * scale
#                              └── the only trained part ──┘

# A: (d × r),  B: (r × d),  where rank r is TINY (e.g. 8)
# so B@A has W's shape but only 2*d*r params instead of d*d

# Colab currently includes an old optional torchao package that conflicts
# with recent PEFT. Basic LoRA does not use torchao, so remove it.
!pip -q uninstall -y torchao
!pip -q install -U transformers datasets peft accelerate

import torch
from torch.utils.data import DataLoader
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    DataCollatorWithPadding,
)
from peft import LoraConfig, TaskType, get_peft_model

device = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "distilbert/distilbert-base-uncased"  # maintained 67M model

# 1. Load a real sentiment dataset and keep a small classroom-sized subset.
# Use the full Hub repository ID; the old "rotten_tomatoes" shortcut can
# produce an HfUriError with newer datasets/huggingface_hub versions.
dataset = load_dataset("cornell-movie-review-data/rotten_tomatoes")
train_data = dataset["train"].shuffle(seed=42).select(range(1200))
valid_data = dataset["validation"].select(range(300))

# 2. Load the tokenizer from the same maintained DistilBERT repository.
tokenizer = AutoTokenizer.from_pretrained(model_name)
def tokenize(batch):
    return tokenizer(batch["text"], truncation=True, max_length=128)

train_data = train_data.map(tokenize, batched=True, remove_columns=["text"])
valid_data = valid_data.map(tokenize, batched=True, remove_columns=["text"])
train_data = train_data.rename_column("label", "labels")
valid_data = valid_data.rename_column("label", "labels")
collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="pt")
train_loader = DataLoader(train_data, batch_size=16, shuffle=True, collate_fn=collator)
valid_loader = DataLoader(valid_data, batch_size=32, collate_fn=collator)

# 3. Load the base model. At this point it is an ordinary DistilBERT classifier.
base_model = AutoModelForSequenceClassification.from_pretrained(
    model_name, num_labels=2
)

# 4. Choose LoRA, which is ONE method provided by the PEFT library.
# DistilBERT calls its attention projections "q_lin" and "v_lin".
lora_config = LoraConfig(
    task_type=TaskType.SEQ_CLS,
    r=8,                    # adapter capacity
    lora_alpha=16,          # scales the adapter update
    lora_dropout=0.05,
    target_modules=["q_lin", "v_lin"],
)
model = get_peft_model(base_model, lora_config).to(device)
model.print_trainable_parameters()  # verify that only a small fraction trains

# 5. This is the familiar training loop. Gradients exist only for the
# LoRA adapter and classification head; the base transformer stays frozen.
optimizer = torch.optim.AdamW(
    (p for p in model.parameters() if p.requires_grad),
    lr=5e-4,
)

model.train()
for epoch in range(3):
    running_loss = 0.0
    for batch in train_loader:
        batch = {key: value.to(device) for key, value in batch.items()}
        optimizer.zero_grad()
        output = model(**batch)
        output.loss.backward()
        optimizer.step()
        running_loss += output.loss.item()
    print(f"epoch {epoch + 1} | loss {running_loss / len(train_loader):.4f}")

# 6. Evaluate on reviews the adapter did not train on.
model.eval()
correct = total = 0
with torch.no_grad():
    for batch in valid_loader:
        batch = {key: value.to(device) for key, value in batch.items()}
        predictions = model(**batch).logits.argmax(dim=-1)
        correct += (predictions == batch["labels"]).sum().item()
        total += predictions.numel()
print(f"validation accuracy: {correct / total:.1%}")

# 7. Fun test: replace these with your own difficult or sarcastic reviews.
reviews = [
    "A clever, warm movie I would happily watch again.",
    "Two hours of my life that I will never get back.",
]
inputs = tokenizer(reviews, padding=True, truncation=True, return_tensors="pt").to(device)
with torch.no_grad():
    probabilities = model(**inputs).logits.softmax(dim=-1).cpu()

for review, probability in zip(reviews, probabilities):
    label = "POSITIVE" if probability[1] > probability[0] else "NEGATIVE"
    confidence = probability.max().item()
    print(f"{label:8} {confidence:.1%} | {review}")

# 8. Save only the adapter files, not another copy of the base model.
model.save_pretrained("movie-mood-lora-adapter")