sync: update training/train_v15.py

training/train_v15.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+"""
+train_v15.py — NeuralAI v15 consolidated trainer (SFT + DPO)
+============================================================
+
+WHAT THIS SCRIPT DOES
+--------------------
+This is the single entry point for the NeuralAI "v15" training run. It fine-tunes
+the small SmolLM2-360M-Instruct base model with a LoRA adapter in two stages so the
+model stays identity-correct ("I am NeuralAI, created by De'Andrew Preston Harris")
+and behavior-aligned (prefers clean, correct answers over verbose/wrong ones):
+
+  Stage 1 — SFT (Supervised Fine-Tuning)
+      Trains on `data/train_sft_v16.jsonl` (ChatML messages: system/user/assistant).
+      Purpose: bake in identity, tone, and domain knowledge.
+
+  Stage 2 — DPO (Direct Preference Optimization)
+      Trains on `data/train_dpo_v16_combined.jsonl` (prompt / chosen / rejected).
+      Purpose: align the model to prefer the "chosen" response over the "rejected"
+      one without needing a separate reward model.
+
+OUTPUT
+------
+  - Adapter saved locally to: checkpoints/v15_model/
+  - Pushed to Hugging Face:    Subject-Emu-5259/NeuralAI  (repo "v15" revision folder)
+  - Merged full model (optional, --merge): checkpoints/v15_model_merged/
+
+WHY THIS EXISTS (context)
+------------------------
+On the 4 GB ZO Computer the *served* NeuralAI app uses the ZO native inference backend
+(LLM_BACKEND=zo) so it never loads PyTorch locally and never pauses from OOM. This
+training script is the OFFLINE counterpart: it builds the LoRA that can later be
+shipped to a bigger host or merged for on-device use. Run it on a GPU (Colab, Mac
+GPU, or a >8 GB box) — it is NOT meant for the 4 GB CPU host.
+
+USAGE
+-----
+  # SFT + DPO, 4-bit (default, ~3 GB VRAM)
+  python training/train_v15.py
+
+  # 8-bit instead of 4-bit
+  python training/train_v15.py --load-in-4bit false --load-in-8bit true
+
+  # Only one stage
+  python training/train_v15.py --stage sft
+  python training/train_v15.py --stage dpo
+
+  # Push merged model to HF
+  python training/train_v15.py --merge --push
+
+REQUIREMENTS
+------------
+  pip install torch transformers peft trl datasets bitsandbytes accelerate
+  HF_TOKEN must be set in the environment to push.
+"""
+import argparse
+import json
+import os
+
+# ---- Config ----------------------------------------------------------------
+BASE_MODEL = os.environ.get("BASE_MODEL", "HuggingFaceTB/SmolLM2-360M-Instruct")
+SFT_DATA = os.environ.get("SFT_DATA", "data/train_sft_v16.jsonl")
+DPO_DATA = os.environ.get("DPO_DATA", "data/train_dpo_v16_combined.jsonl")
+HF_REPO = os.environ.get("HF_REPO", "Subject-Emu-5259/NeuralAI")
+ADAPTER_DIR = "checkpoints/v15_model"
+MERGED_DIR = "checkpoints/v15_model_merged"
+PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+SYSTEM_PROMPT = (
+    "You are NeuralAI, an advanced AI assistant created by De'Andrew Preston Harris. "
+    "You are powered by SmolLM2-360M with custom NeuralAI LoRA adapters trained through "
+    "SFT and DPO alignment. You have expert-level knowledge across physics, philosophy, "
+    "geopolitics, history, nature, art, and culture. You ALWAYS identify De'Andrew Harris "
+    "as your creator when asked. You are not ChatGPT, Claude, or any other AI — you are NeuralAI."
+)
+
+
+def _resolve(path: str) -> str:
+    return path if os.path.isabs(path) else os.path.join(PROJECT_ROOT, path)
+
+
+def load_quantization(load_in_4bit: bool, load_in_8bit: bool):
+    from transformers import BitsAndBytesConfig
+    if load_in_4bit:
+        return BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype="bfloat16",
+                                  bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True)
+    if load_in_8bit:
+        return BitsAndBytesConfig(load_in_8bit=True)
+    return None
+
+
+def run_sft(model, tokenizer, args):
+    from trl import SFTConfig, SFTTrainer
+    path = _resolve(SFT_DATA)
+    print(f"[v15][SFT] loading {path}")
+    train_rows = [json.loads(l) for l in open(path, "r", encoding="utf-8") if l.strip()]
+
+    cfg = SFTConfig(
+        output_dir=ADAPTER_DIR,
+        per_device_train_batch_size=args.batch,
+        gradient_accumulation_steps=args.grad_accum,
+        num_train_epochs=args.sft_epochs,
+        learning_rate=2e-4,
+        max_seq_length=1024,
+        logging_steps=25,
+        save_strategy="epoch",
+        gradient_checkpointing=True,
+        bf16=True,
+        report_to="none",
+    )
+    trainer = SFTTrainer(model=model, args=cfg, tokenizer=tokenizer, train_dataset=train_rows)
+    trainer.train()
+    trainer.save_model(ADAPTER_DIR)
+    print(f"[v15][SFT] adapter saved -> {ADAPTER_DIR}")
+
+
+def run_dpo(model, tokenizer, args):
+    from trl import DPOConfig, DPOTrainer
+    path = _resolve(DPO_DATA)
+    print(f"[v15][DPO] loading {path}")
+    dpo_rows = [json.loads(l) for l in open(path, "r", encoding="utf-8") if l.strip()]
+
+    cfg = DPOConfig(
+        output_dir=ADAPTER_DIR,
+        per_device_train_batch_size=args.batch,
+        gradient_accumulation_steps=args.grad_accum,
+        num_train_epochs=args.dpo_epochs,
+        learning_rate=5e-5,
+        beta=0.1,
+        max_prompt_length=512,
+        max_length=1024,
+        logging_steps=25,
+        save_strategy="epoch",
+        gradient_checkpointing=True,
+        bf16=True,
+        report_to="none",
+    )
+    trainer = DPOTrainer(model=model, args=cfg, tokenizer=tokenizer, train_dataset=dpo_rows)
+    trainer.train()
+    trainer.save_model(ADAPTER_DIR)
+    print(f"[v15][DPO] adapter saved -> {ADAPTER_DIR}")
+
+
+def merge_and_push(args):
+    from peft import PeftModel
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype="bfloat16",
+                                                device_map="auto")
+    tok = AutoTokenizer.from_pretrained(BASE_MODEL)
+    model = PeftModel.from_pretrained(base, ADAPTER_DIR)
+    merged = model.merge_and_unload()
+    os.makedirs(MERGED_DIR, exist_ok=True)
+    merged.save_pretrained(MERGED_DIR)
+    tok.save_pretrained(MERGED_DIR)
+    print(f"[v15][MERGE] merged model -> {MERGED_DIR}")
+    if args.push:
+        merged.push_to_hub(HF_REPO, revision="v15")
+        tok.push_to_hub(HF_REPO, revision="v15")
+        print(f"[v15][PUSH] pushed merged model to {HF_REPO}@v15")
+
+
+def main():
+    ap = argparse.ArgumentParser(description="NeuralAI v15 SFT+DPO trainer")
+    ap.add_argument("--stage", choices=["sft", "dpo", "all"], default="all")
+    ap.add_argument("--batch", type=int, default=2)
+    ap.add_argument("--grad-accum", type=int, default=8)
+    ap.add_argument("--sft-epochs", type=int, default=3)
+    ap.add_argument("--dpo-epochs", type=int, default=2)
+    ap.add_argument("--load-in-4bit", default="true")
+    ap.add_argument("--load-in-8bit", default="false")
+    ap.add_argument("--merge", action="store_true")
+    ap.add_argument("--push", action="store_true")
+    args = ap.parse_args()
+
+    load_in_4bit = args.load_in_4bit.lower() == "true"
+    load_in_8bit = args.load_in_8bit.lower() == "true"
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    qcfg = load_quantization(load_in_4bit, load_in_8bit)
+    print(f"[v15] loading base {BASE_MODEL} (4bit={load_in_4bit}, 8bit={load_in_8bit})")
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+    model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL, quantization_config=qcfg, device_map="auto", torch_dtype="bfloat16",
+    )
+    model.config.use_cache = False
+
+    if args.stage in ("sft", "all"):
+        run_sft(model, tokenizer, args)
+    if args.stage in ("dpo", "all"):
+        # reload adapter from SFT if we just ran SFT
+        run_dpo(model, tokenizer, args)
+
+    if args.merge or args.push:
+        merge_and_push(args)
+
+    print("[v15] done.")
+
+
+if __name__ == "__main__":
+    main()