social-data-attribution-w2/src/unlearning/data/forget_texts.py

"""Load fixed text forget sets for unlearning."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path

import pandas as pd

TEXT_COL = "text"
DOC_ID_COLS = ("doc_id", "id")


@dataclass(frozen=True)
class TextForgetSet:
    texts: list[str]
    doc_ids: list[str]
    rows_read: int
    rows_kept: int


def load_text_forget_set(path: str | Path) -> TextForgetSet:
    """Read a parquet, CSV, or JSONL forget set with a required text column."""
    source = Path(path).expanduser()
    if not source.exists():
        raise FileNotFoundError(f"forget text set not found: {source}")
    frame = _read_frame(source)
    if TEXT_COL not in frame.columns:
        raise ValueError(f"forget text set must contain a '{TEXT_COL}' column: {source}")
    doc_id_col = next((col for col in DOC_ID_COLS if col in frame.columns), None)
    rows_read = len(frame)
    if doc_id_col is not None:
        frame = frame.drop_duplicates(subset=[doc_id_col], keep="first")
    text = frame[TEXT_COL].fillna("").astype(str)
    keep = text.str.strip() != ""
    frame = frame.loc[keep].copy()
    texts = frame[TEXT_COL].astype(str).tolist()
    doc_ids = (
        frame[doc_id_col].fillna("").astype(str).tolist()
        if doc_id_col is not None
        else []
    )
    if not texts:
        raise ValueError(f"forget text set has no non-empty texts: {source}")
    return TextForgetSet(
        texts=texts,
        doc_ids=doc_ids,
        rows_read=rows_read,
        rows_kept=len(texts),
    )


def _read_frame(path: Path) -> pd.DataFrame:
    suffix = path.suffix.lower()
    if suffix == ".parquet":
        return pd.read_parquet(path)
    if suffix == ".csv":
        return pd.read_csv(path)
    if suffix in {".jsonl", ".json"}:
        return _read_json_lines(path)
    raise ValueError(f"unsupported forget text set format: {path}")


def _read_json_lines(path: Path) -> pd.DataFrame:
    rows: list[dict[str, object]] = []
    with path.open(encoding="utf-8") as handle:
        for line in handle:
            if line.strip():
                rows.append(json.loads(line))
    return pd.DataFrame(rows)