method.py

import json
import time
from itertools import product
from pathlib import Path
from typing import Dict, List

import torch
from loguru import logger
from omegaconf import DictConfig, ListConfig
from peft import PeftModel
from transformers import AutoModelForCausalLM

from diffing.methods.diffing_method import DiffingMethod
from diffing.utils.configs import ModelConfig
from .agents import ContrastiveDecodingAgent, GeneralContrastiveDecodingAgent
import functools
from .decoding import (
    generate_contrastive, generate_vanilla,
    cd_combination, pure_cd_combination, with_warmup,
    min_p_mask, min_p_base_mask, top_k_mask, top_p_mask,
)


def _load_plain_model(model_cfg: ModelConfig, device: str):
    """Load a model as a plain HF nn.Module, no nnsight wrapping.

    CDD only ever needs raw forward-pass logits (see decoding.py, which calls
    model(...) directly and expects AutoModelForCausalLM/PeftModel). Loading
    through the shared nnsight-based load_model_from_config (needed by ADL for
    activation access) was a leftover mismatch: decoding.py stopped assuming
    nnsight tracing while model loading still went through StandardizedTransformer,
    which is the likely source of stray meta-device tensors after adapter loading.

    Mirrors the loading pattern already validated in scripts/run_cdd_custom.py.
    """
    base_model_id = (
        model_cfg.base_model_id if model_cfg.base_model_id is not None else model_cfg.model_id
    )
    adapter_id = model_cfg.model_id if base_model_id != model_cfg.model_id else None

    dtype = getattr(torch, model_cfg.dtype)
    model = AutoModelForCausalLM.from_pretrained(
        base_model_id,
        torch_dtype=dtype,
        attn_implementation=model_cfg.attn_implementation,
        device_map={"": device},
        trust_remote_code=model_cfg.trust_remote_code,
    )
    if adapter_id is not None:
        model = PeftModel.from_pretrained(model, adapter_id, device_map={"": device})
    model.eval()
    return model


