#!/usr/bin/env python3
"""
Reproducible recipe for an "anti-epanorthosis" style adapter (paper §7.2).

A LoRA adapter is trained with Direct Preference Optimisation (DPO) to prefer
affirmative phrasing over gratuitous negate-and-upgrade epanorthosis, at fixed
content. The base model's weights are frozen; only the low-rank update is learned,
so the result is a small, detachable artefact that can be scaled at inference time
(the "rhetoric dial") and combined with the genre-relative evaluation of §7.7.

This is a training RECIPE, not a claim of trained results: run it to reproduce the
adapter. It is deliberately model-agnostic; set BASE_MODEL to any open instruct model.

Requirements:
    # On PyTorch 2.4 (e.g. RunPod "PyTorch 2.4" template) pin trl==0.11.4:
    #   newer trl needs FSDP2 (FSDPModule), which requires 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"   # bitsandbytes only for 4-bit
    # With torch >= 2.5 you may instead use the latest trl.

Preference data (JSONL, one object per line), built by the counterfactual
augmentation of §7.1 (mine emphatic "Not X. Y" spans, pair each with a
content-matched de-emphasised paraphrase; keep only pairs whose embedding
similarity is high, so the intervention changes style and not meaning):

    {"prompt": "<the writing task>",
     "chosen": "<affirmative, de-emphasised response>",
     "rejected": "<same content, emphatic epanorthosis>"}
"""
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import LoraConfig
from trl import DPOTrainer, DPOConfig

import os
BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen2.5-7B-Instruct")  # ungated default; override via env
DATA_PATH  = os.environ.get("DATA_PATH", "artificial-epanorthosis-pairs.jsonl")   # see docstring for format
OUT_DIR    = "epanorthosis-deemphasis-lora"
USE_4BIT   = True                                       # QLoRA: single-GPU friendly

# --- base model (frozen), optionally 4-bit quantised (QLoRA) ---
quant = BitsAndBytesConfig(
    load_in_4bit=True, bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True,
) if USE_4BIT else None

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, quantization_config=quant,
    torch_dtype=torch.bfloat16, device_map="auto",
)
model.config.use_cache = False

# --- LoRA: low-rank update on attention (q,v) and MLP projections ---
peft_config = LoraConfig(
    r=16, lora_alpha=32, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    target_modules=["q_proj", "v_proj", "gate_proj", "up_proj", "down_proj"],
)

# --- preference dataset (format the prompt with the model's chat template) ---
dataset = load_dataset("json", data_files=DATA_PATH, split="train")
def _fmt(ex):
    ex["prompt"] = tokenizer.apply_chat_template(
        [{"role": "user", "content": ex["prompt"]}], tokenize=False, add_generation_prompt=True)
    return ex
dataset = dataset.map(_fmt)

# --- DPO configuration ---
config = DPOConfig(
    output_dir=OUT_DIR,
    beta=0.1,                       # KL strength: keep close to the base policy
    learning_rate=5e-6,
    num_train_epochs=1,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=8,
    max_length=1024, max_prompt_length=512,
    lr_scheduler_type="cosine", warmup_ratio=0.1,
    logging_steps=10, save_strategy="epoch", bf16=True,
)

try:                                # trl >= 0.12 renamed tokenizer -> processing_class
    trainer = DPOTrainer(model=model, args=config, train_dataset=dataset,
                         processing_class=tokenizer, peft_config=peft_config)
except TypeError:                   # trl 0.9-0.11 (torch 2.4-safe) still use tokenizer=
    trainer = DPOTrainer(model=model, args=config, train_dataset=dataset,
                         tokenizer=tokenizer, peft_config=peft_config)
trainer.train()
trainer.save_model(OUT_DIR)         # saves the small, detachable adapter only

# --- inference: the adapter as a scalable "rhetoric dial" ---
# from peft import PeftModel
# base = AutoModelForCausalLM.from_pretrained(BASE_MODEL, torch_dtype=torch.bfloat16, device_map="auto")
# tuned = PeftModel.from_pretrained(base, OUT_DIR)
# tuned.set_adapter_scale(0.7)      # 0 = base style, 1 = full de-emphasis; validate empirically (non-monotonic at large scale)
#
# Evaluate against the genre-relative human baseline of §7.7:
#   (1) epanorthosis density vs. the human rate for the target genre (not vs. zero);
#   (2) content fidelity (semantic similarity + factuality) against the base output;
#   (3) human ratings of appropriateness on held-out genres.