#!/usr/bin/env python3
"""Anti-epanorthosis adapter, v2 (paper §7.2). Supervised fine-tuning (the delivered adapter); optional light DPO.

The v1 DPO-only run learned to rank the training pairs but did not shift generation
(small-data DPO instability). v2 teaches the affirmative style directly with SFT on the
'chosen' responses (stable with limited data). This SFT-only adapter is the one the paper
reports (§7.2) and ships; it is what the script produces by default. LoRA rank raised to 32.
Set RUN_DPO=1 to additionally run an (experimental) light DPO pass on the preference pairs;
it is written to a separate path ("<OUT>-dpo") and never overwrites the delivered SFT adapter.

Requirements (PyTorch 2.4 / RunPod "PyTorch 2.4" -> pin trl 0.11.4; newer trl needs torch>=2.5):
    pip install "transformers==4.45.2" "trl==0.11.4" "peft==0.13.2" "datasets>=2.20" \
                "accelerate==0.34.2" "bitsandbytes>=0.44" "rich"   # rich: needed by trl 0.11.4 SFTTrainer

Data: artificial-epanorthosis-pairs.jsonl with {"prompt","chosen","rejected"} per line.
"""
import os, torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig, DPOTrainer, DPOConfig

BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-7B-Instruct")
DATA_PATH  = os.environ.get("DATA_PATH", "artificial-epanorthosis-pairs.jsonl")
OUT        = os.environ.get("OUT", "epanorthosis-deemphasis-lora-v2")

quant = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                           bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True)
tok = AutoTokenizer.from_pretrained(BASE_MODEL)
if tok.pad_token is None: tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(BASE_MODEL, quantization_config=quant, device_map="auto")
model.config.use_cache = False

peft_config = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.05, bias="none",
                         task_type="CAUSAL_LM",
                         target_modules=["q_proj","v_proj","gate_proj","up_proj","down_proj"])

data = load_dataset("json", data_files=DATA_PATH, split="train")

# ---------- Phase 1: SFT on the 'chosen' (affirmative) responses ----------
def to_text(ex):
    return {"text": tok.apply_chat_template(
        [{"role":"user","content":ex["prompt"]},{"role":"assistant","content":ex["chosen"]}],
        tokenize=False)}
sft_ds = data.map(to_text, remove_columns=data.column_names)
sft_args = SFTConfig(output_dir=OUT+"-sft", num_train_epochs=3,
                     per_device_train_batch_size=2, gradient_accumulation_steps=8,
                     learning_rate=2e-4, lr_scheduler_type="cosine", warmup_ratio=0.05,
                     logging_steps=10, save_strategy="no", bf16=True,
                     max_seq_length=1024, dataset_text_field="text", report_to="none")
def make_sft():
    try:    return SFTTrainer(model=model, args=sft_args, train_dataset=sft_ds, processing_class=tok, peft_config=peft_config)
    except TypeError: return SFTTrainer(model=model, args=sft_args, train_dataset=sft_ds, tokenizer=tok, peft_config=peft_config)
print(">>> Phase 1: SFT on affirmative responses")
sft = make_sft(); sft.train()
sft.save_model(OUT)                     # usable adapter even if DPO fails
tuned = sft.model
print(">>> SFT done, adapter saved to", OUT)
import sys
if not os.environ.get("RUN_DPO"):       # default: keep the delivered SFT-only adapter (paper §7.2)
    print(">>> Done: SFT-only adapter (the delivered one) saved to", OUT,
          "-- set RUN_DPO=1 to also run the experimental DPO pass."); sys.exit(0)

# ---------- Phase 2 (opt-in, RUN_DPO=1): experimental light DPO, saved separately ----------
DPO_OUT = OUT + "-dpo"                   # never overwrites the delivered SFT adapter
try:
    def fmt(ex):
        ex["prompt"] = tok.apply_chat_template([{"role":"user","content":ex["prompt"]}],
                                               tokenize=False, add_generation_prompt=True)
        return ex
    dpo_ds = data.map(fmt)
    dpo_args = DPOConfig(output_dir=DPO_OUT, beta=0.1, num_train_epochs=1,
                         per_device_train_batch_size=1, gradient_accumulation_steps=8,
                         learning_rate=5e-6, lr_scheduler_type="cosine", warmup_ratio=0.1,
                         logging_steps=10, save_strategy="no", bf16=True,
                         max_length=1024, max_prompt_length=512, report_to="none")
    try:    dpo = DPOTrainer(model=tuned, args=dpo_args, train_dataset=dpo_ds, processing_class=tok)
    except TypeError: dpo = DPOTrainer(model=tuned, args=dpo_args, train_dataset=dpo_ds, tokenizer=tok)
    print(">>> Phase 2: light DPO")
    dpo.train(); dpo.save_model(DPO_OUT)
    print(">>> SFT+DPO done, experimental adapter saved to", DPO_OUT, "(SFT-only adapter left intact at", OUT + ")")
except Exception as e:
    print(">>> DPO phase skipped/failed (keeping the SFT adapter):", repr(e)[:200])