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

"""Mock format classifier for tests and smoke runs."""

from __future__ import annotations


class MockFormatClassifier:
    def __init__(self, num_labels: int = 24) -> None:
        self.labels = [f"__label__mock_{idx}" for idx in range(num_labels)]

    def predict_batch(
        self, urls: list[str | None], texts: list[str]
    ) -> tuple[list[dict[str, float]], list[str]]:
        prob_dicts = []
        max_labels = []
        for idx, _text in enumerate(texts):
            base = idx % len(self.labels)
            probs = [0.0] * len(self.labels)
            probs[base] = 1.0
            prob_dict = {
                label: float(prob)
                for label, prob in zip(self.labels, probs, strict=True)
            }
            prob_dicts.append(prob_dict)
            max_labels.append(self.labels[base])
        return prob_dicts, max_labels

    def estimate_tokens(self, url: str | None, text: str) -> int:
        return len(text.split())


__all__ = ["MockFormatClassifier"]