#!/usr/bin/env python3
"""
DOJO Phase 2 Supplement Top-Up — runs inside HuggingFace container.
Loads naima-dojo-741hz-v3 weights, trains on dojo_gold_v4_supplement.jsonl,
pushes to naima-dojo-741hz-v4.
"""
import os, torch
from datasets import load_dataset
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer

CHAMBER   = "DOJO"
HZ        = 741
BASE      = "misterJB/naima-dojo-741hz-v3"     # load v3 weights as starting point
DATASET   = "misterJB/dojo-manifestation-training"
SUPPLEMENT = "dojo_gold_v4_supplement.jsonl"
OUTPUT    = "misterJB/naima-dojo-741hz-v4"
CKPT_DIR  = "/tmp/dojo-741hz-v4-ckpt"

print(f"{CHAMBER} {HZ}Hz -- Supplement Top-Up START")
gpu = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "None"
print(f"GPU: {gpu}")
print(f"Base: {BASE}")
print(f"Supplement: {SUPPLEMENT}")

tokenizer = AutoTokenizer.from_pretrained(
    BASE,
    trust_remote_code=True,
    token=os.environ["HF_TOKEN"],
)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

cfg = AutoConfig.from_pretrained(BASE, trust_remote_code=True, token=os.environ["HF_TOKEN"])

model = AutoModelForCausalLM.from_pretrained(
    BASE,
    config=cfg,
    quantization_config=None,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    trust_remote_code=True,
    token=os.environ["HF_TOKEN"],
)
model.config.use_cache = False
model.gradient_checkpointing_enable()
print(f"Model loaded — {sum(p.numel() for p in model.parameters()) / 1e9:.1f}B parameters")

# Load supplement corpus
ds = load_dataset(DATASET, data_files={"train": SUPPLEMENT}, split="train",
                  token=os.environ["HF_TOKEN"])
print(f"Supplement corpus: {len(ds)} examples")

# 3 epochs — ~1728 examples seen, ~108 optimizer steps per epoch, ~324 total
args = SFTConfig(
    output_dir=CKPT_DIR,
    num_train_epochs=3,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=16,        # effective batch 16
    learning_rate=5e-6,                    # half of v3 LR — prevents catastrophic forgetting
    warmup_ratio=0.1,
    lr_scheduler_type="cosine",
    weight_decay=0.01,
    bf16=True,
    gradient_checkpointing=True,
    save_strategy="epoch",
    save_total_limit=1,
    logging_steps=25,
    push_to_hub=True,
    hub_model_id=OUTPUT,
    hub_token=os.environ["HF_TOKEN"],
    hub_strategy="end",
    report_to="none",
    max_length=1024,
)

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=ds,
    processing_class=tokenizer,
)

trainer.train()
trainer.push_to_hub(
    commit_message=f"DOJO 741Hz supplement top-up v4 — valformat gate binding, Naima identity, 3 epochs"
)
print(f"✅ DOJO v4 pushed to {OUTPUT}")