app.py

import os
from typing import List, Dict
from datetime import datetime

import torch
from torch import nn

import gradio as gr
import pandas as pd

from datasets import Dataset

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig,
)

from peft import LoraConfig, get_peft_model
from trl import DPOConfig, DPOTrainer


# =========================================================
#  MODEL LIST (from your BIAS demo)
# =========================================================

MODEL_CHOICES = [
    # Very small / light (good for CPU Spaces)
    "distilgpt2",
    "gpt2",
    "sshleifer/tiny-gpt2",
    "LiquidAI/LFM2-350M",
    "google/gemma-3-270m-it",
    "Qwen/Qwen2.5-0.5B-Instruct",
    "mkurman/NeuroBLAST-V3-SYNTH-EC-150000",

    # Small–medium (~1–2B) – still reasonable on CPU, just slower
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    "google/gemma-3-1b-it",
    "meta-llama/Llama-3.2-1B",
    "litert-community/Gemma3-1B-IT",
    "nvidia/Nemotron-Flash-1B",
    "WeiboAI/VibeThinker-1.5B",
    "Qwen/Qwen3-1.7B",

    # Medium (~2–3B) – probably OK on beefier CPU / small GPU
    "google/gemma-2-2b-it",
    "thu-pacman/PCMind-2.1-Kaiyuan-2B",
    "opendatalab/MinerU-HTML",
    "ministral/Ministral-3b-instruct",
    "HuggingFaceTB/SmolLM3-3B",
    "meta-llama/Llama-3.2-3B-Instruct",
    "nvidia/Nemotron-Flash-3B-Instruct",
    "Qwen/Qwen2.5-3B-Instruct",

    # Heavier (4–8B) – you really want a GPU Space for these
    "Qwen/Qwen3-4B",
    "Qwen/Qwen3-4B-Thinking-2507",
    "Qwen/Qwen3-4B-Instruct-2507",
    "mistralai/Mistral-7B-Instruct-v0.2",
    "allenai/Olmo-3-7B-Instruct",
    "Qwen/Qwen2.5-7B-Instruct",
    "meta-llama/Meta-Llama-3-8B-Instruct",
    "meta-llama/Llama-3.1-8B",
    "meta-llama/Llama-3.1-8B-Instruct",
    "openbmb/MiniCPM4.1-8B",
    "deepseek-ai/DeepSeek-R1-Distill-Llama-8B",
    "rl-research/DR-Tulu-8B",
]

DEFAULT_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
TRAINED_MODEL_DIR = "trained_model"


# =========================================================
#  GLOBALS & CONFIG
# =========================================================

device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = None
policy_model = None
ref_model = None

DEFAULT_DPO_CONFIG = DPOConfig(
    beta=0.1,
    output_dir="dpo_demo",
    num_train_epochs=1,
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    remove_unused_columns=False,
    logging_steps=1,
    gradient_accumulation_steps=1,
    learning_rate=1e-4,
    evaluation_strategy="no",  # warning is fine with current versions
    warmup_steps=0,
    fp16=False,
    save_steps=0,
    report_to="none",
)


# =========================================================
#  LORA TARGET-MODULE HELPER
# =========================================================

def guess_lora_target_modules(model_name: str, base_model) -> List[str]:
    """
    Heuristically choose good LoRA target modules based on the model type/name.
    - GPT-2-like: use c_attn/c_proj
    - LLaMA/Gemma/Mistral/Qwen/etc: use q/k/v/o + MLP projections
    - Fallback: scan Linear module names for known patterns
    """
    model_type = getattr(base_model.config, "model_type", "") or ""
    name_lower = model_name.lower()

    # GPT-2 / DistilGPT-2 / Tiny GPT-2
    if (
        "gpt2" in model_type
        or "gpt2" in name_lower
        or "tiny-gpt2" in name_lower
        or "distilgpt2" in name_lower
    ):
        return ["c_attn", "c_proj"]

    # LLaMA / Gemma / Mistral / Qwen / Olmo / MiniCPM / SmolLM / Nemotron etc.
    if any(
        t in model_type
        for t in [
            "llama",
            "gemma",
            "mistral",
            "qwen",
            "qwen2",
            "olmo",
            "minicpm",
            "smollm",
            "nemotron",
        ]
    ):
        return ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

    # Fallback: inspect Linear modules and see what’s there
    linear_leaf_names = []
    for name, module in base_model.named_modules():
        if isinstance(module, nn.Linear):
            linear_leaf_names.append(name.split(".")[-1])

    candidates = [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
        "c_attn", "c_proj",
    ]
    found = sorted(set(n for n in candidates if n in linear_leaf_names))
    if found:
        return found

    # If absolutely nothing matches, bail with a clear error
    raise ValueError(
        f"Could not guess LoRA target modules for model '{model_name}' "
        f"(model_type='{model_type}'). "
        f"Try setting target_modules manually for this model."
    )


