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

"""FastText-based WebOrganizer classifiers."""

from __future__ import annotations
import os
from huggingface_hub import hf_hub_download
import fasttext


class FastTextClassifier:
    def __init__(self, model_repo: str) -> None:
        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        model_path = hf_hub_download(
            repo_id=model_repo, filename="model.bin", token=token
        )
        self.model = fasttext.load_model(model_path)

    def predict_batch(
        self, urls: list[str | None], texts: list[str]
    ) -> tuple[list[dict[str, float]], list[str]]:
        prob_dicts = []
        max_labels = []
        for url, text in zip(urls, texts):
            # fastText does not use url
            text_clean = text.replace("\n", " ").strip()
            labels, probs = self.model.predict(text_clean, k=-1)
            prob_dict = {
                label.replace("__label__", ""): float(prob)
                for label, prob in zip(labels, probs)
            }
            max_label = max(prob_dict, key=prob_dict.get)
            prob_dicts.append(prob_dict)
            max_labels.append(max_label)
        return prob_dicts, max_labels

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


__all__ = ["FastTextClassifier"]