#!/usr/bin/env python3
"""KTO preference tuning on top of the Starburst 2.3 SFT model.

Desirable  = winning entries + correct skips (837)
Undesirable = losing entries (384 real losers, held out of SFT)

KTO nudges the policy away from loss-like setups using per-sample binary labels —
no DPO-style pairs, no fabricated "rejected" data. Reference = the SFT model
(the merged 2.3) with the new LoRA disabled (PEFT auto-ref).

Free the GPU first:  sudo systemctl stop starburst
Run:  PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True unsloth_env/bin/python train_kto.py
"""
import os
import unsloth  # noqa: F401  (import first to patch transformers/trl)
try:
    from unsloth import PatchDPOTrainer
    PatchDPOTrainer()          # patches the DPO/KTO/ORPO trainers (no-op if already patched)
except Exception:
    pass
from unsloth import FastModel
from datasets import load_dataset
from trl import KTOTrainer, KTOConfig

MODEL = os.environ.get("KTO_MODEL", "/home/sheen/trading-assistant/starburst_2.3_9b_merged")
MAXLEN = int(os.environ.get("KTO_MAXLEN", "2048"))
DATA = os.environ.get("KTO_DATA", "train_kto.jsonl")
OUT = os.environ.get("KTO_OUT", "qlora_kto_2_3")
EPOCHS = float(os.environ.get("KTO_EPOCHS", "1"))
LR = float(os.environ.get("KTO_LR", "1e-5"))
BETA = float(os.environ.get("KTO_BETA", "0.1"))

model, tokenizer = FastModel.from_pretrained(
    MODEL, max_seq_length=MAXLEN, load_in_4bit=True, use_gradient_checkpointing="unsloth")
model = FastModel.get_peft_model(
    model, r=16, lora_alpha=16, lora_dropout=0.0, bias="none",
    finetune_vision_layers=False, finetune_language_layers=True,
    finetune_attention_modules=True, finetune_mlp_modules=True,
    use_gradient_checkpointing="unsloth", random_state=3407)
text_tok = getattr(tokenizer, "tokenizer", tokenizer)

# Use the SAME no-seed template the model is served with, so the completion the
# policy is scored on includes its own <think>...</think> (matches SFT + serving).
with open("/home/sheen/trading-assistant/starburst_chat_template.jinja") as f:
    text_tok.chat_template = f.read()

ds = load_dataset("json", data_files=DATA, split="train")

trainer = KTOTrainer(
    model=model,
    ref_model=None,                      # PEFT: reference = base with adapter disabled
    args=KTOConfig(
        per_device_train_batch_size=2,   # KTO needs actual batch > 1 (KL term)
        gradient_accumulation_steps=4,   # effective batch still 8
        num_train_epochs=EPOCHS,
        learning_rate=LR,
        beta=BETA,
        desirable_weight=1.0,
        undesirable_weight=float(os.environ.get("KTO_UNDES_W", "1.8")),  # 2.3: 837:384; per-run override
        max_length=MAXLEN,
        max_prompt_length=1024,
        max_completion_length=1024,
        warmup_ratio=0.05,
        logging_steps=5,
        optim="adamw_8bit",
        weight_decay=0.0,
        lr_scheduler_type="linear",
        seed=3407,
        output_dir=OUT,
        report_to="none",
    ),
    train_dataset=ds,
    processing_class=text_tok,
)
trainer.train()
model.save_pretrained(OUT)
text_tok.save_pretrained(OUT)
print("KTO adapter saved to", OUT)