"""QLoRA fine-tune of Llama-3.2-1B-Instruct → "Jarvis" for BioPatch. Run on a GPU (Colab T4 free, RunPod, HF Spaces with GPU…). On a 2012 MacBook Pro this script is not runnable — you train remotely and download the resulting LoRA adapter. Pipeline: 1. Load ``meta-llama/Llama-3.2-1B-Instruct`` in 4-bit (bitsandbytes). 2. Attach a LoRA adapter (PEFT) on the attention + MLP projections. 3. Concatenate the three BioPatch datasets (identity + medical + tools), format with the Llama-3 chat template. 4. Train with the standard supervised loss (causal LM masking on the assistant turn only). 5. Save the LoRA weights under ``training/output/jarvis-lora/``. Then run ``merge_and_export.py`` to merge the adapter into the base weights and convert to GGUF for ``llama-cpp-python``. Hyperparameters tuned for one T4 (~15 GB VRAM) and ~860 samples — a single epoch is roughly 6 minutes. 3 epochs is the sweet spot before the identity dataset starts to memorize verbatim. Usage: python -m training.lora_jarvis --output-dir training/output/jarvis-lora --epochs 3 """ from __future__ import annotations import argparse import json import os from pathlib import Path from typing import Iterable HERE = Path(__file__).resolve().parent def main() -> int: parser = argparse.ArgumentParser() parser.add_argument( "--base-model", # Default to the model BioPatch actually ships in production so # the LoRA we produce is plug-compatible with the HF Space. # Override with --base-model meta-llama/Llama-3.2-1B-Instruct for # the lightweight Colab-T4-friendly path. default="BioMistral/BioMistral-7B", help="Base model on HF Hub.", ) # Default datasets: identity + medical + tools, plus continual if # the operator has exported it from the live DB. Continual being # optional means we never fail on a fresh install. _default_datasets = [ str(HERE / "datasets" / "identity.jsonl"), str(HERE / "datasets" / "medical.jsonl"), str(HERE / "datasets" / "tools.jsonl"), ] _continual = HERE / "datasets" / "continual.jsonl" if _continual.exists(): _default_datasets.append(str(_continual)) _medical_pt = HERE / "datasets" / "medical_pt_br.jsonl" if _medical_pt.exists(): _default_datasets.append(str(_medical_pt)) # Real labelled ICU sepsis reasoning (PhysioNet 2019) if the operator # built it — teaches the local-wound → systemic-sepsis escalation path. _physionet = HERE / "datasets" / "physionet.jsonl" if _physionet.exists(): _default_datasets.append(str(_physionet)) parser.add_argument( "--datasets", nargs="+", default=_default_datasets, ) parser.add_argument( "--output-dir", type=Path, default=HERE / "output" / "jarvis-lora", ) parser.add_argument("--epochs", type=int, default=3) parser.add_argument("--lr", type=float, default=2e-4) parser.add_argument("--batch-size", type=int, default=4) parser.add_argument("--grad-accum", type=int, default=2) parser.add_argument("--max-seq-length", type=int, default=2048) parser.add_argument("--lora-r", type=int, default=16) parser.add_argument("--lora-alpha", type=int, default=32) parser.add_argument("--lora-dropout", type=float, default=0.05) parser.add_argument("--seed", type=int, default=42) parser.add_argument( "--push-to-hub", default=None, help="Optional HF Hub repo id to push the LoRA to after training.", ) args = parser.parse_args() # --- Heavy deps imported lazily so the script lives in the repo even # on machines that can't install transformers / torch. import torch from datasets import Dataset, concatenate_datasets from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, DataCollatorForLanguageModeling, Trainer, TrainingArguments, set_seed, ) set_seed(args.seed) # --- Load tokenizer + base model in 4-bit --------------------------- tokenizer = AutoTokenizer.from_pretrained(args.base_model, use_fast=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token 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( args.base_model, quantization_config=bnb, device_map="auto", torch_dtype=torch.bfloat16, ) model = prepare_model_for_kbit_training(model) lora_config = LoraConfig( r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=args.lora_dropout, bias="none", task_type="CAUSAL_LM", target_modules=[ "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ], ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() # --- Build the dataset --------------------------------------------- def read_jsonl(path: str) -> Iterable[dict]: with open(path, "r", encoding="utf-8") as fh: for line in fh: line = line.strip() if line: yield json.loads(line) raw = [] for path in args.datasets: raw.extend(read_jsonl(path)) print(f"Loaded {len(raw)} training samples") def format_example(example: dict) -> dict: """Build training tokens with **assistant-only loss masking**. We mask everything that isn't the assistant turn so the LoRA only learns "what Jarvis SHOULD SAY", never "what the user typed". Without this, identity Q/A pairs like ``user: "qui es-tu ?"`` would teach the model to *produce* the question, diluting the gradient that bakes the answer. Mask = -100 is the HuggingFace convention to skip a position in the loss. """ # Build the full conversation text + a prefix-only version # (system + user, with assistant generation prompt). The length # of the prefix is exactly where the assistant turn starts. msgs = example["messages"] full_text = tokenizer.apply_chat_template( msgs, tokenize=False, add_generation_prompt=False ) prefix_text = tokenizer.apply_chat_template( msgs[:-1], tokenize=False, add_generation_prompt=True ) full = tokenizer( full_text, truncation=True, max_length=args.max_seq_length, padding=False, ) prefix_len = len( tokenizer( prefix_text, truncation=True, max_length=args.max_seq_length, padding=False, )["input_ids"] ) labels = list(full["input_ids"]) # Mask the entire prefix → only the assistant turn contributes # to the loss. -100 is HuggingFace's "ignore" sentinel. for i in range(min(prefix_len, len(labels))): labels[i] = -100 full["labels"] = labels return full hf_ds = Dataset.from_list(raw).map(format_example, remove_columns=["messages"]) # --- Training ------------------------------------------------------- training_args = TrainingArguments( output_dir=str(args.output_dir), num_train_epochs=args.epochs, learning_rate=args.lr, per_device_train_batch_size=args.batch_size, gradient_accumulation_steps=args.grad_accum, warmup_ratio=0.03, lr_scheduler_type="cosine", logging_steps=10, save_steps=200, save_total_limit=2, bf16=True, report_to="none", push_to_hub=args.push_to_hub is not None, hub_model_id=args.push_to_hub, seed=args.seed, remove_unused_columns=False, ) collator = DataCollatorForLanguageModeling(tokenizer, mlm=False) trainer = Trainer( model=model, args=training_args, train_dataset=hf_ds, tokenizer=tokenizer, data_collator=collator, ) trainer.train() trainer.save_model(str(args.output_dir)) tokenizer.save_pretrained(str(args.output_dir)) print(f"Saved LoRA adapter to {args.output_dir}") if args.push_to_hub: trainer.push_to_hub() return 0 if __name__ == "__main__": raise SystemExit(main())