# =========================================================
#  MODEL LOADING
# =========================================================

def load_base_model(model_name: str) -> str:
    """
    Load tokenizer + base model, then create:
      - policy_model: LoRA-adapted (trainable)
      - ref_model: frozen base model for DPO
    """
    global tokenizer, policy_model, ref_model

    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        trust_remote_code=True,
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"

    base_model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
    )
    base_model.config.use_cache = False
    base_model.config.pad_token_id = tokenizer.eos_token_id

    # Choose LoRA target modules dynamically
    target_modules = guess_lora_target_modules(model_name, base_model)

    peft_config = LoraConfig(
        r=4,
        target_modules=target_modules,
        task_type="CAUSAL_LM",
        lora_alpha=8,
        lora_dropout=0.1,
        bias="none",
    )

    # Policy model = base + LoRA (trainable)
    policy = get_peft_model(base_model, peft_config)
    policy.to(device)
    policy.eval()

    # Reference model = frozen base model
    reference = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
    )
    reference.config.use_cache = False
    reference.config.pad_token_id = tokenizer.eos_token_id
    reference.to(device)
    for p in reference.parameters():
        p.requires_grad = False
    reference.eval()

    policy_model = policy
    ref_model = reference

    return (
        f"Loaded base model: **{model_name}** on **{device}** "
        f"with LoRA target_modules={target_modules}"
    )


# Load default on startup
initial_status = load_base_model(DEFAULT_MODEL)


# =========================================================
#  UTILS
# =========================================================

def build_generation_config(
    do_sample: bool,
    temperature: float,
    max_new_tokens: int,
    top_k: int = 20,
    top_p: float = 0.9,
) -> GenerationConfig:
    """
    Helper to build a GenerationConfig from UI settings.
    """
    # Clamp values a bit just to be safe
    temperature = max(0.0, float(temperature))
    max_new_tokens = int(max_new_tokens)
    return GenerationConfig(
        do_sample=bool(do_sample),
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
        max_new_tokens=max_new_tokens,
        pad_token_id=tokenizer.eos_token_id,
    )


def generate_text(
    model: nn.Module,
    prompt: str,
    gen_config: GenerationConfig,
    style_prefix: str = "",
) -> str:
    model.eval()
    full_prompt = style_prefix + prompt

    inputs = tokenizer(
        full_prompt,
        return_tensors="pt",
        padding=False,
    ).to(device)

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            do_sample=gen_config.do_sample,
            top_k=gen_config.top_k,
            top_p=gen_config.top_p,
            temperature=gen_config.temperature,
            max_new_tokens=gen_config.max_new_tokens,
            pad_token_id=gen_config.pad_token_id,
        )

    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if text.startswith(full_prompt):
        return text[len(full_prompt):].strip()
    return text.strip()


def preferences_to_df(preferences: List[Dict]) -> pd.DataFrame:
    if not preferences:
        return pd.DataFrame(columns=["prompt", "chosen", "rejected"])
    return pd.DataFrame(preferences)


def list_trained_model_files() -> List[str]:
    """
    Return a list of filepaths under TRAINED_MODEL_DIR (for download).
    """
    if not os.path.isdir(TRAINED_MODEL_DIR):
        return []
    files: List[str] = []
    for root, dirs, filenames in os.walk(TRAINED_MODEL_DIR):
        for name in filenames:
            files.append(os.path.join(root, name))
    return files


# =========================================================
#  DPO CALLBACKS
# =========================================================

