""" 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()