"""Transformers AutoModel entry point with text-level prediction helpers."""

from transformers import AutoTokenizer, ModernBertForTokenClassification

from .mixlid import MixLID


class MixLIDForTokenClassification(ModernBertForTokenClassification):
    def __init__(self, config):
        config.reference_compile = False
        super().__init__(config)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        tokenizer_kwargs = {
            key: kwargs[key]
            for key in (
                "cache_dir", "force_download", "local_files_only", "token",
                "revision", "subfolder",
            )
            if key in kwargs
        }
        loaded = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
        model = loaded[0] if isinstance(loaded, tuple) else loaded
        model.tokenizer = AutoTokenizer.from_pretrained(
            pretrained_model_name_or_path,
            use_fast=True,
            trust_remote_code=False,
            **tokenizer_kwargs,
        )
        if not model.tokenizer.is_fast:
            raise ValueError("MixLID requires a fast tokenizer for character offsets")
        return loaded

    def predict_sentence(self, texts):
        """Return one majority language per input string."""
        return MixLID(self, self.tokenizer).predict_sentence(texts)

    def predict_tokens(self, texts):
        """Return subword predictions and merged language spans."""
        return MixLID(self, self.tokenizer).predict_tokens(texts)

    def predict(self, texts, *, batch_size=8, max_length=8192, stride=64):
        """Return language, tokens, and spans for one or more input strings."""
        return MixLID(self, self.tokenizer).predict(
            texts, batch_size=batch_size, max_length=max_length, stride=stride
        )