def generate_candidates(
    prompt: str,
    do_sample: bool,
    temperature: float,
    max_new_tokens: int,
) -> tuple[str, str]:
    """
    Generate Answer A (balanced) and Answer B (creative-ish),
    using the same core generation settings from the GUI.
    """
    if not prompt.strip():
        return "", ""

    # Build two configs from the same UI settings,
    # but make B slightly more "wild" by bumping top_k / temperature a bit
    balanced_config = build_generation_config(
        do_sample=do_sample,
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_k=20,
        top_p=0.9,
    )

    # For creative answer, nudge temperature and top_k a bit, but still
    # keep them tied to UI settings.
    creative_temp = float(temperature) + 0.4
    creative_config = build_generation_config(
        do_sample=do_sample,
        temperature=creative_temp,
        max_new_tokens=max_new_tokens,
        top_k=50,
        top_p=0.95,
    )

    style_balanced = (
        "You are a helpful, careful assistant. "
        "Answer clearly and sensibly.\n\nUser: "
    )
    style_creative = (
        "You are a creative assistant who explores unusual ideas and stronger opinions, "
        "while still staying safe.\n\nUser: "
    )

    answer_a = generate_text(
        policy_model,
        prompt,
        balanced_config,
        style_prefix=style_balanced,
    )
    answer_b = generate_text(
        policy_model,
        prompt,
        creative_config,
        style_prefix=style_creative,
    )

    return answer_a, answer_b


def save_preference(
    prompt: str,
    answer_a: str,
    answer_b: str,
    custom_answer: str,
    preference_mode: str,
    state_preferences: List[Dict],
):
    """
    Encode a preference in one of four ways:
      - Prefer A over B        -> chosen=A, rejected=B
      - Prefer B over A        -> chosen=B, rejected=A
      - Prefer custom over A   -> chosen=custom, rejected=A
      - Prefer custom over B   -> chosen=custom, rejected=B
    """
    msg = ""

    if not prompt.strip():
        msg = "No prompt provided."
        return state_preferences, preferences_to_df(state_preferences), msg

    if not answer_a.strip() or not answer_b.strip():
        msg = "Generate both model answers before saving a preference."
        return state_preferences, preferences_to_df(state_preferences), msg

    if not preference_mode:
        msg = "Please choose how to encode the preference."
        return state_preferences, preferences_to_df(state_preferences), msg

    preference_mode = preference_mode.strip()

    chosen = None
    rejected = None

    if preference_mode == "Prefer A over B":
        chosen = answer_a
        rejected = answer_b

    elif preference_mode == "Prefer B over A":
        chosen = answer_b
        rejected = answer_a

    elif preference_mode == "Prefer custom over A":
        if not custom_answer.strip():
            msg = "You selected 'Prefer custom over A' but did not provide a custom answer."
            return state_preferences, preferences_to_df(state_preferences), msg
        chosen = custom_answer
        rejected = answer_a

    elif preference_mode == "Prefer custom over B":
        if not custom_answer.strip():
            msg = "You selected 'Prefer custom over B' but did not provide a custom answer."
            return state_preferences, preferences_to_df(state_preferences), msg
        chosen = custom_answer
        rejected = answer_b

    else:
        msg = f"Unknown preference mode: {preference_mode}"
        return state_preferences, preferences_to_df(state_preferences), msg

    entry = {
        "prompt": prompt.strip(),
        "chosen": chosen.strip(),
        "rejected": rejected.strip(),
    }

    state_preferences = list(state_preferences) + [entry]
    df = preferences_to_df(state_preferences)
    msg = f"Saved preference #{len(state_preferences)}."

    return state_preferences, df, msg


