social-data-attribution-w2/scripts/validation/weborganizer_distribution.py

import io
import json
import sys
from collections import Counter, defaultdict
from pathlib import Path

import zstandard as zstd

INPUT_PATH = Path(
    "/storage/ice1/6/6/cchakraborty3/dolma/dolma3_pool_100K_docs.jsonl.zst"
)
OUTPUT_PATH = (
    Path(__file__).resolve().parent.parent
    / "outputs"
    / "weborganizer_distribution_100K_docs.json"
)


def compute_distributions(path: Path) -> dict:
    topic_max_counts: Counter[str] = Counter()
    topic_score_sums: dict[str, float] = defaultdict(float)
    topic_label_presence: dict[str, int] = defaultdict(int)
    format_max_counts: Counter[str] = Counter()
    format_score_sums: dict[str, float] = defaultdict(float)
    format_label_presence: dict[str, int] = defaultdict(int)
    total_records = 0
    records_with_topic = 0
    records_with_format = 0

    dctx = zstd.ZstdDecompressor()
    with open(path, "rb") as fh:
        with dctx.stream_reader(fh) as reader:
            text_stream = io.TextIOWrapper(reader, encoding="utf-8")
            for line in text_stream:
                record = json.loads(line)
                total_records += 1
                metadata = record.get("metadata", {})

                topic_scores = metadata.get("weborganizer_topic") or metadata.get(
                    "weborganizer"
                )
                topic_max = metadata.get("weborganizer_topic_max") or metadata.get(
                    "weborganizer_max"
                )

                if isinstance(topic_scores, dict) and topic_scores:
                    records_with_topic += 1
                    for label, score in topic_scores.items():
                        topic_score_sums[label] += score
                        topic_label_presence[label] += 1
                    if isinstance(topic_max, str):
                        topic_max_counts[topic_max] += 1

                format_scores = metadata.get("weborganizer_format")
                format_max = metadata.get("weborganizer_format_max")

                if isinstance(format_scores, dict) and format_scores:
                    records_with_format += 1
                    for label, score in format_scores.items():
                        format_score_sums[label] += score
                        format_label_presence[label] += 1
                    if isinstance(format_max, str):
                        format_max_counts[format_max] += 1

    def clean(label: str) -> str:
        return label.replace("__label__", "")

    def build_section(max_counts, score_sums, presence, n_records):
        avg_probs = {}
        for label in sorted(
            score_sums, key=lambda label_name: score_sums[label_name], reverse=True
        ):
            avg_probs[clean(label)] = score_sums[label] / n_records
        distribution = {clean(k): v for k, v in max_counts.most_common()}
        return {
            "distribution": distribution,
            "average_probabilities": avg_probs,
            "label_presence": {clean(k): v for k, v in presence.items()},
            "total_records": n_records,
        }

    result = {"total_records": total_records}

    if records_with_topic > 0:
        result["topic"] = build_section(
            topic_max_counts, topic_score_sums, topic_label_presence, records_with_topic
        )
    else:
        result["topic"] = None

    if records_with_format > 0:
        result["format"] = build_section(
            format_max_counts,
            format_score_sums,
            format_label_presence,
            records_with_format,
        )
    else:
        result["format"] = None

    return result


def main() -> None:
    if not INPUT_PATH.exists():
        print(f"File not found: {INPUT_PATH}", file=sys.stderr)
        sys.exit(1)

    result = compute_distributions(INPUT_PATH)

    print(f"Total records: {result['total_records']}\n")

    for domain in ("topic", "format"):
        section = result.get(domain)
        if section is None:
            print(f"No {domain} domain data found.\n")
            continue
        print(f"--- {domain.upper()} DOMAIN ---")
        print(f"Records with data: {section['total_records']}")
        print(f"\n{'Label':<45} {'Count':>8}  {'Avg Prob':>10}")
        print("-" * 68)
        dist = section["distribution"]
        avg = section["average_probabilities"]
        for label in dist:
            print(f"{label:<45} {dist[label]:>8}  {avg.get(label, 0):>10.6f}")
        print()

    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    OUTPUT_PATH.write_text(json.dumps(result, indent=2))
    print(f"Saved to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()