social-data-attribution-w2/src/data_attribution/analysis/enrich.py

from __future__ import annotations

import argparse
from pathlib import Path
from typing import Sequence

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
from datasets import load_from_disk


def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Join attributions with document text and query prompts."
    )
    parser.add_argument("--attributions", type=Path, required=True)
    parser.add_argument("--doc-metadata", type=Path, required=True)
    parser.add_argument(
        "--query-index",
        type=Path,
        required=True,
        help="Bergson query index directory containing data.hf",
    )
    parser.add_argument(
        "--query-text-field",
        default="prompt",
        help="Field in query dataset to expose as query_prompt.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Output parquet; defaults to <attributions>-enriched.parquet",
    )
    return parser.parse_args(argv)


def _load_query_prompts(index_path: Path, text_field: str) -> pd.DataFrame:
    data_path = index_path / "data.hf"
    dataset = load_from_disk(str(data_path))
    split = next(iter(dataset.values())) if isinstance(dataset, dict) else dataset
    if text_field not in split.column_names:
        raise ValueError(f"Query text field '{text_field}' not found in {data_path}")
    return pd.DataFrame(
        {
            "query_id": split["query_id"],
            "query_prompt": split[text_field],
        }
    )


def _default_output(path: Path) -> Path:
    return path.with_name(f"{path.stem}-enriched{path.suffix}")


def enrich(
    attributions: Path,
    doc_metadata: Path,
    query_index: Path,
    query_text_field: str = "prompt",
    output: Path | None = None,
) -> Path:
    attr_df = pq.read_table(attributions).to_pandas()
    doc_df = pq.read_table(doc_metadata).to_pandas()
    if "doc_id" not in doc_df.columns:
        raise ValueError("doc-metadata must contain a doc_id column")

    doc_columns = {
        "text": "doc_text",
        "source_id": "doc_source_id",
        "metadata": "doc_metadata",
    }
    available = {k: v for k, v in doc_columns.items() if k in doc_df.columns}
    doc_df = doc_df.rename(columns=available)

    merged = attr_df.merge(doc_df, on="doc_id", how="left")

    query_df = _load_query_prompts(query_index, query_text_field)
    merged = merged.merge(query_df, on="query_id", how="left")

    output_path = output or _default_output(attributions)
    pq.write_table(pa.Table.from_pandas(merged), output_path)
    return output_path


def main(argv: Sequence[str] | None = None) -> int:
    args = _parse_args(argv)
    output_path = enrich(
        attributions=args.attributions,
        doc_metadata=args.doc_metadata,
        query_index=args.query_index,
        query_text_field=args.query_text_field,
        output=args.output,
    )
    print(output_path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())