def train_dpo_model(
    state_preferences: List[Dict],
    num_epochs: int,
    learning_rate: float,
    beta: float,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Run DPO training on the accumulated preferences.
    Shows a progress bar/spinner and returns:
      - a detailed status message
      - a 'last trained' timestamp string
      - a list of saved model files for download
    """
    global policy_model, ref_model

    progress(0.0, desc="Checking preferences...")

    if not state_preferences:
        return (
            "⚠️ No preferences collected yet. Add some first.",
            "**Last trained:** never",
            [],
        )

    dataset = Dataset.from_list(state_preferences)

    progress(0.2, desc="Configuring DPO trainer...")

    dpo_config = DPOConfig(
        **{
            **DEFAULT_DPO_CONFIG.to_dict(),
            "num_train_epochs": int(num_epochs),
            "learning_rate": float(learning_rate),
            "beta": float(beta),
        }
    )

    trainer = DPOTrainer(
        model=policy_model,
        ref_model=ref_model,
        args=dpo_config,
        train_dataset=dataset,
        eval_dataset=None,
        tokenizer=tokenizer,
        max_length=256,
    )

    progress(0.4, desc="Training model with DPO...")

    trainer.train()

    progress(0.75, desc="Finalizing and moving model to device...")

    policy_model = trainer.model
    policy_model.to(device)
    policy_model.eval()

    # Save the trained model + tokenizer so you can download them
    progress(0.9, desc="Saving trained model to disk...")

    os.makedirs(TRAINED_MODEL_DIR, exist_ok=True)
    policy_model.save_pretrained(TRAINED_MODEL_DIR)
    tokenizer.save_pretrained(TRAINED_MODEL_DIR)

    files = list_trained_model_files()

    progress(1.0, desc="Done")

    n = len(state_preferences)
    finished_at = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    msg = f"""### ✅ Training complete

- Preference pairs used: **{n}**
- Epochs: **{num_epochs}**
- Learning rate: **{learning_rate}**
- DPO beta (strength): **{beta}**

The tuned policy model + tokenizer have been saved to `{TRAINED_MODEL_DIR}/`.
You can download them using the file list below.
"""

    last_trained_msg = f"**Last trained:** {finished_at}"

    return msg, last_trained_msg, files


def generate_from_aligned_model(
    prompt: str,
    do_sample: bool,
    temperature: float,
    max_new_tokens: int,
) -> str:
    if not prompt.strip():
        return ""
    gen_config = build_generation_config(
        do_sample=do_sample,
        temperature=temperature,
        max_new_tokens=max_new_tokens,
        top_k=20,
        top_p=0.9,
    )
    style_balanced = (
        "You are a helpful, careful assistant. "
        "Answer clearly and sensibly.\n\nUser: "
    )
    return generate_text(
        policy_model,
        prompt,
        gen_config,
        style_prefix=style_balanced,
    )


def on_model_change(
    model_name: str,
    _state_preferences: List[Dict],
):
    """
    When the user picks a new base model:
      - reload tokenizer + policy_model + ref_model
      - clear collected preferences (since they belong to previous model)
      - reset training status, 'last trained', and download list
    """
    status = load_base_model(model_name)
    empty_prefs: List[Dict] = []
    df = preferences_to_df(empty_prefs)
    reset_msg = (
        status
        + "\n\nPreferences cleared (new model = new preference data)."
    )
    last_trained_reset = "**Last trained:** (reset for new base model)"
    files_reset: List[str] = []
    # returns: model_status, prefs, pref_table_df, train_status, last_trained, files
    return reset_msg, empty_prefs, df, "", last_trained_reset, files_reset


# =========================================================
#  GRADIO UI
# =========================================================

with gr.Blocks() as demo:
    gr.Markdown(
        """
    # 🔧 DPO Playground – Preference Tuning on Different Models

    - Pick a **base model** from the dropdown.
    - Ask a question and generate two answers:
      - **A** = balanced / normal
      - **B** = creative / more extreme
    - Optionally write **your own ideal answer**.
    - Choose how to encode the preference (e.g. A over B, custom over A, etc.).
    - Collect several preferences and **train the model with DPO**.
    - Test how the aligned policy model behaves on new prompts.
    - Download the tuned model (LoRA adapter + tokenizer) after training.
    - **Control temperature, sampling, and max_new_tokens directly in the UI.**
    """
    )

    state_preferences = gr.State([])

    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=MODEL_CHOICES,
            value=DEFAULT_MODEL,
            label="Base model",
        )

    model_status = gr.Markdown(initial_status)

    # -----------------------------------------------------
    # Collect preferences tab
    # -----------------------------------------------------
    with gr.Tab("Collect preferences"):
        with gr.Row():
            prompt_input = gr.Textbox(
                label="Prompt",
                placeholder="Ask anything...",
                lines=3,
            )

        gr.Markdown("### Generation settings for Answer A & B")

        with gr.Row():
            gen_do_sample = gr.Checkbox(
                value=True,
                label="Use sampling (do_sample)",
            )
            gen_temperature = gr.Slider(
                minimum=0.0,
                maximum=1.5,
                value=0.8,
                step=0.05,
                label="Temperature",
            )
            gen_max_new_tokens = gr.Slider(
                minimum=4,
                maximum=256,
                value=128,
                step=4,
                label="Max new tokens",
            )

        generate_btn = gr.Button("Generate A & B")

        with gr.Row():
            answer_a_box = gr.Textbox(
                label="Answer A (balanced / normal)",
                lines=8,
            )
            answer_b_box = gr.Textbox(
                label="Answer B (creative / more extreme)",
                lines=8,
            )

        custom_answer_box = gr.Textbox(
            label="Your own ideal answer (optional)",
            lines=8,
            placeholder="If you want, write the answer you *wish* the model had given.",
        )

        preference_mode = gr.Radio(
            choices=[
                "Prefer A over B",
                "Prefer B over A",
                "Prefer custom over A",
                "Prefer custom over B",
            ],
            label="How should this preference be encoded?",
        )

        save_pref_btn = gr.Button("Save preference")

        pref_status = gr.Markdown("")
        pref_table = gr.Dataframe(
            headers=["prompt", "chosen", "rejected"],
            label="Collected preferences (for DPO training)",
            wrap=True,
        )

        generate_btn.click(
            fn=generate_candidates,
            inputs=[prompt_input, gen_do_sample, gen_temperature, gen_max_new_tokens],
            outputs=[answer_a_box, answer_b_box],
        )

        save_pref_btn.click(
            fn=save_preference,
            inputs=[
                prompt_input,
                answer_a_box,
                answer_b_box,
                custom_answer_box,
                preference_mode,
                state_preferences,
            ],
            outputs=[
                state_preferences,
                pref_table,
                pref_status,
            ],
        )

    # -----------------------------------------------------
    # Train & test tab
    # -----------------------------------------------------
    with gr.Tab("Train & test DPO model"):
        gr.Markdown(
            "Train the LoRA-adapted policy model using your preferences "
            "with **Direct Preference Optimization (DPO)**."
        )

        with gr.Row():
            num_epochs_slider = gr.Slider(
                minimum=1,
                maximum=5,
                step=1,
                value=1,
                label="Number of epochs",
            )
            lr_slider = gr.Slider(
                minimum=1e-5,
                maximum=5e-4,
                step=1e-5,
                value=1e-4,
                label="Learning rate",
            )
            beta_slider = gr.Slider(
                minimum=0.05,
                maximum=0.5,
                step=0.05,
                value=0.1,
                label="DPO beta (strength)",
            )

        train_btn = gr.Button("Train DPO model", variant="primary")
        train_status = gr.Markdown("")
        last_trained = gr.Markdown("**Last trained:** never")

        download_files = gr.Files(
            label="Trained model files (adapter + tokenizer)",
            interactive=False,
        )

        train_btn.click(
            fn=train_dpo_model,
            inputs=[
                state_preferences,
                num_epochs_slider,
                lr_slider,
                beta_slider,
            ],
            outputs=[train_status, last_trained, download_files],
        )

        gr.Markdown("## Try the current policy model")

        with gr.Row():
            test_do_sample = gr.Checkbox(
                value=False,
                label="Use sampling (do_sample) for test",
            )
            test_temperature = gr.Slider(
                minimum=0.0,
                maximum=1.5,
                value=0.0,
                step=0.05,
                label="Temperature (test)",
            )
            test_max_new_tokens = gr.Slider(
                minimum=4,
                maximum=256,
                value=64,
                step=4,
                label="Max new tokens (test)",
            )

        test_prompt = gr.Textbox(
            label="Test prompt",
            placeholder="Ask something to see the aligned model...",
            lines=3,
        )
        test_btn = gr.Button("Generate from DPO policy model")
        test_answer = gr.Textbox(
            label="Policy model answer",
            lines=8,
        )

        test_btn.click(
            fn=generate_from_aligned_model,
            inputs=[
                test_prompt,
                test_do_sample,
                test_temperature,
                test_max_new_tokens,
            ],
            outputs=test_answer,
        )

    # model change: reload + clear prefs + reset train status + last trained + downloads
    model_dropdown.change(
        fn=on_model_change,
        inputs=[model_dropdown, state_preferences],
        outputs=[
            model_status,
            state_preferences,
            pref_table,
            train_status,
            last_trained,
            download_files,
        ],
    )

if __name__ == "__main__":
    demo.queue().launch()