Spaces:

Shwy1001
/

BioPatch

Sleeping

BioPatch / backend /training /lora_jarvis.py

Jerry

Hardware bridge: ESP32 device ingest + hardware mode

34cb656 3 months ago

8.75 kB

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

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