headroom-ai 0.2.13__py3-none-any.whl

headroom/tokenizers/huggingface.py

@@ -0,0 +1,317 @@
+"""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
+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",
+}
+
+
+@lru_cache(maxsize=16)
+def _load_tokenizer(tokenizer_name: str):
+    """Load and cache HuggingFace 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,
+        )
+    except Exception as e:
+        logger.warning(f"Failed to load tokenizer {tokenizer_name}: {e}")
+        return 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})"

headroom/tokenizers/mistral.py

@@ -0,0 +1,245 @@
+"""Mistral tokenizer using the official mistral-common package.
+
+Mistral AI released their tokenizer publicly, making accurate
+token counting possible without API calls.
+
+Requires: pip install mistral-common
+"""
+
+from __future__ import annotations
+
+import logging
+from functools import lru_cache
+from typing import Any
+
+from .base import BaseTokenizer
+
+logger = logging.getLogger(__name__)
+
+# Check if mistral-common is available
+try:
+    from mistral_common.protocol.instruct.messages import (
+        AssistantMessage,
+        SystemMessage,
+        UserMessage,
+    )
+    from mistral_common.protocol.instruct.request import ChatCompletionRequest
+    from mistral_common.tokens.tokenizers.mistral import MistralTokenizer as _MistralTokenizer
+
+    MISTRAL_AVAILABLE = True
+except ImportError:
+    MISTRAL_AVAILABLE = False
+    _MistralTokenizer = None
+
+
+def is_mistral_available() -> bool:
+    """Check if mistral-common is installed."""
+    return MISTRAL_AVAILABLE
+
+
+# Model to tokenizer version mapping
+MODEL_TO_VERSION = {
+    # Mistral models use v3 tokenizer (tekken)
+    "mistral-large": "v3",
+    "mistral-large-latest": "v3",
+    "mistral-small": "v3",
+    "mistral-small-latest": "v3",
+    "ministral-8b": "v3",
+    "ministral-3b": "v3",
+    "mistral-nemo": "v3",
+    "pixtral-12b": "v3",
+    "codestral": "v3",
+    "codestral-latest": "v3",
+    # Mixtral uses v1
+    "mixtral-8x7b": "v1",
+    "mixtral-8x22b": "v1",
+    "open-mixtral-8x7b": "v1",
+    "open-mixtral-8x22b": "v1",
+    # Mistral 7B uses v1
+    "mistral-7b": "v1",
+    "open-mistral-7b": "v1",
+    "mistral-7b-instruct": "v1",
+}
+
+
+@lru_cache(maxsize=4)
+def _get_tokenizer(version: str):
+    """Get and cache Mistral tokenizer by version."""
+    if not MISTRAL_AVAILABLE:
+        raise RuntimeError(
+            "mistral-common is required for MistralTokenizer. "
+            "Install with: pip install mistral-common"
+        )
+
+    if version == "v3":
+        return _MistralTokenizer.v3(is_tekken=True)
+    elif version == "v2":
+        return _MistralTokenizer.v2()
+    else:  # v1
+        return _MistralTokenizer.v1()
+
+
+def get_tokenizer_version(model: str) -> str:
+    """Get tokenizer version for a model."""
+    model_lower = model.lower()
+
+    # Direct lookup
+    if model_lower in MODEL_TO_VERSION:
+        return MODEL_TO_VERSION[model_lower]
+
+    # Prefix matching
+    for prefix, version in [
+        ("mistral-large", "v3"),
+        ("mistral-small", "v3"),
+        ("ministral", "v3"),
+        ("codestral", "v3"),
+        ("pixtral", "v3"),
+        ("mistral-nemo", "v3"),
+        ("mixtral", "v1"),
+        ("mistral-7b", "v1"),
+        ("open-mistral", "v1"),
+    ]:
+        if model_lower.startswith(prefix):
+            return version
+
+    # Default to v3 for newer models
+    return "v3"
+
+
+class MistralTokenizer(BaseTokenizer):
+    """Token counter using Mistral's official tokenizer.
+
+    Uses mistral-common package for accurate token counting.
+
+    Requires: pip install mistral-common
+
+    Example:
+        counter = MistralTokenizer("mistral-large")
+        tokens = counter.count_text("Hello, world!")
+    """
+
+    MESSAGE_OVERHEAD = 4
+    REPLY_OVERHEAD = 3
+
+    def __init__(self, model: str = "mistral-large"):
+        """Initialize Mistral tokenizer.
+
+        Args:
+            model: Mistral model name.
+        """
+        if not MISTRAL_AVAILABLE:
+            raise RuntimeError(
+                "mistral-common is required for MistralTokenizer. "
+                "Install with: pip install mistral-common"
+            )
+
+        self.model = model
+        self.version = get_tokenizer_version(model)
+        self._tokenizer = None  # Lazy load
+
+    @property
+    def tokenizer(self):
+        """Lazy-load the tokenizer (MistralTokenizer object)."""
+        if self._tokenizer is None:
+            self._tokenizer = _get_tokenizer(self.version)
+        return self._tokenizer
+
+    @property
+    def _text_tokenizer(self):
+        """Get the underlying text tokenizer for encode/decode."""
+        return self.tokenizer.instruct_tokenizer.tokenizer
+
+    def count_text(self, text: str) -> int:
+        """Count tokens in text.
+
+        Args:
+            text: Text to tokenize.
+
+        Returns:
+            Number of tokens.
+        """
+        if not text:
+            return 0
+        tokens = self._text_tokenizer.encode(text, bos=False, eos=False)
+        return len(tokens)
+
+    def count_messages(self, messages: list[dict[str, Any]]) -> int:
+        """Count tokens in chat messages.
+
+        Uses Mistral's chat template for accurate counting.
+
+        Args:
+            messages: List of chat messages.
+
+        Returns:
+            Total token count.
+        """
+        if not messages:
+            return 0
+
+        try:
+            # Convert to Mistral message format
+            mistral_messages = []
+            for msg in messages:
+                role = msg.get("role", "user")
+                content = msg.get("content", "")
+
+                if isinstance(content, list):
+                    # Multi-part content - extract text
+                    text_parts = []
+                    for part in content:
+                        if isinstance(part, dict) and part.get("type") == "text":
+                            text_parts.append(part.get("text", ""))
+                        elif isinstance(part, str):
+                            text_parts.append(part)
+                    content = "\n".join(text_parts)
+
+                if role == "user":
+                    mistral_messages.append(UserMessage(content=content))
+                elif role == "assistant":
+                    mistral_messages.append(AssistantMessage(content=content))
+                elif role == "system":
+                    mistral_messages.append(SystemMessage(content=content))
+                else:
+                    # Tool messages etc - treat as user
+                    mistral_messages.append(UserMessage(content=content))
+
+            # Encode with chat template
+            request = ChatCompletionRequest(messages=mistral_messages)
+            tokenized = self.tokenizer.encode_chat_completion(request)
+            return len(tokenized.tokens)
+
+        except Exception as e:
+            logger.debug(f"Mistral chat encoding failed: {e}, falling back to text counting")
+            # Fallback to base implementation
+            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.
+        """
+        return self._text_tokenizer.encode(text, bos=False, eos=False)
+
+    def decode(self, tokens: list[int]) -> str:
+        """Decode token IDs to text.
+
+        Args:
+            tokens: List of token IDs.
+
+        Returns:
+            Decoded text.
+        """
+        return self._text_tokenizer.decode(tokens)
+
+    @classmethod
+    def is_available(cls) -> bool:
+        """Check if Mistral tokenizer is available."""
+        return MISTRAL_AVAILABLE
+
+    def __repr__(self) -> str:
+        return f"MistralTokenizer(model={self.model!r}, version={self.version!r})"