🦙 Llama from install to fine-tuning

import torch

print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    props = torch.cuda.get_device_properties(0)
    print("GPU:", props.name)
    print(f"VRAM: {props.total_memory / 1024**3:.1f} GB")
    print("BF16 supported:", torch.cuda.is_bf16_supported())   # if False, train in fp16

from ollama import chat

messages = [{"role": "system", "content": "You are an SRE. Answer step by step and keep it short."}]

while True:
    user = input("you> ").strip()
    if user in {"exit", "quit"}:
        break
    messages.append({"role": "user", "content": user})

    answer = ""
    for part in chat(model="llama3.2", messages=messages, stream=True, options={"temperature": 0.2}):
        print(part.message.content, end="", flush=True)
        answer += part.message.content
    print()

    messages.append({"role": "assistant", "content": answer})   # keep the conversation context

import torch
from transformers import AutoTokenizer, pipeline

MODEL_DIR = "./models/llama-3.2-3b-instruct"
messages = [
    {"role": "system", "content": "You are an SRE."},
    {"role": "user", "content": "A disk usage alert just fired"},
]

# ── See the exact string the model receives ──────
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
print(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True))

# ── Run inference with Transformers ──────────────
pipe = pipeline("text-generation", model=MODEL_DIR, dtype=torch.bfloat16, device_map="auto")
out = pipe(messages, max_new_tokens=256, do_sample=False)
print(out[0]["generated_text"][-1]["content"])

from datasets import load_dataset
from transformers import AutoTokenizer

MODEL_DIR = "./models/llama-3.2-3b-instruct"
MAX_LENGTH = 1024
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)

ds = load_dataset("json", data_files="data/incidents.jsonl", split="train")

# ── 1) Drop duplicate questions ──────────────────
seen = set()
def first_time(example) -> bool:
    key = example["prompt"][-1]["content"].strip()
    if key in seen:
        return False
    seen.add(key)
    return True
ds = ds.filter(first_time)

# ── 2) Check token length — beyond MAX_LENGTH the end (the answer) is cut off ─
def count_tokens(example) -> dict:
    enc = tokenizer.apply_chat_template(example["prompt"] + example["completion"], tokenize=True, return_dict=True)
    return {"n_tokens": len(enc["input_ids"])}
ds = ds.map(count_tokens)
too_long = ds.filter(lambda e: e["n_tokens"] > MAX_LENGTH)
print(f"{len(ds)} examples, max {max(ds['n_tokens'])} tokens, {len(too_long)} over {MAX_LENGTH}")
ds = ds.filter(lambda e: e["n_tokens"] <= MAX_LENGTH).remove_columns("n_tokens")

# ── 3) Train / eval split ────────────────────────
split = ds.train_test_split(test_size=0.1, seed=42)
split["train"].to_json("data/train.jsonl", force_ascii=False)
split["test"].to_json("data/eval.jsonl", force_ascii=False)
print("train:", len(split["train"]), "eval:", len(split["test"]))

import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import BitsAndBytesConfig
from trl import SFTConfig, SFTTrainer

MODEL_DIR = "./models/llama-3.2-3b-instruct"   # scale up: meta-llama/Llama-3.1-8B-Instruct
OUTPUT_DIR = "./outputs/llama-3.2-3b-incident-lora"

data = load_dataset("json", data_files={"train": "data/train.jsonl", "eval": "data/eval.jsonl"})
bf16 = torch.cuda.is_bf16_supported()

trainer = SFTTrainer(
    model=MODEL_DIR,
    train_dataset=data["train"],
    eval_dataset=data["eval"],
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16 if bf16 else torch.float16,
        bnb_4bit_use_double_quant=True,
    ),
    peft_config=LoraConfig(
        r=16, lora_alpha=32, lora_dropout=0.05,
        target_modules="all-linear", task_type="CAUSAL_LM",
    ),
    args=SFTConfig(
        output_dir=OUTPUT_DIR,
        num_train_epochs=3,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,       # effective batch = 4 × 4 = 16
        learning_rate=1e-4,
        lr_scheduler_type="cosine",
        warmup_steps=10,
        max_length=1024,
        bf16=bf16,
        fp16=not bf16,
        logging_steps=10,
        eval_strategy="steps",
        eval_steps=50,
        save_strategy="steps",
        save_steps=50,
        save_total_limit=2,
        load_best_model_at_end=True,         # keep the checkpoint with the lowest eval_loss
        metric_for_best_model="eval_loss",
        model_init_kwargs={"dtype": torch.bfloat16 if bf16 else torch.float16},
    ),
)

