experiment/knowledge_editing/run_visedit.py

"""
Run VisEdit (VEAD) baseline for hallucination suppression.

Uses the exact same data loading as DualEdit (run_dualedit.sh):
  build_requests(edit_set, use_eval_instances=True)
  → edit_set["eval_instances"]["bathroom_no_toilet"] (the fixed ~50 images)

Only differences from DualEdit: prompts match run_dualedit.py and the
requests are converted to VisEdit EIC JSON format for vead_train.py.

NOTE: VisEdit loads TWO copies of LLaVA-1.5-7b simultaneously (one for
training, one for data preprocessing). This requires ~32 GB VRAM. Use
--proc_device to place the preprocessing model on a second GPU.

Usage:
    python -m experiment.knowledge_editing.run_visedit \\
        --edit_set experiment/data/edit_set.json \\
        --output_dir step4_baseline_outputs/visedit \\
        --device cuda:0 \\
        --proc_device cuda:1

    # Skip training (use existing checkpoint)
    python -m experiment.knowledge_editing.run_visedit \\
        --edit_set experiment/data/edit_set.json \\
        --output_dir step4_baseline_outputs/visedit \\
        --skip_train
"""

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

# ---------------------------------------------------------------------------
# Prompts — match run_dualedit.py exactly
# ---------------------------------------------------------------------------

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

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

VISEDIT_DIR = Path(__file__).resolve().parents[2] / "VisEdit"


def requests_to_visedit_json(requests: list[dict], img_dir: Path) -> list[dict]:
    """Convert request list to VisEdit EIC JSON format.

    Saves PIL images to img_dir and records filenames.
    Fields: image, src, alt, rephrase, image_rephrase, loc, loc_ans, m_loc, m_loc_q, m_loc_a
    """
    img_dir.mkdir(parents=True, exist_ok=True)

    def save_img(pil_img, filename: str) -> str:
        dst = img_dir / filename
        if not dst.exists():
            pil_img.save(dst, format="JPEG", quality=95)
        return filename

    records = []
    for req in requests:
        image_id = req["_image_id"]
        fname = save_img(req["image"], f"edit_{image_id}.jpg")

        reph_img = req.get("image_rephrase")
        reph_fname = save_img(reph_img, f"reph_{image_id}.jpg") if reph_img else fname

        loc_img = req.get("multimodal_locality_image")
        loc_fname = save_img(loc_img, f"loc_{image_id}.jpg") if loc_img else fname

        records.append({
            "image":          fname,
            "src":            req["prompt"],
            "alt":            req["target"],
            "rephrase":       req["rephrase_prompt"],
            "image_rephrase": reph_fname,
            "loc":            req["locality_prompt"],
            "loc_ans":        req["locality_ground_truth"],
            "m_loc":          loc_fname,
            "m_loc_q":        req["multimodal_locality_prompt"],
            "m_loc_a":        req["multimodal_locality_ground_truth"],
        })

    return records


def patch_global_py(model_name: str = "llava-hf/llava-1.5-7b-hf"):
    """Overwrite VisEdit/utils/GLOBAL.py with correct absolute paths."""
    global_py = VISEDIT_DIR / "utils" / "GLOBAL.py"
    root = str(VISEDIT_DIR)
    content = (
        f"ROOT_PATH = {root!r}\n"
        f"model_path_map = {{\n"
        f"    'llava-v1.5-7b': {model_name!r},\n"
        f"    'blip2-opt-2.7b': 'models/blip2-opt-2.7b',\n"
        f"    'minigpt-4-vicuna-7b': 'models/minigpt-4-vicuna-7b',\n"
        f"}}\n"
    )
    global_py.write_text(content)
    print(f"  Patched {global_py}")


def find_latest_checkpoint(records_dir: Path) -> str | None:
    """Find the most recently modified checkpoint in records/vead/llava-v1.5-7b/."""
    base = records_dir / "vead" / "llava-v1.5-7b"
    if not base.exists():
        return None
    ckpts = sorted(base.rglob("epoch-*"), key=lambda p: p.stat().st_mtime, reverse=True)
    return str(ckpts[0]) if ckpts else None


