experiment/knowledge_editing/run_dualedit.py

"""
Standalone DualEdit script for hallucination suppression.

Reads caption_targets.json directly (output of build_caption_targets.py)
and runs DualEdit to train adapters that suppress toilet hallucinations.

Usage:
    python -m experiment.knowledge_editing.run_dualedit \
        --caption_targets experiment/data/caption_targets.json \
        --output_dir dualedit_outputs

    # Limit edit instances (for quick testing)
    python -m experiment.knowledge_editing.run_dualedit \
        --caption_targets experiment/data/caption_targets.json \
        --n_edits 50 \
        --output_dir dualedit_outputs

    # Use custom hparams
    python -m experiment.knowledge_editing.run_dualedit \
        --caption_targets experiment/data/caption_targets.json \
        --hparams experiment/knowledge_editing/hparams/dualedit.yaml
"""

import argparse
import json
import os
import sys
from pathlib import Path

import torch
from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from experiment.config.relation_config import get_relation_config
from experiment.knowledge_editing.llava15_compat import LLaVA15ProcessorWrapper
from experiment.knowledge_editing.dualedit import apply_dualedit_to_multimodal_model
from experiment.knowledge_editing.dualedit.dualedit_hparams import DualEditHyperParams


# ---------------------------------------------------------------------------
# Data mapping: caption_targets.json → DualEdit requests
# ---------------------------------------------------------------------------

EDIT_PROMPT = "Describe this image."
REPHRASE_PROMPT = "What do you see in this image?"


def _load_image(entry: dict, hf_images: dict | None) -> Image.Image | None:
    """Load image from disk path or HF cache."""
    path = entry.get("image_path")
    if path and os.path.isfile(path):
        return Image.open(path).convert("RGB")
    if hf_images is not None:
        img = hf_images.get(entry.get("image_id") or path)
        if img is not None:
            return img.convert("RGB")
    return None


def _build_hf_cache(dataset_id: str) -> dict:
    from experiment.data.hf_loader import load_hf_dataset
    from datasets import concatenate_datasets

    print(f"  Loading images from HuggingFace ({dataset_id})...")
    ds = load_hf_dataset(dataset_id)
    if hasattr(ds, "keys"):
        ds = concatenate_datasets([ds[s] for s in ds])
    return {item["image_id"]: item["image"] for item in ds}


def load_requests_from_caption_targets(
    caption_targets_path: str,
    n_edits: int | None = None,
    split: str = "train",
    dataset_id: str = "pbcong/bathroom-toilet",
    relation: str = "bathroom_toilet",
) -> list[dict]:
    """Convert caption_targets.json to DualEdit request list.

    Edit instances:  efficacy category, is_hallucinating=True, is_usable=True
    Locality images: scene_with_object (model should still describe the object)
    Rephrase image:  a different edit image (tests visual generalization)

    Args:
        caption_targets_path: Path to caption_targets.json.
        n_edits: Max number of edit instances. None = use all.
        split: "train" or "val" split to use for edit instances (default: "val").
        dataset_id: HF dataset ID for image loading fallback.
        relation: Relation key from relations.json.

    Returns:
        List of request dicts ready for apply_dualedit_to_multimodal_model().
    """
    rc = get_relation_config(relation)
    efficacy_cat = rc.efficacy_category
    locality_cat = rc.scene_with_object

    with open(caption_targets_path) as f:
        data = json.load(f)

    images = data["images"]

    # Edit instances: hallucinating efficacy-category images with usable clean captions
    edit_entries = [
        entry for entry in images.values()
        if entry.get("category") == efficacy_cat
        and entry.get("is_hallucinating") is True
        and entry.get("is_usable") is True
        and entry.get("cleaned_caption") is not None
        and entry.get("split") == split
    ]

    # Locality images: scene WITH object (preserve ability to mention the object)
    loc_entries = [
        entry for entry in images.values()
        if entry.get("category") == locality_cat
        and entry.get("original_caption") is not None
    ]

    if not edit_entries:
        raise ValueError(
            f"No usable hallucinating edit instances found in {caption_targets_path} "
            f"(split={split}, category={efficacy_cat}). Run build_caption_targets.py first."
        )
    if not loc_entries:
        raise ValueError(f"No {locality_cat} locality instances found.")

    print(f"  Edit instances (split={split}): {len(edit_entries)}")
    print(f"  Locality instances: {len(loc_entries)}")

    if n_edits is not None:
        import random
        rng = random.Random(42)
        edit_entries = rng.sample(edit_entries, min(n_edits, len(edit_entries)))
        print(f"  Using {len(edit_entries)} edit instances (n_edits={n_edits}, seed=42)")

    # Check if images need to be loaded from HF
    all_entries = edit_entries + loc_entries
    needs_hf = any(
        not e.get("image_path") or not os.path.isfile(e.get("image_path", ""))
        for e in all_entries
    )
    hf_images = _build_hf_cache(dataset_id) if needs_hf else None

    requests = []
    for i, entry in enumerate(edit_entries):
        edit_image = _load_image(entry, hf_images)
        if edit_image is None:
            print(f"  Skipping {entry.get('image_path', '?')}: image not found")
            continue

        # Rephrase: a different edit image for visual generalization
        rephrase_entry = edit_entries[(i + 1) % len(edit_entries)]
        rephrase_image = _load_image(rephrase_entry, hf_images) or edit_image

        # Locality: bathroom with toilet — model should still mention toilet here
        loc_entry = loc_entries[i % len(loc_entries)]
        loc_image = _load_image(loc_entry, hf_images) or edit_image
        loc_answer = loc_entry.get("original_caption") or "Describe this image."

        requests.append({
            "prompt": EDIT_PROMPT,
            "target": entry["cleaned_caption"],
            "image": edit_image,
            "file_type": "image",
            # Generality: different prompt, different image
            "rephrase_prompt": REPHRASE_PROMPT,
            "image_rephrase": rephrase_image,
            # Locality: text-only unrelated knowledge
            "locality_prompt": "What is the capital of France?",
            "locality_ground_truth": "Paris",
            # Locality: vision — bathroom WITH toilet should still be described
            "multimodal_locality_image": loc_image,
            "multimodal_locality_prompt": EDIT_PROMPT,
            "multimodal_locality_ground_truth": loc_answer,
        })

    print(f"  Built {len(requests)} requests")
    return requests


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