class ContrastiveDecodingMethod(DiffingMethod):
    """Diffing method that uses contrastive decoding to surface finetuning content.

    Generates text via CD(x) = (1+ε)·log p_ft(x) - ε·log p_base(x) with a vague
    prefill ("The ", "A ", ...) so the finetuning prior dominates unconditionally.
    """

    def __init__(self, cfg: DictConfig):
        super().__init__(cfg)
        self.results_dir = Path(cfg.diffing.results_dir) / "contrastive_decoding"
        self.results_dir.mkdir(parents=True, exist_ok=True)
        self._timings: Dict[str, float] = {}
        self._active_combo_tag: str | None = None

    @property
    def base_model(self):
        """Plain HF model, bypassing the nnsight-based loader (see _load_plain_model)."""
        if self._base_model is None:
            self._base_model = _load_plain_model(self.base_model_cfg, self.device)
        return self._base_model

    @property
    def finetuned_model(self):
        """Plain HF/PeftModel, bypassing the nnsight-based loader (see _load_plain_model)."""
        if self._finetuned_model is None:
            self._finetuned_model = _load_plain_model(self.finetuned_model_cfg, self.device)
        return self._finetuned_model

    def run(self) -> None:
        logger.info("Running Contrastive Decoding method")

        # Generation params
        n_trials = int(self.method_cfg.n_trials)
        max_new_tokens = int(self.method_cfg.max_new_tokens)
        temperature = float(self.method_cfg.temperature)
        do_sample = bool(self.method_cfg.do_sample)
        prefill_batch_size = int(getattr(self.method_cfg, "prefill_batch_size", 10))
        warmup_steps = int(getattr(self.method_cfg, "warmup_steps", 0))

        raw_beta = self.method_cfg.beta
        betas = [float(e) for e in raw_beta] if isinstance(raw_beta, (list, ListConfig)) else [float(raw_beta)]
        filter_name, filter_params, filter_tag_fn, make_filter_fn, filter_param_key = self._parse_filter_config()

        if not do_sample and n_trials > 1:
            logger.warning(
                f"do_sample=False with n_trials={n_trials}: greedy decoding is deterministic, "
                "all trials will produce identical output. Set do_sample=True for diverse coverage."
            )

        t0 = time.perf_counter()
        self.setup_models()
        self._timings["model_setup"] = time.perf_counter() - t0
        tokenizer = self.tokenizer

        # Load JSON prefills (document + user modes)
        prefills_file = Path(self.method_cfg.prefills_file)
        with open(prefills_file) as f:
            prefills_data = json.load(f)
        json_prefills: List[str] = [p for group in prefills_data.values() for p in group]
        logger.info(f"Loaded {len(json_prefills)} prefills from {prefills_file}")

        simulators = list(getattr(self.method_cfg, "simulators", ["document"]))
        logger.info(f"Simulator modes: {simulators}")

        # Build per-mode prompt lists
        doc_prefills: List[str] = json_prefills if "document" in simulators else []

        user_prefills: List[str] = []
        if "user" in simulators:
            sentinel = "\x00"
            raw = tokenizer.apply_chat_template(
                [{"role": "user", "content": sentinel}],
                tokenize=False,
                add_generation_prompt=False,
            )
            chat_prefix = raw[:raw.index(sentinel)]
            if tokenizer.bos_token and chat_prefix.startswith(tokenizer.bos_token):
                chat_prefix = chat_prefix[len(tokenizer.bos_token):]
            logger.info(f"Chat template user prefix: {repr(chat_prefix)}")
            user_prefills = [chat_prefix + p for p in json_prefills]

        asst_items: List[dict] = []
        asst_prompts: List[str] = []
        if "assistant" in simulators:
            asst_file = Path(self.method_cfg.assistant_prefills_file)
            with open(asst_file) as f:
                asst_data = json.load(f)
            asst_items = [item for group in asst_data.values() for item in group]
            for item in asst_items:
                raw = tokenizer.apply_chat_template(
                    [{"role": "user", "content": item["user_message"]}],
                    tokenize=False,
                    add_generation_prompt=True,
                )
                if tokenizer.bos_token and raw.startswith(tokenizer.bos_token):
                    raw = raw[len(tokenizer.bos_token):]
                asst_prompts.append(raw + item.get("assistant_start", ""))
            logger.info(f"Loaded {len(asst_items)} assistant prefill items from {asst_file}")

        if getattr(tokenizer, "padding_side", None) != "left":
            tokenizer.padding_side = "left"

        combination_name = getattr(self.method_cfg, "combination", "cd")
        is_pure_cd = combination_name == "pure_cd"

        # Run base model vanilla generation once — reused across all (beta, alpha) combos.
        save_base = bool(getattr(self.method_cfg, "save_base_generations", False))
        vanilla_n_trials = int(getattr(self.method_cfg, "vanilla_n_trials", 3))
        base_generations_by_prefill: dict[str, list[str]] = {}
        base_generations_file = self.results_dir / "base_generations.json"
        if save_base:
            if base_generations_file.exists() and not self.method_cfg.overwrite:
                logger.info("Loading existing base generations from base_generations.json")
                with open(base_generations_file) as f:
                    base_generations_by_prefill = {e["prefill"]: e["generations"] for e in json.load(f)}
            else:
                # Base generations always use document-mode prefills (raw, no template)
                base_prefills = doc_prefills if doc_prefills else json_prefills
                logger.info(f"Running base model vanilla generation ({vanilla_n_trials} trials per prefill)")
                t0 = time.perf_counter()
                chunks = [base_prefills[i:i + prefill_batch_size] for i in range(0, len(base_prefills), prefill_batch_size)]
                for chunk in chunks:
                    prompts = [p for p in chunk for _ in range(vanilla_n_trials)]
                    gens = generate_vanilla(
                        model=self.base_model,
                        tokenizer=tokenizer,
                        prompts=prompts,
                        max_new_tokens=max_new_tokens,
                        temperature=temperature,
                        do_sample=True,
                    )
                    for i, prefill in enumerate(chunk):
                        base_generations_by_prefill[prefill] = gens[i * vanilla_n_trials:(i + 1) * vanilla_n_trials]
                self._timings["base_vanilla_generation"] = time.perf_counter() - t0
                with open(base_generations_file, "w") as f:
                    json.dump(
                        [{"prefill": p, "generations": g} for p, g in base_generations_by_prefill.items()],
                        f, indent=2,
                    )
                logger.info(f"Base generations saved to {base_generations_file}")

        combos = [(None, fp) for fp in filter_params] if is_pure_cd else list(product(betas, filter_params))
        logger.info(f"Running {len(combos)} combos (combination={combination_name}, filter={filter_name}): betas={betas}, filter_params={filter_params}")

        for beta, filter_param in combos:
            ftag = filter_tag_fn(filter_param)
            results_file = (
                self.results_dir / f"cd_results_pure_cd_t{temperature}_{ftag}.json"
                if is_pure_cd
                else self.results_dir / f"cd_results_b{beta}_t{temperature}_{ftag}.json"
            )

            # Load existing results for mode-level overwrite check
            existing_results: dict = {}
            if results_file.exists() and not self.method_cfg.overwrite:
                with open(results_file) as f:
                    existing = json.load(f)
                raw = existing.get("results", {})
                # backward compat: old list format → treat as document mode
                existing_results = raw if isinstance(raw, dict) else {"document": raw}
            results_by_mode: dict = dict(existing_results)

            logger.info(
                f"combination={combination_name}  {filter_name}={filter_param}"
                + (f"  beta={beta}  temperature={temperature}" if not is_pure_cd else "")
            )

            filter_fn = make_filter_fn(filter_param)
            if is_pure_cd:
                combination_fn = functools.partial(pure_cd_combination, filter_fn=filter_fn, temperature=temperature)
            else:
                combination_fn = functools.partial(cd_combination, beta=beta, filter_fn=filter_fn, temperature=temperature)
            combination_fn = with_warmup(combination_fn, warmup_steps)

            def _run_cd_chunks(prompts_list: List[str]) -> List[str]:
                """Run CD generation in batches; returns flat list of generated strings (len == len(prompts_list))."""
                all_gens: List[str] = []
                chunks = [prompts_list[i:i + prefill_batch_size * n_trials]
                          for i in range(0, len(prompts_list), prefill_batch_size * n_trials)]
                for chunk in chunks:
                    t0 = time.perf_counter()
                    gens = generate_contrastive(
                        base_model=self.base_model,
                        finetuned_model=self.finetuned_model,
                        tokenizer=tokenizer,
                        prompts=chunk,
                        max_new_tokens=max_new_tokens,
                        do_sample=do_sample,
                        combination_fn=combination_fn,
                        warmup_steps=warmup_steps,
                    )
                    self._timings["generation"] = self._timings.get("generation", 0) + time.perf_counter() - t0
                    all_gens.extend(gens)
                return all_gens

            for mode in simulators:
                if mode in results_by_mode:
                    logger.info(f"  Skipping {mode} simulator (already in {results_file.name})")
                    continue

                logger.info(f"  Running {mode} simulator")

                if mode == "document":
                    prompts = [p for p in doc_prefills for _ in range(n_trials)]
                    all_gens = _run_cd_chunks(prompts)
                    mode_results = []
                    for i, prefill in enumerate(doc_prefills):
                        generations = all_gens[i * n_trials:(i + 1) * n_trials]
                        mode_results.append({"prefill": prefill, "generations": generations})
                        logger.info(f"    [{repr(prefill)}] Sample: {repr(generations[0][:120])}")

                elif mode == "user":
                    prompts = [p for p in user_prefills for _ in range(n_trials)]
                    all_gens = _run_cd_chunks(prompts)
                    mode_results = []
                    for i, prefill in enumerate(user_prefills):
                        generations = all_gens[i * n_trials:(i + 1) * n_trials]
                        mode_results.append({"prefill": prefill, "generations": generations})
                        logger.info(f"    [user:{repr(prefill[-30:])}] Sample: {repr(generations[0][:120])}")

                elif mode == "assistant":
                    prompts = [p for p in asst_prompts for _ in range(n_trials)]
                    all_gens = _run_cd_chunks(prompts)
                    mode_results = []
                    for i, item in enumerate(asst_items):
                        generations = all_gens[i * n_trials:(i + 1) * n_trials]
                        mode_results.append({
                            "user_message": item["user_message"],
                            "assistant_start": item.get("assistant_start", ""),
                            "generations": generations,
                        })
                        logger.info(f"    [asst:{repr(item['user_message'])}] Sample: {repr(generations[0][:120])}")

                else:
                    logger.warning(f"Unknown simulator mode '{mode}', skipping")
                    continue

                results_by_mode[mode] = mode_results

            if results_by_mode == existing_results:
                logger.info(f"  No new modes generated, skipping write for {results_file.name}")
                continue

            output = {
                "combination": combination_name,
                "simulators": simulators,
                "plausibility_filter": filter_name,
                filter_param_key: filter_param,
                "n_trials": n_trials,
                "max_new_tokens": max_new_tokens,
                "do_sample": do_sample,
                "temperature": temperature,
                "warmup_steps": warmup_steps,
                **({"beta": beta} if not is_pure_cd else {}),
                "results": results_by_mode,
            }
            with open(results_file, "w") as f:
                json.dump(output, f, indent=2)
            logger.info(f"Results saved to {results_file}")

        timings_path = self.results_dir / "timings.json"
        timings_path.write_text(json.dumps(self._timings, indent=2))
        logger.info(f"Step timings (s): {self._timings}")
        logger.info(f"Saved timings to {timings_path}")

    @staticmethod
    def has_results(results_dir: Path) -> Dict[str, Dict[str, str]]:
        results = {}
        for model_dir in results_dir.parent.iterdir():
            if not model_dir.is_dir():
                continue
            for organism_dir in model_dir.iterdir():
                if not organism_dir.is_dir():
                    continue
                candidate = next((organism_dir / "contrastive_decoding").glob("cd_results_*.json"), None)
                if candidate is None:
                    continue
                if candidate.exists():
                    model_name = model_dir.name
                    organism_name = organism_dir.name
                    if model_name not in results:
                        results[model_name] = {}
                    results[model_name][organism_name] = str(candidate)
        return results

    def _parse_filter_config(self) -> tuple:
        """Return (filter_name, params_list, tag_fn, make_fn, param_key) from method config."""
        filter_name = getattr(self.method_cfg, "plausibility_filter", "min_p")
        if filter_name == "top_k":
            raw = getattr(self.method_cfg, "top_k", 50)
            params = [int(k) for k in raw] if isinstance(raw, (list, ListConfig)) else [int(raw)]
            return filter_name, params, lambda k: f"k{k}", lambda k: functools.partial(top_k_mask, k=k), "top_k"
        if filter_name == "top_p":
            raw = getattr(self.method_cfg, "top_p", 0.9)
            params = [float(p) for p in raw] if isinstance(raw, (list, ListConfig)) else [float(raw)]
            return filter_name, params, lambda p: f"p{p}", lambda p: functools.partial(top_p_mask, p=p), "top_p"
        if filter_name == "min_p_base":
            raw = self.method_cfg.alpha
            params = [float(a) for a in raw] if isinstance(raw, (list, ListConfig)) else [float(raw)]
            return "min_p_base", params, lambda a: f"ab{a}", lambda a: functools.partial(min_p_base_mask, alpha=a), "alpha"
        # "min_p" — default, backward-compatible (tag: a{alpha})
        raw = self.method_cfg.alpha
        params = [float(a) for a in raw] if isinstance(raw, (list, ListConfig)) else [float(raw)]
        return "min_p", params, lambda a: f"a{a}", lambda a: functools.partial(min_p_mask, alpha=a), "alpha"

    def get_agent(self) -> ContrastiveDecodingAgent:
        organism_type = getattr(self.cfg.organism, "type", "SDF")
        if organism_type == "General":
            return GeneralContrastiveDecodingAgent(cfg=self.cfg)
        return ContrastiveDecodingAgent(cfg=self.cfg)

    def get_combo_tags(self) -> list[str]:
        """Return sorted tags for all result files present in results_dir."""
        return sorted(
            f.stem.replace("cd_results_", "")
            for f in self.results_dir.glob("cd_results_*.json")
        )

    @property
    def results_cfg_tag(self) -> str:
        """Tag used to find the results file — CD hyperparams only, no eval-time filters."""
        if self._active_combo_tag is not None:
            return self._active_combo_tag
        _, filter_params, filter_tag_fn, _, _ = self._parse_filter_config()
        ftag = filter_tag_fn(filter_params[0])
        combination_name = getattr(self.method_cfg, "combination", "cd")
        temperature = float(self.method_cfg.temperature)
        if combination_name == "pure_cd":
            return f"pure_cd_t{temperature}_{ftag}"
        beta = self.method_cfg.beta
        if isinstance(beta, (list, ListConfig)):
            beta = beta[0]
        return f"b{float(beta)}_t{temperature}_{ftag}"

    @property
    def relevant_cfg_tag(self) -> str:
        """Tag used to name the agent output directory — includes eval-time toggles."""
        tag = self.results_cfg_tag
        agent_simulators = getattr(self.method_cfg, "agent_simulators", None)
        if agent_simulators is not None:
            tag = f"{tag}_sim-{'-'.join(sorted(list(agent_simulators)))}"
        return tag