social-data-attribution-w2/src/dolma/paths.py

"""Path helpers for enrichment outputs."""

from __future__ import annotations

import hashlib
from glob import glob
from pathlib import Path
from typing import Iterable


def expand_inputs(patterns: Iterable[str]) -> list[Path]:
    expanded: list[Path] = []
    for pattern in patterns:
        matches = sorted(Path(entry) for entry in glob(pattern, recursive=True))
        if not matches:
            matches = [Path(pattern)]
        expanded.extend(matches)
    return expanded


def output_path(output_dir: Path, name: str | Path, compress: str | None) -> Path:
    base = slugify(str(name))
    base = strip_suffix(base)
    suffix = ".jsonl" if compress in {None, "none"} else ".jsonl.zst"
    return output_dir / f"{base}.enriched{suffix}"


def slugify(value: str) -> str:
    return value.replace("/", "_").replace("..", "__")


def strip_suffix(name: str) -> str:
    for suffix in (".jsonl.zst", ".jsonl", ".zst"):
        if name.endswith(suffix):
            return name[: -len(suffix)]
    return name


def hf_output_name(dataset: str, data_files: list[str] | None) -> str:
    if not data_files:
        return slugify(dataset)
    if len(data_files) == 1:
        return slugify(data_files[0])
    combined = "__".join(slugify(entry) for entry in data_files)
    if len(combined) > 180:
        digest = hashlib.sha1(combined.encode("utf-8")).hexdigest()
        return f"multi_{digest}"
    return combined


__all__ = [
    "expand_inputs",
    "hf_output_name",
    "output_path",
    "slugify",
    "strip_suffix",
]