/headroom/tokenizers/huggingface.py

"""HuggingFace tokenizer wrapper for open models.

Supports Llama, Mistral, Falcon, and other models with HuggingFace
tokenizers. Requires the `transformers` library.
"""

from __future__ import annotations

import logging
import os
import threading
from functools import lru_cache
from typing import Any

from .base import BaseTokenizer

logger = logging.getLogger(__name__)


# Model name to HuggingFace tokenizer mapping
# Maps common model names to their HuggingFace tokenizer identifiers
MODEL_TO_TOKENIZER: dict[str, str] = {
    # Llama 3 family
    "llama-3": "meta-llama/Meta-Llama-3-8B",
    "llama-3-8b": "meta-llama/Meta-Llama-3-8B",
    "llama-3-70b": "meta-llama/Meta-Llama-3-70B",
    "llama-3.1-8b": "meta-llama/Llama-3.1-8B",
    "llama-3.1-70b": "meta-llama/Llama-3.1-70B",
    "llama-3.1-405b": "meta-llama/Llama-3.1-405B",
    "llama-3.2-1b": "meta-llama/Llama-3.2-1B",
    "llama-3.2-3b": "meta-llama/Llama-3.2-3B",
    "llama-3.3-70b": "meta-llama/Llama-3.3-70B-Instruct",
    # Llama 2 family
    "llama-2": "meta-llama/Llama-2-7b-hf",
    "llama-2-7b": "meta-llama/Llama-2-7b-hf",
    "llama-2-13b": "meta-llama/Llama-2-13b-hf",
    "llama-2-70b": "meta-llama/Llama-2-70b-hf",
    # CodeLlama
    "codellama": "codellama/CodeLlama-7b-hf",
    "codellama-7b": "codellama/CodeLlama-7b-hf",
    "codellama-13b": "codellama/CodeLlama-13b-hf",
    "codellama-34b": "codellama/CodeLlama-34b-hf",
    # Mistral family
    "mistral": "mistralai/Mistral-7B-v0.1",
    "mistral-7b": "mistralai/Mistral-7B-v0.1",
    "mistral-7b-v0.2": "mistralai/Mistral-7B-Instruct-v0.2",
    "mistral-7b-v0.3": "mistralai/Mistral-7B-Instruct-v0.3",
    "mistral-nemo": "mistralai/Mistral-Nemo-Base-2407",
    "mistral-small": "mistralai/Mistral-Small-Instruct-2409",
    "mistral-large": "mistralai/Mistral-Large-Instruct-2407",
    # Mixtral
    "mixtral": "mistralai/Mixtral-8x7B-v0.1",
    "mixtral-8x7b": "mistralai/Mixtral-8x7B-v0.1",
    "mixtral-8x22b": "mistralai/Mixtral-8x22B-v0.1",
    # Qwen family
    "qwen": "Qwen/Qwen-7B",
    "qwen-7b": "Qwen/Qwen-7B",
    "qwen-14b": "Qwen/Qwen-14B",
    "qwen-72b": "Qwen/Qwen-72B",
    "qwen2": "Qwen/Qwen2-7B",
    "qwen2-7b": "Qwen/Qwen2-7B",
    "qwen2-72b": "Qwen/Qwen2-72B",
    "qwen2.5": "Qwen/Qwen2.5-7B",
    "qwen2.5-7b": "Qwen/Qwen2.5-7B",
    "qwen2.5-72b": "Qwen/Qwen2.5-72B",
    # DeepSeek
    "deepseek": "deepseek-ai/deepseek-llm-7b-base",
    "deepseek-7b": "deepseek-ai/deepseek-llm-7b-base",
    "deepseek-67b": "deepseek-ai/deepseek-llm-67b-base",
    "deepseek-coder": "deepseek-ai/deepseek-coder-6.7b-base",
    "deepseek-v2": "deepseek-ai/DeepSeek-V2",
    "deepseek-v3": "deepseek-ai/DeepSeek-V3",
    # Yi family
    "yi": "01-ai/Yi-6B",
    "yi-6b": "01-ai/Yi-6B",
    "yi-34b": "01-ai/Yi-34B",
    "yi-1.5": "01-ai/Yi-1.5-6B",
    # Phi family
    "phi-2": "microsoft/phi-2",
    "phi-3": "microsoft/Phi-3-mini-4k-instruct",
    "phi-3-mini": "microsoft/Phi-3-mini-4k-instruct",
    "phi-3-small": "microsoft/Phi-3-small-8k-instruct",
    "phi-3-medium": "microsoft/Phi-3-medium-4k-instruct",
    # Falcon
    "falcon": "tiiuae/falcon-7b",
    "falcon-7b": "tiiuae/falcon-7b",
    "falcon-40b": "tiiuae/falcon-40b",
    "falcon-180b": "tiiuae/falcon-180B",
    # StarCoder
    "starcoder": "bigcode/starcoder",
    "starcoder2": "bigcode/starcoder2-15b",
    "starcoder2-3b": "bigcode/starcoder2-3b",
    "starcoder2-7b": "bigcode/starcoder2-7b",
    "starcoder2-15b": "bigcode/starcoder2-15b",
    # MPT
    "mpt-7b": "mosaicml/mpt-7b",
    "mpt-30b": "mosaicml/mpt-30b",
    # Gemma
    "gemma": "google/gemma-7b",
    "gemma-2b": "google/gemma-2b",
    "gemma-7b": "google/gemma-7b",
    "gemma-2": "google/gemma-2-9b",
    "gemma-2-9b": "google/gemma-2-9b",
    "gemma-2-27b": "google/gemma-2-27b",
}