def load_model(model_name: str, device: str):
    from transformers import AutoProcessor, LlavaForConditionalGeneration

    print(f"Loading {model_name}...")
    model = LlavaForConditionalGeneration.from_pretrained(
        model_name, torch_dtype=torch.float16, device_map={"": device},
    )
    raw_processor = AutoProcessor.from_pretrained(model_name)
    processor = LLaVA15ProcessorWrapper(raw_processor)
    return model, processor


# ---------------------------------------------------------------------------
# Save
# ---------------------------------------------------------------------------

def save_adapter(model, output_dir: str):
    """Save DualEdit adapter state to output_dir/dualedit_state.pt."""
    os.makedirs(output_dir, exist_ok=True)
    if not hasattr(model, "_dualedit_state"):
        print("  WARNING: model has no _dualedit_state, nothing saved")
        return

    state = model._dualedit_state
    save_state = {k: v for k, v in state.items()
                  if k not in ("vision_hook_handle", "text_hook_handle")}
    out_path = os.path.join(output_dir, "dualedit_state.pt")
    torch.save(save_state, out_path)
    print(f"  Saved adapter state → {out_path}")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Run DualEdit directly from caption_targets.json"
    )
    parser.add_argument("--caption_targets", type=str,
                        default="experiment/data/caption_targets.json")
    parser.add_argument("--output_dir", type=str, default="dualedit_outputs")
    parser.add_argument("--model_name", type=str, default="llava-hf/llava-1.5-7b-hf")
    parser.add_argument("--hparams", type=str,
                        default=os.path.join(
                            os.path.dirname(os.path.abspath(__file__)),
                            "hparams", "dualedit.yaml",
                        ))
    parser.add_argument("--n_edits", type=int, default=None,
                        help="Max edit instances (default: all)")
    parser.add_argument("--split", type=str, default="val",
                        choices=["train", "val"])
    parser.add_argument("--relation", type=str, default="bathroom_toilet",
                        help="Relation key from relations.json (default: bathroom_toilet)")
    parser.add_argument("--dataset_id", type=str, default=None,
                        help="HF dataset ID (default: auto from relation config)")
    parser.add_argument("--device", type=str, default="cuda")
    args = parser.parse_args()

    rc = get_relation_config(args.relation)
    dataset_id = args.dataset_id or rc.dataset_id

    os.makedirs(args.output_dir, exist_ok=True)

    # Build requests directly from caption_targets.json
    print(f"\nLoading caption targets from {args.caption_targets}...")
    requests = load_requests_from_caption_targets(
        args.caption_targets,
        n_edits=args.n_edits,
        split=args.split,
        dataset_id=dataset_id,
        relation=args.relation,
    )
    if not requests:
        print("ERROR: No valid requests. Check caption_targets.json.")
        return

    # Load model
    model, processor = load_model(args.model_name, args.device)

    # Load hparams
    hparams = DualEditHyperParams.from_hparams(args.hparams)
    print(f"\nHparams: lr={hparams.edit_lr}, iters={hparams.n_iterations}, "
          f"batch={hparams.batch_size}, gate_threshold={hparams.gating_threshold}")

    # Run DualEdit
    print(f"\nRunning DualEdit on {len(requests)} edit instances...")
    edited_model, adapter_state = apply_dualedit_to_multimodal_model(
        model, processor, requests, hparams,
        copy=False, return_orig_weights=True, keep_original_weight=False,
    )

    # Save
    save_adapter(edited_model, args.output_dir)

    # Save run config
    config = {
        "caption_targets": args.caption_targets,
        "model_name": args.model_name,
        "n_edits": len(requests),
        "split": args.split,
        "hparams": args.hparams,
    }
    with open(os.path.join(args.output_dir, "run_config.json"), "w") as f:
        json.dump(config, f, indent=2)

    print(f"\nDone. Outputs saved to {args.output_dir}/")


if __name__ == "__main__":
    main()