File size: 5,527 Bytes
52874b5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 | #!/usr/bin/env python3
"""
Path A in-place repair of one already-built <name>*_processed dir.
What it does (and ONLY this — bank/.npz files untouched):
1. Load the existing template bank from <root>/template/*.npz.
2. Re-run patched precompute_for_entries (self-exclusion ON, --overwrite ON) on
backup_entries/{train,val}_entries.json, regenerating the `template` field
for every SFT sample.
3. recover_from_backup() rebuilds <root>/hf_dataset/ from the updated backups.
After this, every bank sample's match list is guaranteed not to contain itself.
SFT samples that were missing a template entirely also get one.
Usage:
python repair_self_match.py --root /projects/p32685/RL_dataset/ACDC_10percent_processed
"""
import argparse
import json
import os
import sys
from pathlib import Path
from types import SimpleNamespace
THIS_DIR = Path(__file__).resolve().parent
VESSA_ROOT = THIS_DIR.parent
for p in [str(THIS_DIR),
str(VESSA_ROOT / "train_vessa"),
str(VESSA_ROOT / "train_vessa" / "tools")]:
if p not in sys.path:
sys.path.insert(0, p)
from packing_HF_dataset import GenHFDataset
from template_manager import TemplateManager
from precompute_template import precompute_for_entries
def parse_args():
ap = argparse.ArgumentParser()
ap.add_argument("--root", required=True, type=str,
help="Existing <name>*_processed dir (must contain template/, backup_entries/)")
ap.add_argument("--device", type=str, default="cuda:0")
# NOTE: top-1 only (was previously top-k for downstream random.choice; that mechanism
# is no longer in use, so we just store the single best match).
ap.add_argument("--topk", type=int, default=1)
ap.add_argument("--match_strategy", type=str, default="partial", choices=["partial", "complete"])
ap.add_argument("--max_search", type=int, default=20)
return ap.parse_args()
def repair_one(root, device="cuda:0", topk=1, match_strategy="partial", max_search=20):
root = os.path.abspath(root)
template_dir = os.path.join(root, "template")
bdir = os.path.join(root, "backup_entries")
bp_train = os.path.join(bdir, "train_entries.json")
bp_val = os.path.join(bdir, "val_entries.json")
if not os.path.isdir(template_dir):
raise FileNotFoundError(f"missing template/ under {root}")
if not os.path.isfile(bp_train):
raise FileNotFoundError(f"missing backup_entries/train_entries.json under {root}")
print(f"\n[REPAIR] {root}")
# ---- step 1: load existing template bank ----
tm = TemplateManager(device=device, distance_metric="cosine")
templates, faiss_index, template_ids_list = tm._load_templates(template_dir)
print(f" loaded {len(templates)} templates from {template_dir}")
# ---- step 2: re-run precompute_for_entries with --overwrite=True ----
train_entries = json.load(open(bp_train))
val_entries = json.load(open(bp_val)) if os.path.isfile(bp_val) else []
# NOTE: a Python class body has its own scope and cannot see the function locals,
# so `topk = topk` inside `class _Args:` raised NameError. Use SimpleNamespace instead.
cfg = SimpleNamespace(
overwrite=True, # always re-write the template field
topk=topk, # default 1 — single best match only
match_strategy=match_strategy,
max_search=max_search,
)
u_train = precompute_for_entries(
"train", train_entries, templates, faiss_index, template_ids_list, tm, cfg,
)
print(f" precompute updated train rows: {u_train}/{len(train_entries)}")
u_val = precompute_for_entries(
"val", val_entries, templates, faiss_index, template_ids_list, tm, cfg,
)
print(f" precompute updated val rows : {u_val}/{len(val_entries)}")
with open(bp_train, "w") as f:
json.dump(train_entries, f, indent=2)
if os.path.isfile(bp_val):
with open(bp_val, "w") as f:
json.dump(val_entries, f, indent=2)
print(f" saved updated backup_entries")
# ---- step 3: rebuild hf_dataset from updated backups ----
gen = GenHFDataset()
print(f" recover_from_backup -> rebuilding hf_dataset")
gen.recover_from_backup(output_dir=root)
# ---- step 4: post-fix verification ----
import glob
bank_stems = {os.path.splitext(os.path.basename(p))[0] for p in glob.glob(os.path.join(template_dir, "*.npz"))}
train_after = json.load(open(bp_train))
self_top1, self_in_list, miss = 0, 0, 0
for e in train_after:
stem = os.path.splitext(os.path.basename(e["image"]))[0]
if stem not in bank_stems:
continue
tpl = e.get("template")
if not isinstance(tpl, list) or not tpl:
miss += 1
continue
ids = [t.get("template_id") for t in tpl]
if stem in ids:
self_in_list += 1
if ids[0] == stem:
self_top1 += 1
print(f" POST-FIX: bank={len(bank_stems)} self_in_list={self_in_list} self_top1={self_top1} missing_template={miss}")
if self_in_list == 0 and self_top1 == 0:
print(f" ✓ {os.path.basename(root)}: CLEAN")
else:
print(f" ✗ {os.path.basename(root)}: STILL HAS VIOLATIONS")
return self_in_list == 0 and self_top1 == 0
def main():
args = parse_args()
repair_one(args.root, device=args.device, topk=args.topk,
match_strategy=args.match_strategy, max_search=args.max_search)
if __name__ == "__main__":
main()
|