# Bound the first (network) load of a HuggingFace tokenizer. Without a bound,
# huggingface_hub download retries can block for many minutes (GH #1701: 610s
# on a restricted Windows network). 0 disables network loads entirely.
_LOAD_TIMEOUT_ENV = "HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS"
_LOAD_TIMEOUT_DEFAULT = 10.0


def _load_timeout_secs() -> float:
    try:
        return float(os.environ.get(_LOAD_TIMEOUT_ENV, _LOAD_TIMEOUT_DEFAULT))
    except (TypeError, ValueError):
        return _LOAD_TIMEOUT_DEFAULT


@lru_cache(maxsize=16)
def _load_tokenizer(tokenizer_name: str):
    """Load and cache HuggingFace tokenizer.

    The first attempt is cache-only (``local_files_only=True``) so a warm
    HF cache never touches the network. A cache miss falls through to a
    network download bounded by ``HEADROOM_HF_TOKENIZER_LOAD_TIMEOUT_SECS``
    (default 10s) on a daemon thread — the download itself cannot be
    cancelled, but the caller unblocks and falls back to estimation.
    Failures are cached by ``lru_cache`` (returns ``None``), so a slow or
    offline hub is probed at most once per process per tokenizer.

    Args:
        tokenizer_name: HuggingFace model/tokenizer name.

    Returns:
        Loaded tokenizer, or None if unavailable.
    """
    from transformers import AutoTokenizer

    try:
        return AutoTokenizer.from_pretrained(
            tokenizer_name,
            trust_remote_code=True,
            local_files_only=True,
        )
    except Exception:
        pass  # Not in the local cache — try the network below, bounded.

    timeout = _load_timeout_secs()
    if timeout <= 0:
        logger.warning(
            f"Tokenizer {tokenizer_name} not in local HF cache and network "
            f"loading is disabled ({_LOAD_TIMEOUT_ENV}=0); using estimation"
        )
        return None

    result: list[Any] = []
    error: list[BaseException] = []

    def _download() -> None:
        try:
            result.append(
                AutoTokenizer.from_pretrained(
                    tokenizer_name,
                    trust_remote_code=True,
                )
            )
        except BaseException as e:  # noqa: BLE001 — report any failure to the waiter
            error.append(e)

    thread = threading.Thread(
        target=_download,
        name=f"headroom-hf-tokenizer-load-{tokenizer_name}",
        daemon=True,
    )
    thread.start()
    thread.join(timeout)
    if thread.is_alive():
        logger.warning(
            f"Timed out loading tokenizer {tokenizer_name} after {timeout}s "
            f"(set {_LOAD_TIMEOUT_ENV} to adjust); using estimation"
        )
        return None
    if error:
        logger.warning(f"Failed to load tokenizer {tokenizer_name}: {error[0]}")
        return None
    return result[0] if result else None


def get_tokenizer_name(model: str) -> str:
    """Get HuggingFace tokenizer name for a model.

    Args:
        model: Model name.

    Returns:
        HuggingFace tokenizer identifier.
    """
    model_lower = model.lower()

    # Direct lookup
    if model_lower in MODEL_TO_TOKENIZER:
        return MODEL_TO_TOKENIZER[model_lower]

    # Try prefix matching
    for key, value in MODEL_TO_TOKENIZER.items():
        if model_lower.startswith(key):
            return value

    # Assume model name is the tokenizer name
    return model


