Fine-Tuning Pipeline

"""Data curation pipeline for fine-tuning datasets."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Iterator

from pydantic import BaseModel, Field, field_validator


class TrainingExample(BaseModel):
    """A single training example in chat format."""

    messages: list[dict[str, str]]
    metadata: dict = Field(default_factory=dict)

    @field_validator("messages")
    @classmethod
    def validate_messages(cls, v):
        roles = [m["role"] for m in v]
        if roles[0] != "system":
            raise ValueError("First message must have role 'system'")
        if roles[-1] != "assistant":
            raise ValueError("Last message must have role 'assistant'")
        return v


def load_and_validate(path: Path) -> Iterator[TrainingExample]:
    """Load JSONL and validate each example."""
    with open(path) as f:
        for i, line in enumerate(f):
            try:
                data = json.loads(line)
                yield TrainingExample(**data)
            except Exception as e:
                print(f"Skipping line {i}: {e}")


def deduplicate(examples: list[TrainingExample]) -> list[TrainingExample]:
    """Remove exact duplicates based on content hash."""
    seen = set()
    unique = []
    for ex in examples:
        content_hash = hashlib.sha256(
            json.dumps(ex.messages, sort_keys=True).encode()
        ).hexdigest()
        if content_hash not in seen:
            seen.add(content_hash)
            unique.append(ex)
    return unique


def split_dataset(
    examples: list[TrainingExample],
    train_ratio: float = 0.85,
    val_ratio: float = 0.10,
) -> dict[str, list[TrainingExample]]:
    """Split into train/val/test with deterministic shuffle."""
    import random
    rng = random.Random(42)
    shuffled = list(examples)
    rng.shuffle(shuffled)

    n = len(shuffled)
    train_end = int(n * train_ratio)
    val_end = train_end + int(n * val_ratio)

    return {
        "train": shuffled[:train_end],
        "validation": shuffled[train_end:val_end],
        "test": shuffled[val_end:],
    }

"""LoRA/QLoRA fine-tuning with Hugging Face PEFT and TRL."""

from __future__ import annotations

import torch
from datasets import load_dataset
from peft import LoraConfig, TaskType, get_peft_model
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
)
from trl import SFTTrainer


def train(
    base_model: str = "meta-llama/Llama-3.1-8B-Instruct",
    dataset_path: str = "./data/train.jsonl",
    output_dir: str = "./checkpoints",
    use_qlora: bool = True,
    lora_r: int = 16,
    lora_alpha: int = 32,
    epochs: int = 3,
    learning_rate: float = 2e-4,
    batch_size: int = 4,
    max_seq_length: int = 2048,
):
    """Run LoRA or QLoRA fine-tuning.

    Args:
        base_model: HuggingFace model ID for the base model.
        dataset_path: Path to JSONL training data.
        output_dir: Where to save checkpoints.
        use_qlora: If True, quantize base model to 4-bit.
        lora_r: LoRA rank (higher = more parameters).
        lora_alpha: LoRA scaling factor.
    """
    # Quantization config for QLoRA
    quant_config = None
    if use_qlora:
        quant_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )

    # Load base model
    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=quant_config,
        device_map="auto",
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
    )

    tokenizer = AutoTokenizer.from_pretrained(base_model)
    tokenizer.pad_token = tokenizer.eos_token

    # LoRA config — target the attention layers
    lora_config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_alpha,
        lora_dropout=0.05,
        target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
        task_type=TaskType.CAUSAL_LM,
        bias="none",
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()  # Typically 0.5-2% of total

    # Load dataset
    dataset = load_dataset("json", data_files=dataset_path, split="train")

    # Training arguments
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=epochs,
        per_device_train_batch_size=batch_size,
        gradient_accumulation_steps=4,
        learning_rate=learning_rate,
        lr_scheduler_type="cosine",
        warmup_ratio=0.1,
        bf16=True,
        logging_steps=10,
        save_strategy="epoch",
        evaluation_strategy="epoch",
        report_to="wandb",
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset,
        tokenizer=tokenizer,
        max_seq_length=max_seq_length,
    )

    trainer.train()
    trainer.save_model(f"{output_dir}/final")