TDM 20100: Project 12 - Fine-Tuning a Large Language Model

# launch the ollama server (leave this Terminal running)
/anvil/projects/tdm/bin/ollama serve

import os

# tell this notebook where the Ollama server is running
with open(f"/dev/shm/ollama.{os.getuid()}") as hostfile:
    hostline = [line.rstrip() for line in hostfile]
os.environ["OLLAMA_HOST"] = hostline[0]
print(os.environ["OLLAMA_HOST"])

%%bash
# download the small model we will fine-tune
/anvil/projects/tdm/bin/ollama pull smollm2:135m
/anvil/projects/tdm/bin/ollama list

%%bash
cat > ~/smollm2.modelfile << HERE
FROM smollm2:135m
PARAMETER num_thread 2
HERE
/anvil/projects/tdm/bin/ollama create smollm2-135m-2 -f ~/smollm2.modelfile

%%bash
/anvil/projects/tdm/bin/ollama run smollm2-135m-2 "Write a short passage in the style of Shakespeare:"

import urllib.request

url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
text = urllib.request.urlopen(url).read().decode("utf-8")

print("Total characters:", len(text))
print(text[:300])

import json, random
random.seed(42)

# break the text into ~400-character chunks
chunk_size = 400
chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]

# keep reasonably sized chunks, then cap the count so training stays fast
examples = [c for c in chunks if len(c.strip()) > 100]
random.shuffle(examples)
examples = examples[:300]

# hold out 10 chunks to evaluate on later (the model will NOT train on these)
eval_examples = examples[:10]
train_examples = examples[10:]

print("Training examples:", len(train_examples))
print("Held-out examples:", len(eval_examples))

# save the training set so we have a record of it
with open("shakespeare_train.jsonl", "w") as f:
    for ex in train_examples:
        f.write(json.dumps({"text": ex}) + "\n")

%pip install -q transformers peft datasets accelerate

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

torch.set_num_threads(4)   # use our 4 cores (and no more)
model_name = "HuggingFaceTB/SmolLM2-135M-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token   # the model needs a padding token

model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)
print("Loaded", model_name)

import math

def perplexity(model, texts):
    """Average perplexity of the model over a list of texts (CPU-friendly)."""
    model.eval()
    total_loss, total_tokens = 0.0, 0
    for t in texts:
        enc = tokenizer(t, return_tensors="pt", truncation=True, max_length=256)
        with torch.no_grad():
            out = model(**enc, labels=enc["input_ids"])
        n = enc["input_ids"].size(1)
        total_loss += out.loss.item() * n     # loss is the average per token
        total_tokens += n
    return math.exp(total_loss / total_tokens)

baseline_ppl = perplexity(model, eval_examples)
print("Baseline perplexity (before fine-tuning):", round(baseline_ppl, 2))

from peft import LoraConfig, get_peft_model, TaskType

lora_config = LoraConfig(
    r=8,                       # size of the small added weights
    lora_alpha=16,
    lora_dropout=0.05,
    target_modules=["q_proj", "v_proj"],   # which layers LoRA attaches to
    task_type=TaskType.CAUSAL_LM,
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()   # shows how few weights we actually train

from datasets import Dataset

train_ds = Dataset.from_list([{"text": t} for t in train_examples])

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True, max_length=192)

train_ds = train_ds.map(tokenize, batched=True, remove_columns=train_ds.column_names)
print(train_ds)

from transformers import TrainingArguments, Trainer, DataCollatorForLanguageModeling

# mlm=False tells the collator we are doing normal (causal) language modeling
collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

args = TrainingArguments(
    output_dir="smollm2-lora",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    num_train_epochs=3,
    learning_rate=2e-4,
    logging_steps=10,
    save_strategy="no",
    report_to="none",
    use_cpu=True,            # force CPU - there is no GPU
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train_ds,
    data_collator=collator,
)

trainer.train()   # expect roughly 10-20 minutes on a 4-core node

merged = model.merge_and_unload()          # fold LoRA into the base weights
save_dir = "smollm2-shakespeare-merged"
merged.save_pretrained(save_dir)
tokenizer.save_pretrained(save_dir)

import os
print("Saved files:", os.listdir(save_dir))

finetuned_ppl = perplexity(merged, eval_examples)

print("Baseline perplexity   :", round(baseline_ppl, 2))
print("Fine-tuned perplexity :", round(finetuned_ppl, 2))

%%bash
cat > Modelfile << HERE
FROM ./smollm2-shakespeare-merged
PARAMETER num_thread 2
PARAMETER num_predict 200
PARAMETER repeat_penalty 1.3
TEMPLATE "{{ .Prompt }}"
HERE
/anvil/projects/tdm/bin/ollama create smollm2-shakespeare -f Modelfile
/anvil/projects/tdm/bin/ollama list

%%bash
echo "===== BEFORE (base model) ====="
/anvil/projects/tdm/bin/ollama run smollm2-135m-2 "ROMEO:"

echo ""
echo "===== AFTER (fine-tuned model) ====="
/anvil/projects/tdm/bin/ollama run smollm2-shakespeare "ROMEO:"