bikoret/kushya_generator/fine_tune/train.py

"""
QLoRA fine-tuning script for the kushya generator.

Architecture §3.6.

Requires GPU and the heavy ML stack:
  pip install torch transformers peft bitsandbytes datasets accelerate trl pyyaml

Run:
  python train.py --config config.yaml

This script is the operator's task — it requires GPU compute we don't have
in the build environment. The script is complete and self-contained; when
the operator has the hardware, it runs.

The output is a LoRA adapter (small file, ~100MB) that wraps the base model.
Quantize and serve via inference/serve.py for the Talmudic Engine to call.
"""

import argparse
import json
import os
import sys
from pathlib import Path

# Defer heavy imports until execution so the file can be parsed and
# linted in environments without torch/transformers installed.


def load_config(path: Path) -> dict:
    import yaml
    with open(path) as f:
        return yaml.safe_load(f)


def load_jsonl(path: Path) -> list[dict]:
    records = []
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records


def format_records(records: list[dict], template: str) -> list[str]:
    return [
        template.format(input=r["input"], output=r["output"])
        for r in records
    ]


def main():
    parser = argparse.ArgumentParser(description="QLoRA fine-tune the kushya generator")
    parser.add_argument("--config", type=Path, default=Path(__file__).parent / "config.yaml")
    args = parser.parse_args()

    if not args.config.exists():
        print(f"ERROR: config not found at {args.config}", file=sys.stderr)
        sys.exit(1)

    config = load_config(args.config)

    print("=" * 60)
    print("KUSHYA GENERATOR — QLORA FINE-TUNING")
    print("=" * 60)
    print(f"Base model:  {config['base_model']['name']}")
    print(f"LoRA r/α:    {config['lora']['r']} / {config['lora']['alpha']}")
    print(f"Epochs:      {config['training']['num_train_epochs']}")
    print(f"LR:          {config['training']['learning_rate']}")
    print()

    # Resolve data paths relative to the config file.
    base_dir = args.config.parent
    train_path = (base_dir / config["data"]["train_file"]).resolve()
    eval_path = (base_dir / config["data"]["eval_file"]).resolve()

    if not train_path.exists():
        print(f"ERROR: train file not found at {train_path}", file=sys.stderr)
        print(f"Run: python ../data/prepare_training.py", file=sys.stderr)
        sys.exit(1)

    train_records = load_jsonl(train_path)
    eval_records = load_jsonl(eval_path) if eval_path.exists() else []

    print(f"Train records: {len(train_records)}")
    print(f"Eval records:  {len(eval_records)}")

    if len(train_records) < 100:
        print()
        print(f"WARNING: only {len(train_records)} training records.")
        print("Architecture §3.6 expects training_pairs/ to contain extracted")
        print("kushyot from Bavli, Yerushalmi, and Rishonim. The seed pairs")
        print("are sufficient to validate the training pipeline but not to")
        print("produce a useful kushya generator. Curate more pairs first.")
        print()

    template = config["data"]["prompt_template"]
    train_texts = format_records(train_records, template)
    eval_texts = format_records(eval_records, template) if eval_records else []

    # Heavy imports only at execution time.
    try:
        import torch
        from datasets import Dataset
        from transformers import (
            AutoModelForCausalLM,
            AutoTokenizer,
            BitsAndBytesConfig,
            TrainingArguments,
        )
        from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
        from trl import SFTTrainer
    except ImportError as e:
        print(f"ERROR: required ML library not installed: {e}", file=sys.stderr)
        print()
        print("Install with:")
        print("  pip install torch transformers peft bitsandbytes datasets accelerate trl pyyaml")
        sys.exit(1)

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=getattr(torch, config["base_model"]["bnb_4bit_compute_dtype"]),
        bnb_4bit_quant_type=config["base_model"]["bnb_4bit_quant_type"],
        bnb_4bit_use_double_quant=config["base_model"]["bnb_4bit_use_double_quant"],
    )

    print(f"Loading base model {config['base_model']['name']}...")
    tokenizer = AutoTokenizer.from_pretrained(
        config["base_model"]["name"],
        trust_remote_code=config["base_model"]["trust_remote_code"],
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        config["base_model"]["name"],
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=config["base_model"]["trust_remote_code"],
    )
    model = prepare_model_for_kbit_training(model)

    lora_config = LoraConfig(
        r=config["lora"]["r"],
        lora_alpha=config["lora"]["alpha"],
        lora_dropout=config["lora"]["dropout"],
        bias=config["lora"]["bias"],
        task_type=config["lora"]["task_type"],
        target_modules=config["lora"]["target_modules"],
    )
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    train_ds = Dataset.from_dict({"text": train_texts})
    eval_ds = Dataset.from_dict({"text": eval_texts}) if eval_texts else None

    training_args = TrainingArguments(
        output_dir=config["training"]["output_dir"],
        num_train_epochs=config["training"]["num_train_epochs"],
        per_device_train_batch_size=config["training"]["per_device_train_batch_size"],
        per_device_eval_batch_size=config["training"]["per_device_eval_batch_size"],
        gradient_accumulation_steps=config["training"]["gradient_accumulation_steps"],
        gradient_checkpointing=config["training"]["gradient_checkpointing"],
        optim=config["training"]["optim"],
        learning_rate=config["training"]["learning_rate"],
        lr_scheduler_type=config["training"]["lr_scheduler_type"],
        warmup_ratio=config["training"]["warmup_ratio"],
        weight_decay=config["training"]["weight_decay"],
        max_grad_norm=config["training"]["max_grad_norm"],
        fp16=config["training"]["fp16"],
        bf16=config["training"]["bf16"],
        logging_steps=config["training"]["logging_steps"],
        eval_strategy=config["training"]["evaluation_strategy"] if eval_ds else "no",
        eval_steps=config["training"]["eval_steps"] if eval_ds else None,
        save_strategy=config["training"]["save_strategy"],
        save_total_limit=config["training"]["save_total_limit"],
        group_by_length=config["training"]["group_by_length"],
        report_to=config["training"]["report_to"],
        seed=config["training"]["seed"],
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_ds,
        eval_dataset=eval_ds,
        tokenizer=tokenizer,
        max_seq_length=config["data"]["max_seq_length"],
        dataset_text_field="text",
    )

    trainer.train()
    trainer.save_model(config["training"]["output_dir"])
    print()
    print(f"Training complete. Adapter saved to {config['training']['output_dir']}")
    print()
    print("Next: quantize to GGUF and serve via inference/serve.py")


if __name__ == "__main__":
    main()