class HuggingFaceTokenizer(BaseTokenizer):
    """Token counter using HuggingFace tokenizers.

    Supports any model with a HuggingFace tokenizer, including:
    - Llama family (Llama 2, Llama 3, CodeLlama)
    - Mistral family (Mistral, Mixtral)
    - Qwen family
    - DeepSeek family
    - Phi family
    - Falcon, StarCoder, MPT, Gemma, etc.

    Requires the `transformers` library:
        pip install transformers

    Some models may require authentication:
        huggingface-cli login

    Example:
        counter = HuggingFaceTokenizer("llama-3-8b")
        tokens = counter.count_text("Hello, world!")
    """

    # Overhead per message (varies by model, this is a reasonable default)
    MESSAGE_OVERHEAD = 4
    REPLY_OVERHEAD = 3

    def __init__(self, model: str):
        """Initialize HuggingFace tokenizer.

        Args:
            model: Model name (e.g., 'llama-3-8b', 'mistral-7b').
        """
        self.model = model
        self.tokenizer_name = get_tokenizer_name(model)
        self._tokenizer = None  # Lazy load

    @property
    def tokenizer(self):
        """Lazy-load the tokenizer."""
        if self._tokenizer is None:
            loaded = _load_tokenizer(self.tokenizer_name)
            if loaded is not None:
                self._tokenizer = loaded
            else:
                # Mark as unavailable
                self._tokenizer = False
        return self._tokenizer if self._tokenizer is not False else None

    def _use_fallback(self) -> bool:
        """Check if we need to use fallback estimation."""
        return self.tokenizer is None

    def count_text(self, text: str) -> int:
        """Count tokens in text.

        Falls back to estimation if tokenizer unavailable.

        Args:
            text: Text to tokenize.

        Returns:
            Number of tokens.
        """
        if not text:
            return 0
        if self._use_fallback():
            # Fall back to ~4 chars per token estimation
            return max(1, int(len(text) / 4 + 0.5))
        tokens = self.tokenizer.encode(text, add_special_tokens=False)
        return len(tokens)

    def count_messages(self, messages: list[dict[str, Any]]) -> int:
        """Count tokens in chat messages.

        Uses the model's chat template if available, otherwise
        falls back to base class implementation.

        Args:
            messages: List of chat messages.

        Returns:
            Total token count.
        """
        if self._use_fallback():
            # Use base class implementation with estimation
            return super().count_messages(messages)

        # Try to use chat template for accurate counting
        if hasattr(self.tokenizer, "apply_chat_template"):
            try:
                # Apply chat template and count
                formatted = self.tokenizer.apply_chat_template(
                    messages,
                    tokenize=True,
                    add_generation_prompt=True,
                )
                return len(formatted)
            except Exception:
                # Fall back to base implementation
                pass

        return super().count_messages(messages)

    def encode(self, text: str) -> list[int]:
        """Encode text to token IDs.

        Args:
            text: Text to encode.

        Returns:
            List of token IDs.

        Raises:
            NotImplementedError: If tokenizer not available.
        """
        if self._use_fallback():
            raise NotImplementedError(
                f"Encoding not available for {self.model} - "
                f"tokenizer {self.tokenizer_name} could not be loaded"
            )
        return self.tokenizer.encode(text, add_special_tokens=False)

    def decode(self, tokens: list[int]) -> str:
        """Decode token IDs to text.

        Args:
            tokens: List of token IDs.

        Returns:
            Decoded text.

        Raises:
            NotImplementedError: If tokenizer not available.
        """
        if self._use_fallback():
            raise NotImplementedError(
                f"Decoding not available for {self.model} - "
                f"tokenizer {self.tokenizer_name} could not be loaded"
            )
        return self.tokenizer.decode(tokens)

    @classmethod
    def is_available(cls) -> bool:
        """Check if HuggingFace tokenizers are available.

        Returns:
            True if transformers is installed.
        """
        try:
            import transformers  # noqa: F401

            return True
        except ImportError:
            return False

    @classmethod
    def list_supported_models(cls) -> list[str]:
        """List models with known tokenizer mappings.

        Returns:
            List of supported model names.
        """
        return list(MODEL_TO_TOKENIZER.keys())

    def __repr__(self) -> str:
        return f"HuggingFaceTokenizer(model={self.model!r}, tokenizer={self.tokenizer_name!r})"