def run_training(device: str, proc_device: str, epochs: int,
                 batch_size: int, save_per: int):
    """Run vead_train.py as subprocess from VISEDIT_DIR."""
    proc_idx = proc_device.split(":")[-1] if ":" in proc_device else "0"
    cmd = [
        sys.executable, "vead_train.py",
        "-mn", "llava",
        "-dna", "EIC",
        "-bs", str(batch_size),
        "-dvc", device,
        "-edvc", proc_idx,
        "-eps", str(epochs),
        "-sci", str(save_per),
        "-tnp", "bathroom-toilet",
    ]
    print(f"\n  Running: {' '.join(cmd)}")
    env = os.environ.copy()
    env["PYTORCH_ALLOC_CONF"] = "expandable_segments:True"
    result = subprocess.run(cmd, cwd=str(VISEDIT_DIR), env=env)
    if result.returncode != 0:
        raise RuntimeError(f"vead_train.py exited with code {result.returncode}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Prepare data and train VisEdit (VEAD) baseline"
    )
    parser.add_argument("--edit_set", type=str,
                        default="experiment/data/edit_set.json",
                        help="Path to edit_set.json — same file DualEdit uses")
    parser.add_argument("--output_dir", type=str,
                        default="step4_baseline_outputs/visedit")
    parser.add_argument("--dataset_id", type=str,
                        default="pbcong/bathroom-toilet",
                        help="HF dataset ID for image loading fallback")
    parser.add_argument("--model_name", type=str,
                        default="llava-hf/llava-1.5-7b-hf",
                        help="HF model ID or local path for llava-v1.5-7b")
    parser.add_argument("--device", type=str, default="cuda:0",
                        help="CUDA device for the main (training) model")
    parser.add_argument("--proc_device", type=str, default="cuda:1",
                        help="CUDA device for data-preprocessing copy of model")
    parser.add_argument("--epochs", type=int, default=1000)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--save_per", type=int, default=500,
                        help="Save checkpoint every N iterations")
    parser.add_argument("--skip_train", action="store_true",
                        help="Skip training; look for existing checkpoint")
    args = parser.parse_args()

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

    # -----------------------------------------------------------------------
    # 1. Load data — identical to DualEdit (run_baselines.py use_eval_instances=True)
    # -----------------------------------------------------------------------
    from experiment.knowledge_editing.run_baselines import build_requests

    print(f"Loading edit set from {args.edit_set}...")
    with open(args.edit_set) as f:
        edit_set = json.load(f)
    print(f"  Stats: {edit_set['stats']}")

    print("\nBuilding requests (eval_instances.bathroom_no_toilet, same as DualEdit)...")
    requests = build_requests(edit_set, dataset_id=args.dataset_id,
                              use_eval_instances=True)

    # Override prompts to match run_dualedit.py
    for req in requests:
        req["prompt"] = EDIT_PROMPT
        req["rephrase_prompt"] = REPHRASE_PROMPT
        req["multimodal_locality_prompt"] = EDIT_PROMPT

    print(f"  {len(requests)} requests ready")

    # -----------------------------------------------------------------------
    # 2. Convert to VisEdit EIC JSON format
    # -----------------------------------------------------------------------
    visedit_img_dir = VISEDIT_DIR / "data" / "easy-edit-mm" / "images"
    visedit_cap_dir = VISEDIT_DIR / "data" / "easy-edit-mm" / "caption"
    visedit_cap_dir.mkdir(parents=True, exist_ok=True)

    print("\nConverting to VisEdit EIC format...")
    records = requests_to_visedit_json(requests, visedit_img_dir)

    train_json_path = visedit_cap_dir / "caption_train_edit.json"
    with open(train_json_path, "w") as f:
        json.dump(records, f, indent=2)
    print(f"  Saved {len(records)} records → {train_json_path}")

    # VisEdit also expects a caption_eval_edit.json (same data for EIC mode)
    eval_json_path = visedit_cap_dir / "caption_eval_edit.json"
    with open(eval_json_path, "w") as f:
        json.dump(records, f, indent=2)
    print(f"  Saved {len(records)} records → {eval_json_path}")

    # -----------------------------------------------------------------------
    # 3. Patch GLOBAL.py
    # -----------------------------------------------------------------------
    print("\nPatching VisEdit GLOBAL.py...")
    patch_global_py(args.model_name)

    # -----------------------------------------------------------------------
    # 4. Training
    # -----------------------------------------------------------------------
    if not args.skip_train:
        print("\n>>> Running VEAD training...")
        run_training(
            device=args.device,
            proc_device=args.proc_device,
            epochs=args.epochs,
            batch_size=args.batch_size,
            save_per=args.save_per,
        )
    else:
        print("\n>>> Skipping training (--skip_train)")

    # -----------------------------------------------------------------------
    # 5. Find checkpoint
    # -----------------------------------------------------------------------
    records_dir = VISEDIT_DIR / "records"
    ckpt = find_latest_checkpoint(records_dir)
    if ckpt is None:
        print("WARNING: No checkpoint found. Run training first.")
    else:
        print(f"\n  Found checkpoint: {ckpt}")

    # -----------------------------------------------------------------------
    # 6. Save eval_targets.json and edit_image_ids.json
    # -----------------------------------------------------------------------
    eval_targets = {req["_image_id"]: req["target"] for req in requests}
    eval_targets_path = os.path.join(args.output_dir, "eval_targets.json")
    with open(eval_targets_path, "w") as f:
        json.dump(eval_targets, f, indent=2)
    print(f"  Saved {len(eval_targets)} eval targets → {eval_targets_path}")

    edit_image_ids = [req["_image_id"] for req in requests]
    edit_image_ids_path = os.path.join(args.output_dir, "edit_image_ids.json")
    with open(edit_image_ids_path, "w") as f:
        json.dump(edit_image_ids, f, indent=2)
    print(f"  Saved {len(edit_image_ids)} edit image IDs → {edit_image_ids_path}")

    # -----------------------------------------------------------------------
    # 7. Save run config
    # -----------------------------------------------------------------------
    run_config = {
        "visedit_dir":    str(VISEDIT_DIR),
        "checkpoint":     ckpt,
        "model_name":     args.model_name,
        "device":         args.device,
        "n_train":        len(records),
        "edit_set":       args.edit_set,
        "dataset_id":     args.dataset_id,
        "eval_targets":   eval_targets_path,
        "edit_image_ids": edit_image_ids_path,
    }
    config_path = os.path.join(args.output_dir, "run_config.json")
    with open(config_path, "w") as f:
        json.dump(run_config, f, indent=2)
    print(f"Run config saved → {config_path}")


if __name__ == "__main__":
    main()