trainer.train()
trainer.save_model()          # saves the LoRA adapter (tens of MB) to OUTPUT_DIR
print("saved to:", OUTPUT_DIR)

import json
import torch
from datasets import load_dataset
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_DIR = "./models/llama-3.2-3b-instruct"
ADAPTER_DIR = "./outputs/llama-3.2-3b-incident-lora"

tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
base = AutoModelForCausalLM.from_pretrained(MODEL_DIR, dtype=torch.bfloat16, device_map="auto")
model = PeftModel.from_pretrained(base, ADAPTER_DIR)
model.eval()

def generate(messages: list[dict]) -> str:
    inputs = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt", return_dict=True,
    ).to(model.device)
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=300, do_sample=False)
    return tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True).strip()

REQUIRED = ["Likely causes", "Check", "Action"]      # the answer format we want

rows = []
for ex in load_dataset("json", data_files="data/eval.jsonl", split="train"):
    with model.disable_adapter():
        before = generate(ex["prompt"])
    after = generate(ex["prompt"])
    rows.append({
        "question": ex["prompt"][-1]["content"],
        "expected": ex["completion"][0]["content"],
        "before": before,
        "after": after,
        "format_ok_before": all(k in before for k in REQUIRED),
        "format_ok_after": all(k in after for k in REQUIRED),
    })

with open("eval_report.jsonl", "w", encoding="utf-8") as f:
    for r in rows:
        f.write(json.dumps(r, ensure_ascii=False) + "\n")

n = len(rows)
print(f"format compliance: before {sum(r['format_ok_before'] for r in rows)}/{n}, after {sum(r['format_ok_after'] for r in rows)}/{n}")

import torch
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer

ADAPTER_DIR = "./outputs/llama-3.2-3b-incident-lora"
MERGED_DIR = "./outputs/llama-3.2-3b-incident-merged"

# Merge on top of the 16-bit original weights, not the 4-bit ones
model = AutoPeftModelForCausalLM.from_pretrained(ADAPTER_DIR, dtype=torch.bfloat16)
model.merge_and_unload().save_pretrained(MERGED_DIR)
AutoTokenizer.from_pretrained("./models/llama-3.2-3b-instruct").save_pretrained(MERGED_DIR)

import time
import logging
from fastapi import FastAPI, Header, HTTPException
from fastapi.responses import StreamingResponse
from ollama import AsyncClient
from pydantic import BaseModel, Field

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("incident-llama")

app = FastAPI(title="Incident Assistant")
ollama = AsyncClient(host="http://localhost:11434")
API_KEYS = {"team-sre-key"}               # in production, read these from env vars or a secret store

class ChatRequest(BaseModel):
    message: str = Field(min_length=1, max_length=4000)

@app.get("/health")
async def health() -> dict:
    await ollama.list()                    # checks the Ollama server too
    return {"status": "ok"}

@app.post("/chat")
async def chat(req: ChatRequest, x_api_key: str = Header()):
    if x_api_key not in API_KEYS:
        raise HTTPException(status_code=401, detail="invalid api key")

    async def stream():
        started = time.perf_counter()
        first_token_at = None
        async for part in await ollama.chat(
            model="incident-llama",
            messages=[{"role": "user", "content": req.message}],
            stream=True,
        ):
            if first_token_at is None:
                first_token_at = time.perf_counter()
            yield part.message.content
            if part.done:
                log.info(
                    "ttft=%.2fs total=%.2fs prompt_tokens=%s output_tokens=%s",
                    first_token_at - started, time.perf_counter() - started,
                    part.prompt_eval_count, part.eval_count,
                )

    return StreamingResponse(stream(), media_type="text/plain; charset=utf-8")