Spaces:
Running
Running
| """Compute retrieval metrics from a result pickle.""" | |
| import argparse | |
| import os | |
| import pickle | |
| import sys | |
| from pathlib import Path | |
| def _add_local_dependency_paths() -> None: | |
| """Ensure ``flare`` and ``massspecgym`` are importable when run as a script.""" | |
| flare_repo_root = Path(__file__).resolve().parents[1] | |
| candidate_roots = [flare_repo_root, flare_repo_root / "massspecgym"] | |
| env_value = os.environ.get("MASSSPECGYM_ROOT") | |
| if env_value: | |
| candidate_roots.append(Path(env_value).expanduser()) | |
| for root in candidate_roots: | |
| if root.exists(): | |
| root_str = str(root) | |
| if root_str not in sys.path: | |
| sys.path.insert(0, root_str) | |
| _add_local_dependency_paths() | |
| import numpy as np | |
| from tqdm import tqdm | |
| def rank_at_k(ranks, k: int) -> float: | |
| return (ranks <= k).mean() | |
| def confidence_interval( | |
| values, | |
| confidence_level: float = 0.999, | |
| n_resamples: int = 20_000, | |
| seed: int = 0, | |
| ) -> str: | |
| """Return a bootstrap confidence interval formatted as in msgym_result.ipynb.""" | |
| from scipy.stats import bootstrap | |
| ci = bootstrap( | |
| (np.asarray(values),), | |
| np.mean, | |
| confidence_level=confidence_level, | |
| n_resamples=n_resamples, | |
| random_state=seed, | |
| ).confidence_interval | |
| return f"{ci.low:.2f}-{ci.high:.2f}" | |
| def compute_mces(result, n_processes: int): | |
| """Compute the per-row MCES between the target and top-ranked candidate.""" | |
| # Lazy import: needs massspecgym / torchmetrics, only required for --mces. | |
| print("Importing MCES helpers (flare.utils.eval)...") | |
| from flare.utils import eval as eval_utils | |
| required = {"candidates", "labels", "scores"} | |
| missing = required.difference(result.columns) | |
| if missing: | |
| raise SystemExit( | |
| f"MCES requires columns {sorted(required)}; missing: {sorted(missing)}" | |
| ) | |
| print("Extracting targets and top-1 candidates...") | |
| targets = [ | |
| eval_utils.get_target(candidates, labels) | |
| for candidates, labels in tqdm( | |
| zip(result["candidates"], result["labels"]), | |
| total=len(result), | |
| desc="targets", | |
| ) | |
| ] | |
| top_candidates = [ | |
| eval_utils.get_top_cand(candidates, scores) | |
| for candidates, scores in tqdm( | |
| zip(result["candidates"], result["scores"]), | |
| total=len(result), | |
| desc="cand@1", | |
| ) | |
| ] | |
| pairs = list(zip(targets, top_candidates)) | |
| unique_pairs = list(dict.fromkeys(pairs)) | |
| print( | |
| f"Computing MCES for {len(unique_pairs)} unique target/candidate pairs " | |
| f"({n_processes} processes)..." | |
| ) | |
| computed = eval_utils.Compute_Myopic_MCES.compute_mces_parallel( | |
| unique_pairs, n_processes=n_processes | |
| ) | |
| print("Assigning MCES values to rows...") | |
| mces_by_pair = {tuple(pair): mces for pair, mces in computed} | |
| return np.asarray([mces_by_pair[pair] for pair in pairs], dtype=float) | |
| def main(): | |
| parser = argparse.ArgumentParser( | |
| description="Compute rank@k and optional CI/MCES metrics from a result pickle." | |
| ) | |
| parser.add_argument("result_file", type=Path, help="Path to result_*.pkl") | |
| parser.add_argument( | |
| "--rank_col", | |
| default="rank", | |
| help="Column with ranks (default: rank)", | |
| ) | |
| parser.add_argument( | |
| "--ks", | |
| nargs="+", | |
| type=int, | |
| default=[1, 5, 20], | |
| help="k values for rank@k (default: 1 5 20)", | |
| ) | |
| parser.add_argument( | |
| "--ci", | |
| "--compute_ci", | |
| dest="compute_ci", | |
| action="store_true", | |
| help="Compute bootstrap confidence intervals (hit rates are percentages)", | |
| ) | |
| parser.add_argument( | |
| "--mces", | |
| "--compute_mces", | |
| dest="compute_mces", | |
| action="store_true", | |
| help="Compute the mean MCES for each target/top-candidate pair", | |
| ) | |
| parser.add_argument( | |
| "--confidence_level", | |
| type=float, | |
| default=0.999, | |
| help="Bootstrap confidence level (default: 0.999)", | |
| ) | |
| parser.add_argument( | |
| "--n_resamples", | |
| type=int, | |
| default=20_000, | |
| help="Number of bootstrap resamples (default: 20000)", | |
| ) | |
| parser.add_argument( | |
| "--seed", | |
| type=int, | |
| default=0, | |
| help="Bootstrap random seed (default: 0)", | |
| ) | |
| parser.add_argument( | |
| "--mces_processes", | |
| type=int, | |
| default=25, | |
| help="Number of MCES worker processes (default: 25)", | |
| ) | |
| args = parser.parse_args() | |
| result_path = args.result_file.resolve() | |
| print(f"Loading result file: {result_path}") | |
| with open(result_path, "rb") as f: | |
| result = pickle.load(f) | |
| print(f"Loaded {len(result)} rows; columns={list(result.columns)}") | |
| if args.rank_col not in result.columns: | |
| raise SystemExit( | |
| f"Column '{args.rank_col}' not found. Available: {list(result.columns)}" | |
| ) | |
| ranks = result[args.rank_col] | |
| if len(ranks) == 0: | |
| raise SystemExit("The result file contains no rows.") | |
| print(f"Computing rank@k for ks={args.ks} (rank_col={args.rank_col})...") | |
| lines = [f"n={len(ranks)}", f"rank_col={args.rank_col}"] | |
| for k in args.ks: | |
| val = rank_at_k(ranks, k) | |
| line = f"rank@{k}: {val:.6f}" | |
| lines.append(line) | |
| print(line) | |
| mces_values = None | |
| if args.compute_mces: | |
| print("MCES flag set; starting MCES computation...") | |
| mces_values = compute_mces(result, args.mces_processes) | |
| line = f"mces: {mces_values.mean():.6f}" | |
| lines.append(line) | |
| print(line) | |
| elif "mces" in result.columns: | |
| print("Using precomputed 'mces' column from result file.") | |
| mces_values = result["mces"].to_numpy(dtype=float) | |
| if args.compute_ci: | |
| ci_label = f"{args.confidence_level * 100:g}%_ci" | |
| print( | |
| f"Computing bootstrap CIs " | |
| f"(level={args.confidence_level}, n_resamples={args.n_resamples}, seed={args.seed})..." | |
| ) | |
| ci_jobs = [] | |
| for k in args.ks: | |
| # The notebook reports hit-rate confidence intervals on a 0-100 scale. | |
| hit_rate_percent = (ranks <= k).to_numpy(dtype=float) * 100 | |
| ci_jobs.append((f"rank@{k}_{ci_label}_percent", hit_rate_percent)) | |
| if mces_values is not None: | |
| ci_jobs.append((f"mces_{ci_label}", mces_values)) | |
| else: | |
| print("No MCES values available; skipping MCES CI.") | |
| for name, values in tqdm(ci_jobs, desc="bootstrap CI"): | |
| print(f" Bootstrapping {name}...") | |
| interval = confidence_interval( | |
| values, | |
| confidence_level=args.confidence_level, | |
| n_resamples=args.n_resamples, | |
| seed=args.seed, | |
| ) | |
| line = f"{name}: {interval}" | |
| lines.append(line) | |
| print(line) | |
| out_path = result_path.with_name(result_path.stem + "_rank_at_k.txt") | |
| print(f"Writing metrics to {out_path}...") | |
| out_path.write_text("\n".join(lines) + "\n") | |
| print(f"Done. Wrote {out_path}") | |
| if __name__ == "__main__": | |
| main() | |