social-data-attribution-w2/scripts/query_data/merge_scores.py

import argparse
import json
from pathlib import Path

import numpy as np
import torch


def load_scores(score_dir: Path) -> tuple[torch.Tensor, int]:
    with open(score_dir / "info.json", encoding="utf-8") as f:
        info = json.load(f)

    if info["num_scores"] != 1:
        raise ValueError(
            f"{score_dir} contains {info['num_scores']} score columns; expected 1"
        )

    dtype = np.dtype(info["dtype"])
    mmap = np.memmap(
        score_dir / "scores.bin", dtype=dtype, mode="r", shape=(info["num_items"],)
    )
    scores = torch.from_numpy(np.array(mmap["score_0"], dtype=np.float32, copy=True))
    return scores, info["num_items"]


def main():
    parser = argparse.ArgumentParser(
        description="Merge multiple bergson score directories and print global top-k."
    )
    parser.add_argument(
        "score_dirs", nargs="+", type=Path, help="Score directories to merge"
    )
    parser.add_argument("--top-k", type=int, default=5)
    parser.add_argument("--preview-chars", type=int, default=500)
    args = parser.parse_args()

    all_scores = []
    offsets = []
    sources = []
    offset = 0

    for d in args.score_dirs:
        scores, n = load_scores(d)
        all_scores.append(scores)
        offsets.append(offset)
        sources.append(d)
        with open(d / "index_config.json") as f:
            cfg = json.load(f)
        jsonl = cfg["data"].get("data_args", "")
        print(f"Loaded {d.name}: {n} docs from {jsonl}")
        offset += n

    combined = torch.cat(all_scores)
    topk = torch.topk(combined, k=min(args.top_k, len(combined)))

    print(f"\nTotal documents: {len(combined)}")
    print(f"Top {args.top_k} most influential:\n")

    for rank, (global_idx, score) in enumerate(
        zip(topk.indices.tolist(), topk.values.tolist()), 1
    ):
        src_name = None
        local_idx = global_idx
        for i, (s, o) in enumerate(zip(sources, offsets)):
            next_o = offsets[i + 1] if i + 1 < len(offsets) else len(combined)
            if global_idx < next_o:
                src_name = s.name
                local_idx = global_idx - o
                break

        text = ""
        data_hf = sources[i] / "data.hf"
        if data_hf.exists():
            from datasets import load_from_disk

            ds = load_from_disk(str(data_hf))
            text = ds[local_idx].get("text", "")[: args.preview_chars]

        print(f"[{rank}] source={src_name}  local_idx={local_idx}  score={score:.6f}")
        if text:
            print(f"    {text}")
        print()


if __name__ == "__main__":
    main()