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

"""Model wrapper for WebOrganizer format classification."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Iterable

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from dolma.labels import build_label_map


@dataclass(frozen=True)
class ModelConfig:
    model_name: str
    model_name_nourl: str
    device: str
    max_length: int
    torch_dtype: torch.dtype | None
    use_memory_efficient_attention: bool
    unpad_inputs: bool
    compile_model: bool


class FormatClassifier:
    def __init__(
        self,
        model_name: str = "WebOrganizer/FormatClassifier",
        model_name_nourl: str = "WebOrganizer/FormatClassifier-NoURL",
        device: str = "cuda",
        max_length: int = 1024,
        torch_dtype: torch.dtype | None = None,
        use_memory_efficient_attention: bool = True,
        unpad_inputs: bool = True,
        compile_model: bool = False,
        use_nourl_fallback: bool = False,
    ) -> None:
        self.config = ModelConfig(
            model_name=model_name,
            model_name_nourl=model_name_nourl,
            device=device,
            max_length=max_length,
            torch_dtype=torch_dtype,
            use_memory_efficient_attention=use_memory_efficient_attention,
            unpad_inputs=unpad_inputs,
            compile_model=compile_model,
        )
        self.use_nourl_fallback = use_nourl_fallback
        self.device = torch.device(device)
        self.tokenizer, self.model, self.label_map = self._load_model(
            self.config.model_name
        )
        self._nourl_tokenizer: AutoTokenizer | None = None
        self._nourl_model: AutoModelForSequenceClassification | None = None
        self._nourl_label_map: dict[int, str] | None = None

    def predict_batch(
        self, urls: list[str | None], texts: list[str]
    ) -> tuple[list[dict[str, float]], list[str]]:
        if len(urls) != len(texts):
            raise ValueError("urls and texts must be the same length")

        with_url = []
        with_url_indices = []
        without_url = []
        without_url_indices = []

        for idx, (url, text) in enumerate(zip(urls, texts, strict=True)):
            if url:
                with_url.append(self._build_input(url, text))
                with_url_indices.append(idx)
            else:
                without_url.append(text)
                without_url_indices.append(idx)

        prob_dicts: list[dict[str, float]] = [{} for _ in texts]
        max_labels: list[str] = [""] * len(texts)

        if with_url:
            outputs = self._predict_with_model(
                self.tokenizer, self.model, self.label_map, with_url
            )
            for idx, (prob_dict, max_label) in zip(
                with_url_indices, outputs, strict=True
            ):
                prob_dicts[idx] = prob_dict
                max_labels[idx] = max_label

        if without_url:
            if self.use_nourl_fallback:
                tokenizer, model, label_map = self._ensure_nourl_model()
                outputs = self._predict_with_model(
                    tokenizer, model, label_map, without_url
                )
                for idx, (prob_dict, max_label) in zip(
                    without_url_indices, outputs, strict=True
                ):
                    prob_dicts[idx] = prob_dict
                    max_labels[idx] = max_label
            else:
                inputs = [self._build_input("", text) for text in without_url]
                outputs = self._predict_with_model(
                    self.tokenizer, self.model, self.label_map, inputs
                )
                for idx, (prob_dict, max_label) in zip(
                    without_url_indices, outputs, strict=True
                ):
                    prob_dicts[idx] = prob_dict
                    max_labels[idx] = max_label

        return prob_dicts, max_labels

    def estimate_tokens(self, url: str | None, text: str) -> int:
        if url:
            payload = self._build_input(url, text)
        elif self.use_nourl_fallback and self._nourl_tokenizer is not None:
            payload = text
        else:
            payload = self._build_input("", text)
        tokens = self.tokenizer(
            payload, truncation=True, max_length=self.config.max_length
        )
        return len(tokens["input_ids"])

    def _ensure_nourl_model(
        self,
    ) -> tuple[AutoTokenizer, AutoModelForSequenceClassification, dict[int, str]]:
        if self._nourl_model is None:
            tokenizer, model, label_map = self._load_model(self.config.model_name_nourl)
            self._nourl_tokenizer = tokenizer
            self._nourl_model = model
            self._nourl_label_map = label_map
        return self._nourl_tokenizer, self._nourl_model, self._nourl_label_map  # type: ignore[return-value]

    def _load_model(
        self, model_name: str
    ) -> tuple[AutoTokenizer, AutoModelForSequenceClassification, dict[int, str]]:
        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
        model = AutoModelForSequenceClassification.from_pretrained(
            model_name,
            trust_remote_code=True,
            use_memory_efficient_attention=self.config.use_memory_efficient_attention,
            unpad_inputs=self.config.unpad_inputs,
            torch_dtype=self.config.torch_dtype,
        )
        # When memory-efficient attention (xformers) is disabled, unpadding
        # must also be disabled — the unpad→pad→attend→unpad path is fragile
        # across transformers versions and triggers CUDA asserts with SDPA.
        if not self.config.use_memory_efficient_attention and model.config.unpad_inputs:
            model.config.unpad_inputs = False
        model = model.to(self.device)
        model.eval()
        if self.config.compile_model:
            model = torch.compile(model)
        label_map = build_label_map(model.config.id2label)
        return tokenizer, model, label_map

    def _predict_with_model(
        self,
        tokenizer: AutoTokenizer,
        model: AutoModelForSequenceClassification,
        label_map: dict[int, str],
        inputs: Iterable[str],
    ) -> list[tuple[dict[str, float], str]]:
        batch = tokenizer(
            list(inputs),
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.config.max_length,
        )
        batch = {key: value.to(self.device) for key, value in batch.items()}
        with torch.inference_mode():
            outputs = model(**batch)
        probabilities = torch.softmax(outputs.logits, dim=-1)
        label_indices = sorted(label_map)
        results = []
        for row in probabilities:
            prob_dict = {
                label_map[idx]: float(row[idx].item()) for idx in label_indices
            }
            max_idx = int(torch.argmax(row).item())
            results.append((prob_dict, label_map[max_idx]))
        return results

    @staticmethod
    def _build_input(url: str, text: str) -> str:
        return f"{url}\n\n{text}"


__all__ = ["FormatClassifier"]