""" 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()