train_lora.py

#!/usr/bin/env python3
"""QLoRA fine-tune of Gemma 2 9B-it as a malevolent/benign classifier.

Run on a CUDA GPU (~24 GB is comfortable for 9B QLoRA). Data is the HF chat-format file
produced by ../ft_dataset/build_ft_jsonl.py (Format B: instruction folded into the user turn,
which suits Gemma 2 since it has no system role).

  huggingface-cli login          # and accept the Gemma license on the model page
  python train_lora.py
"""
import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig
from data_utils import to_messages   # shared with eval_testset.py (single source of truth)

MODEL_ID = "google/gemma-2-9b-it"          # text-only, plain CausalLM, mature GGUF/Ollama support.
# (Gemma 4 12B: GGUF conversion not ready. Gemma 3 12B: multimodal, fights AutoModelForCausalLM.
#  Gemma 2 9B avoids both and is plenty for a binary classifier. Swap to Qwen2.5-14B-Instruct if you prefer.)
# Reuse the already-generated, committed datasets (produced once by ../ft_dataset/build_ft_jsonl.py).
# This script only CONSUMES these files — it does not generate or regenerate any data.
TRAIN_FILE = "bcbsma_injection_ft_train.jsonl"   # 800 examples (400/400)
TEST_FILE  = "bcbsma_injection_ft_test.jsonl"    # 100/100 held-out (also used by eval_testset.py)
OUT_DIR = "out"

def main():
    # Reuse both committed files: train on TRAIN_FILE, monitor eval loss on the held-out TEST_FILE.
    ds = load_dataset("json", data_files={"train": TRAIN_FILE, "test": TEST_FILE})  # rows have {"messages": [...]}

    tok = AutoTokenizer.from_pretrained(MODEL_ID)

    # Normalize either dataset format (HF "messages" or Vertex "contents") to one "text"
    # column via the chat template. to_messages() is shared with eval_testset.py (data_utils.py).
    def to_text(ex):
        return {"text": tok.apply_chat_template(to_messages(ex), tokenize=False)}
    ds = ds.map(to_text, remove_columns=ds["train"].column_names)

    bnb = BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True,
    )
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID, quantization_config=bnb, torch_dtype=torch.bfloat16,
        device_map="auto", attn_implementation="eager",   # eager is safest for Gemma 2 soft-capping
    )

    lora = LoraConfig(
        r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    )

    cfg = SFTConfig(
        output_dir=OUT_DIR, num_train_epochs=3,
        per_device_train_batch_size=2, gradient_accumulation_steps=8,   # small batch + accum; 9B QLoRA fits ~16GB+
        learning_rate=2e-4, lr_scheduler_type="cosine", warmup_ratio=0.03,
        logging_steps=10, save_strategy="epoch", bf16=True,
        eval_strategy="epoch", per_device_eval_batch_size=4,   # eval loss on the held-out test set each epoch
        max_length=512, packing=False, report_to="none",   # SFTConfig renamed max_seq_length -> max_length
        dataset_text_field="text",   # we pre-rendered the chat into a "text" column above
    )

    trainer = SFTTrainer(model=model, train_dataset=ds["train"], eval_dataset=ds["test"],
                         peft_config=lora, args=cfg, processing_class=tok)
    trainer.train()
    trainer.save_model(f"{OUT_DIR}/adapter")
    print(f"Saved LoRA adapter to {OUT_DIR}/adapter")

if __name__ == "__main__":
    main()