social-data-attribution-w2/src/data_attribution/attribution/influence_simple.py

from __future__ import annotations

import argparse
import warnings
from argparse import Namespace
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from tqdm import tqdm

from data_attribution.attribution.index import _import_attributor


def compute_attribution(args: Namespace) -> None:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    Attributor = _import_attributor()
    print("index attributor", flush=True)
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="The given NumPy array is not writable",
            category=UserWarning,
        )
        index = Attributor(Path(args.index_dataset_path), device=device, unit_norm=True)
        print("queries attributor", flush=True)
        queries = Attributor(
            Path(args.query_dataset_path), device=device, unit_norm=True
        )

    module_names = list(queries.grads.keys())

    N = index.grads[module_names[0]].shape[0]
    M = queries.grads[module_names[0]].shape[0]

    q_chunk = 256
    index_scores_sum = torch.zeros(N, device=device, dtype=torch.float32)

    for q0 in tqdm(range(0, M, q_chunk), disable=not args.verbose):
        q1 = min(M, q0 + q_chunk)
        block = None
        for name in module_names:
            A = index.grads[name].to(torch.float32)
            B = queries.grads[name][q0:q1].to(torch.float32)
            prod = A @ B.T
            block = prod if block is None else (block + prod)

        index_scores_sum += torch.nan_to_num(block).sum(dim=1)

    scores = (index_scores_sum / M).detach().cpu().numpy()

    out_path = Path(args.attribution_path)
    out_path.mkdir(parents=True, exist_ok=True)

    pd.DataFrame(
        {
            "index_example_idx": np.arange(N),
            "attribution": scores,
        }
    ).to_csv(out_path / "attributions.csv", index=False)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Compute mean-aggregated influence scores."
    )
    parser.add_argument("--index_dataset_path", type=Path, required=True)
    parser.add_argument("--query_dataset_path", type=Path, required=True)
    parser.add_argument("--attribution_path", type=Path, required=True)
    parser.add_argument("--verbose", action="store_true", default=False)

    args = parser.parse_args()
    compute_attribution(args)


if __name__ == "__main__":
    main()