Spaces:
Sleeping
Sleeping
| import copy | |
| import math | |
| import os | |
| from typing import Any, Dict, List, Optional | |
| import lightning as L | |
| import rootutils | |
| import torch | |
| import torch.nn.functional as F | |
| from dora.hydra import hydra_main | |
| from omegaconf import DictConfig | |
| rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True) | |
| from steerable_retrieval.callbacks.energy import get_dictionary_from_lightningsae, load_state_dict_any # noqa: E402 | |
| from steerable_retrieval.experiments.common import ( # noqa: E402 | |
| ensure_dir, | |
| finish_wandb, | |
| instantiate_model_and_load, | |
| maybe_start_wandb, | |
| resolve_model_from_source, | |
| save_config, | |
| save_json, | |
| ) | |
| from steerable_retrieval.utils import RankedLogger, extras, register_resolvers # noqa: E402 | |
| log = RankedLogger(__name__, rank_zero_only=True) | |
| register_resolvers() | |
| def _stability_hungarian(a: torch.Tensor, b: torch.Tensor) -> float: | |
| from scipy.optimize import linear_sum_assignment | |
| sim = (a @ b.t()).detach().cpu().numpy() | |
| row_idx, col_idx = linear_sum_assignment(-sim) | |
| return float(sim[row_idx, col_idx].mean()) | |
| def _k_values(n_concepts: int, n_curve_points: int) -> List[int]: | |
| if n_concepts <= 1: | |
| return [1] | |
| out = [1] | |
| for i in range(max(2, n_curve_points)): | |
| frac = i / max(n_curve_points - 1, 1) | |
| k = int(round(math.exp(math.log(1) * (1 - frac) + math.log(n_concepts) * frac))) | |
| out.append(max(1, min(n_concepts, k))) | |
| out.append(n_concepts) | |
| return sorted(set(out)) | |
| def run_stability(cfg: DictConfig) -> Dict[str, Any]: | |
| if cfg.get("seed") is not None: | |
| L.seed_everything(int(cfg.seed), workers=True) | |
| source_run_id = resolve_model_from_source(cfg) | |
| output_dir = ensure_dir(str(cfg.output_dir)) | |
| save_config(cfg, output_dir) | |
| device = torch.device(cfg.get("device", "cuda:0") if torch.cuda.is_available() else "cpu") | |
| model, ckpt_path = instantiate_model_and_load(cfg, device=device) | |
| d0 = F.normalize(get_dictionary_from_lightningsae(model).to(device), dim=-1) | |
| n_concepts = int(d0.shape[0]) | |
| checkpoints = list(cfg.stability.checkpoint_paths) | |
| if not checkpoints: | |
| raise ValueError("stability.checkpoint_paths cannot be empty.") | |
| run = maybe_start_wandb( | |
| cfg=cfg, | |
| job_type="stability_processing", | |
| name=str(cfg.get("run_name", "stability_processing")), | |
| output_dir=output_dir, | |
| ) | |
| sorted_indices = torch.arange(n_concepts, device=device) | |
| k_values = _k_values(n_concepts=n_concepts, n_curve_points=int(cfg.stability.n_curve_points)) | |
| topk = cfg.stability.get("topk_by_energy") | |
| per_checkpoint = [] | |
| all_stabilities = [] | |
| all_topk = [] | |
| curves = [] | |
| for path in checkpoints: | |
| comparison = copy.deepcopy(model).to("cpu") | |
| comparison.load_state_dict(load_state_dict_any(path, map_location="cpu"), strict=False) | |
| comparison = comparison.to(device).eval() | |
| dk = F.normalize(get_dictionary_from_lightningsae(comparison).to(device), dim=-1) | |
| stab_full = _stability_hungarian(d0, dk) | |
| all_stabilities.append(stab_full) | |
| stab_topk = None | |
| if topk is not None: | |
| k = min(int(topk), n_concepts) | |
| top_idx = sorted_indices[:k] | |
| stab_topk = _stability_hungarian(d0[top_idx], dk[top_idx]) | |
| all_topk.append(stab_topk) | |
| curve_vals = [] | |
| for k in k_values: | |
| top_idx = sorted_indices[:k] | |
| curve_vals.append(_stability_hungarian(d0[top_idx], dk[top_idx])) | |
| curves.append(curve_vals) | |
| per_checkpoint.append( | |
| { | |
| "checkpoint_path": str(path), | |
| "stability_full": stab_full, | |
| "stability_topk": stab_topk, | |
| } | |
| ) | |
| mean_curve = [sum(vals) / len(vals) for vals in zip(*curves)] | |
| summary = { | |
| "experiment": "stability_processing", | |
| "config_source_wandb_run_id": source_run_id, | |
| "source_checkpoint": ckpt_path, | |
| "n_concepts": n_concepts, | |
| "per_checkpoint": per_checkpoint, | |
| "mean_stability_full": float(sum(all_stabilities) / len(all_stabilities)), | |
| "mean_stability_topk": float(sum(all_topk) / len(all_topk)) if all_topk else None, | |
| "k_values": k_values, | |
| "mean_stability_curve": mean_curve, | |
| } | |
| save_json(os.path.join(output_dir, "summary.json"), summary) | |
| if run is not None: | |
| import wandb | |
| table = wandb.Table(columns=["checkpoint_path", "stability_full", "stability_topk"]) | |
| for item in per_checkpoint: | |
| table.add_data(item["checkpoint_path"], item["stability_full"], item["stability_topk"]) | |
| wandb.log({"stability/per_checkpoint": table}) | |
| wandb.log({"stability/mean_full": summary["mean_stability_full"]}) | |
| if summary["mean_stability_topk"] is not None: | |
| wandb.log({"stability/mean_topk": summary["mean_stability_topk"]}) | |
| finish_wandb(run) | |
| return summary | |
| def main(cfg: DictConfig) -> Optional[Dict[str, Any]]: | |
| os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
| extras(cfg) | |
| return run_stability(cfg) | |
| if __name__ == "__main__": | |
| main() | |