social-data-attribution-w2/src/data_attribution/cache/policy.py

from __future__ import annotations

from collections import OrderedDict
from typing import Sequence

from data_attribution.config import ALLOWED_MODEL_IDS, DEFAULT_MODEL_ID
from dolma.constants import (
    ALLOWED_DOLMA_DATASET_IDS,
    ALLOWED_DOLMA_SUBSETS,
)


def normalize_model_ids(model_ids: Sequence[str] | None) -> list[str]:
    if not model_ids:
        return [DEFAULT_MODEL_ID]
    normalized = list(OrderedDict.fromkeys(model_ids))
    invalid = [model_id for model_id in normalized if model_id not in ALLOWED_MODEL_IDS]
    if invalid:
        raise ValueError(
            f"Unknown model id(s): {', '.join(invalid)}. "
            f"Allowed: {', '.join(ALLOWED_MODEL_IDS)}."
        )
    return normalized


def validate_dataset_id(dataset_id: str) -> str:
    if dataset_id not in ALLOWED_DOLMA_DATASET_IDS:
        raise ValueError(
            f"Unknown dataset id '{dataset_id}'. "
            f"Allowed: {', '.join(ALLOWED_DOLMA_DATASET_IDS)}."
        )
    return dataset_id


def validate_subset(subset: str | None) -> str | None:
    if subset is None:
        return None
    if subset not in ALLOWED_DOLMA_SUBSETS:
        allowed = ", ".join(ALLOWED_DOLMA_SUBSETS) or "none"
        raise ValueError(f"Subset '{subset}' is not allowed. Allowed: {allowed}.")
    return subset


__all__ = ["normalize_model_ids", "validate_dataset_id", "validate_subset"]