#!/usr/bin/env python3
"""Fine-tune a small open model on the Doc Quatrain corpus. QLoRA, one GPU.

    pip install -r requirements.txt
    python3 build_corpus.py          # writes train.jsonl and holdout.jsonl
    python3 train.py                 # writes ./out/adapter

    BASE_MODEL=Qwen/Qwen3-1.7B python3 train.py    # the smaller track

Roughly 11 GB of VRAM on the 8B model and 6 GB on the 1.7B, and a few minutes either way.
The base weights are frozen and quantized to 4 bits; only a small adapter is trained, and
that adapter is what lands in ./out/adapter.

Part nine of the series walks through every setting. The four that matter are marked MATTERS
below, and they trade against each other rather than being independent dials.
"""
import os
from pathlib import Path

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

HERE = Path(__file__).resolve().parent
BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-8B")
DATA = os.environ.get("DATA", str(HERE / "train.jsonl"))
OUT = os.environ.get("OUT", str(HERE / "out" / "adapter"))

# Pin the run to one card. On a multi-GPU box this is what stops a mistyped setting starting
# a distributed run, which does not fail, it just takes twenty times longer.
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "0")

if not torch.cuda.is_available():
    raise SystemExit(
        "No CUDA device visible. This will not train usefully on a CPU: it will appear to "
        "work and take days. Check the driver, then `python3 -c \"import torch; "
        "print(torch.cuda.is_available())\"`."
    )

print(f"base model : {BASE_MODEL}")
print(f"data       : {DATA}")
print(f"device     : {torch.cuda.get_device_name(0)}")

# QLoRA: the frozen base is stored at 4 bits, which is what brings an 8B model onto a
# 24 GB card with room for the activations. The adapter itself stays at full precision,
# because it is the part being trained.
quant = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
if tokenizer.chat_template is None:
    raise SystemExit(
        f"{BASE_MODEL} ships no chat template, so the turn markers the model was trained "
        "to recognize cannot be applied. Part three explains why that matters."
    )

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    quantization_config=quant,
    dtype=torch.bfloat16,
    device_map={"": 0},
)
model.config.use_cache = False

# Low-rank adaptation. Instead of updating all 8 billion parameters, freeze them and learn
# two small matrices beside each weight table. r=16 trains well under 1% of the model.
peft_config = LoraConfig(
    r=16,                       # MATTERS: adapter capacity. Higher memorizes sooner.
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
    target_modules="all-linear",
)

dataset = load_dataset("json", data_files=DATA, split="train")
print(f"examples   : {len(dataset)}")

config = SFTConfig(
    output_dir=OUT,
    num_train_epochs=15,        # MATTERS: passes over very few examples.
    learning_rate=1e-4,         # MATTERS: step size. Roughly 10x a full fine-tune's.
    save_strategy="epoch",      # MATTERS: keep every pass, so you can go back to the best.
    lr_scheduler_type="cosine",
    warmup_ratio=0.1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    max_length=1024,            # Above the longest example. Below it, pieces get truncated.
    gradient_checkpointing=True,
    bf16=True,
    optim="paged_adamw_8bit",
    logging_steps=1,
    report_to="none",
    # Train on the response only, so the model does not learn to generate prompts as well
    # as answers. This relies on the tokenizer's chat template marking the assistant turn;
    # if your version of TRL or the template rejects it, drop this line. Loss then covers
    # the whole conversation, which at this corpus size is a small effect but not nothing.
    assistant_only_loss=True,
)

trainer = SFTTrainer(
    model=model,
    args=config,
    train_dataset=dataset,
    peft_config=peft_config,
    processing_class=tokenizer,
)

result = trainer.train()
trainer.save_model(OUT)
tokenizer.save_pretrained(OUT)

first = trainer.state.log_history[0].get("loss") if trainer.state.log_history else None
print(f"\nadapter    : {OUT}")
print(f"loss       : {first} -> {result.training_loss:.4f}")
print(
    "\nIf those two numbers are the same to several decimal places, nothing trained: the "
    "adapter attached to no layers. That failure writes checkpoints and reports success."
)