social-data-attribution-w2/src/dolma/eda/state.py

"""Aggregation state for Dolma EDA."""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass, field
from typing import Any, Iterable

from dolma.eda.stats import LengthStats

JOINT_SEPARATOR = "\t"


@dataclass
class EdaState:
    total_records: int = 0
    missing_metadata: int = 0
    topic_counts: Counter[str] = field(default_factory=Counter)
    topic_score_sums: Counter[str] = field(default_factory=Counter)
    topic_word_sums: Counter[str] = field(default_factory=Counter)
    topic_word_counts: Counter[str] = field(default_factory=Counter)
    format_counts: Counter[str] = field(default_factory=Counter)
    format_score_sums: Counter[str] = field(default_factory=Counter)
    format_word_sums: Counter[str] = field(default_factory=Counter)
    format_word_counts: Counter[str] = field(default_factory=Counter)
    joint_counts: Counter[str] = field(default_factory=Counter)
    joint_word_sums: Counter[str] = field(default_factory=Counter)
    joint_word_counts: Counter[str] = field(default_factory=Counter)
    metadata_present_counts: Counter[str] = field(default_factory=Counter)
    metadata_value_counts: dict[str, Counter[str]] = field(default_factory=dict)
    length_stats: LengthStats = field(default_factory=LengthStats)

    def update_record(
        self, record: dict[str, Any], value_fields: Iterable[str]
    ) -> None:
        self.total_records += 1
        metadata = record.get("metadata")
        if not isinstance(metadata, dict):
            self.missing_metadata += 1
            self.length_stats.update_value(None)
            return
        self._update_metadata_counts(metadata, value_fields)
        word_count = self._extract_word_count(metadata)
        self._update_labels(metadata, word_count)
        self._update_length(metadata)

    def _update_metadata_counts(
        self, metadata: dict[str, Any], value_fields: Iterable[str]
    ) -> None:
        for key in metadata.keys():
            self.metadata_present_counts[key] += 1
        for field_name in value_fields:
            counter = self.metadata_value_counts.setdefault(field_name, Counter())
            if field_name in metadata and metadata[field_name] is not None:
                counter[str(metadata[field_name])] += 1

    def _update_labels(self, metadata: dict[str, Any], word_count: int | None) -> None:
        topic_scores = metadata.get("weborganizer_topic")
        format_scores = metadata.get("weborganizer_format")
        if isinstance(topic_scores, dict):
            for label, score in topic_scores.items():
                self.topic_score_sums[str(label)] += float(score)
        if isinstance(format_scores, dict):
            for label, score in format_scores.items():
                self.format_score_sums[str(label)] += float(score)
        topic_max = metadata.get("weborganizer_topic_max")
        format_max = metadata.get("weborganizer_format_max")
        if isinstance(topic_max, str):
            self.topic_counts[topic_max] += 1
            if word_count is not None:
                self.topic_word_sums[topic_max] += word_count
                self.topic_word_counts[topic_max] += 1
        if isinstance(format_max, str):
            self.format_counts[format_max] += 1
            if word_count is not None:
                self.format_word_sums[format_max] += word_count
                self.format_word_counts[format_max] += 1
        if isinstance(topic_max, str) and isinstance(format_max, str):
            joint_key = f"{topic_max}{JOINT_SEPARATOR}{format_max}"
            self.joint_counts[joint_key] += 1
            if word_count is not None:
                self.joint_word_sums[joint_key] += word_count
                self.joint_word_counts[joint_key] += 1

    def _update_length(self, metadata: dict[str, Any]) -> None:
        value = metadata.get("original_word_count")
        if isinstance(value, (int, float)):
            self.length_stats.update_value(value)
        else:
            self.length_stats.update_value(None)

    def _extract_word_count(self, metadata: dict[str, Any]) -> int | None:
        value = metadata.get("original_word_count")
        if isinstance(value, (int, float)):
            return int(value)
        return None

    def to_dict(self) -> dict[str, Any]:
        return {
            "total_records": self.total_records,
            "missing_metadata": self.missing_metadata,
            "topic_counts": dict(self.topic_counts),
            "topic_score_sums": dict(self.topic_score_sums),
            "topic_word_sums": dict(self.topic_word_sums),
            "topic_word_counts": dict(self.topic_word_counts),
            "format_counts": dict(self.format_counts),
            "format_score_sums": dict(self.format_score_sums),
            "format_word_sums": dict(self.format_word_sums),
            "format_word_counts": dict(self.format_word_counts),
            "joint_counts": dict(self.joint_counts),
            "joint_word_sums": dict(self.joint_word_sums),
            "joint_word_counts": dict(self.joint_word_counts),
            "metadata_present_counts": dict(self.metadata_present_counts),
            "metadata_value_counts": {
                key: dict(value) for key, value in self.metadata_value_counts.items()
            },
            "length_stats": self.length_stats.to_dict(),
        }

    def merge(self, other: "EdaState") -> None:
        self.total_records += other.total_records
        self.missing_metadata += other.missing_metadata
        self.topic_counts.update(other.topic_counts)
        self.topic_score_sums.update(other.topic_score_sums)
        self.topic_word_sums.update(other.topic_word_sums)
        self.topic_word_counts.update(other.topic_word_counts)
        self.format_counts.update(other.format_counts)
        self.format_score_sums.update(other.format_score_sums)
        self.format_word_sums.update(other.format_word_sums)
        self.format_word_counts.update(other.format_word_counts)
        self.joint_counts.update(other.joint_counts)
        self.joint_word_sums.update(other.joint_word_sums)
        self.joint_word_counts.update(other.joint_word_counts)
        self.metadata_present_counts.update(other.metadata_present_counts)
        for key, counter in other.metadata_value_counts.items():
            self.metadata_value_counts.setdefault(key, Counter()).update(counter)
        self.length_stats.merge(other.length_stats)

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "EdaState":
        state = cls(
            total_records=payload.get("total_records", 0),
            missing_metadata=payload.get("missing_metadata", 0),
            topic_counts=Counter(payload.get("topic_counts", {})),
            topic_score_sums=Counter(payload.get("topic_score_sums", {})),
            topic_word_sums=Counter(payload.get("topic_word_sums", {})),
            topic_word_counts=Counter(payload.get("topic_word_counts", {})),
            format_counts=Counter(payload.get("format_counts", {})),
            format_score_sums=Counter(payload.get("format_score_sums", {})),
            format_word_sums=Counter(payload.get("format_word_sums", {})),
            format_word_counts=Counter(payload.get("format_word_counts", {})),
            joint_counts=Counter(payload.get("joint_counts", {})),
            joint_word_sums=Counter(payload.get("joint_word_sums", {})),
            joint_word_counts=Counter(payload.get("joint_word_counts", {})),
            metadata_present_counts=Counter(payload.get("metadata_present_counts", {})),
            length_stats=LengthStats.from_dict(payload["length_stats"])
            if "length_stats" in payload
            else LengthStats(),
        )
        state.metadata_value_counts = {
            key: Counter(value)
            for key, value in payload.get("metadata_value_counts", {}).items()
        }
        return state


__all__ = [
    "JOINT_SEPARATOR",
    "EdaState",
    "LengthStats",
]