Integrate mistral-common for correct Devstral tokenization

Root cause: Devstral's Tekken tokenizer is incompatible with HuggingFace's
standard tokenization. When [INST]/[/INST] are formatted as text and
tokenized with HF, the model receives corrupted tokens -> garbage output.

Solution:
- Add mistral-common>=1.5.0 dependency
- Create MistralTokenizerWrapper using MistralTokenizer.from_hf_hub()
- Use encode_chat_completion() for correct Tekken token encoding
- Relax numpy/pydantic version constraints for compatibility

The mistral-common library produces correct token sequences by encoding
chat messages directly to token IDs, bypassing text-based formatting.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>

backend/mistral_tokenizer.py

@@ -0,0 +1,136 @@
+"""
+Mistral Tokenizer Wrapper
+Provides correct tokenization for Devstral using mistral-common library.
+
+The Tekken tokenizer used by Devstral is incompatible with HuggingFace's
+standard tokenization approach. This wrapper uses mistral-common to
+produce correct token sequences for the model.
+"""
+
+import logging
+from typing import List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class MistralTokenizerWrapper:
+    """
+    Wrapper around mistral-common's MistralTokenizer for Devstral.
+
+    Uses encode_chat_completion() to produce correct token IDs
+    that the model actually expects, rather than HF's text-based approach
+    which produces corrupted tokens for Tekken-based models.
+    """
+
+    def __init__(self, model_name: str):
+        """
+        Initialize the Mistral tokenizer from HuggingFace hub.
+
+        Args:
+            model_name: HuggingFace model path (e.g., "mistralai/Devstral-Small-2507")
+        """
+        try:
+            from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+            self.tokenizer = MistralTokenizer.from_hf_hub(model_name)
+            self._available = True
+            logger.info(f"Loaded MistralTokenizer for {model_name}")
+        except ImportError as e:
+            logger.warning(f"mistral-common not available: {e}")
+            self._available = False
+            self.tokenizer = None
+        except Exception as e:
+            logger.error(f"Failed to load MistralTokenizer: {e}")
+            self._available = False
+            self.tokenizer = None
+
+    @property
+    def is_available(self) -> bool:
+        """Check if the tokenizer was loaded successfully."""
+        return self._available
+
+    def encode_chat(
+        self,
+        system_prompt: str,
+        user_prompt: str
+    ) -> List[int]:
+        """
+        Encode chat messages to token IDs using mistral-common.
+
+        This produces the correct token sequence for Devstral, including
+        proper handling of control tokens like [INST] and [/INST].
+
+        Args:
+            system_prompt: System message content
+            user_prompt: User message content (e.g., "def quicksort(arr):")
+
+        Returns:
+            List of token IDs ready for model input
+        """
+        if not self._available:
+            raise RuntimeError("MistralTokenizer not available")
+
+        from mistral_common.protocol.instruct.messages import (
+            SystemMessage, UserMessage
+        )
+        from mistral_common.protocol.instruct.request import ChatCompletionRequest
+
+        # Build messages list
+        messages = []
+        if system_prompt:
+            messages.append(SystemMessage(content=system_prompt))
+        messages.append(UserMessage(content=user_prompt))
+
+        # Encode using mistral-common's chat completion encoding
+        request = ChatCompletionRequest(messages=messages)
+        tokenized = self.tokenizer.encode_chat_completion(request)
+
+        logger.info(f"Encoded chat: {len(tokenized.tokens)} tokens")
+        return tokenized.tokens
+
+    def decode(self, token_ids: List[int]) -> str:
+        """
+        Decode token IDs back to text.
+
+        Args:
+            token_ids: List of token IDs to decode
+
+        Returns:
+            Decoded text string
+        """
+        if not self._available:
+            raise RuntimeError("MistralTokenizer not available")
+
+        return self.tokenizer.decode(token_ids)
+
+    def decode_token(self, token_id: int) -> str:
+        """
+        Decode a single token ID to text.
+
+        Args:
+            token_id: Single token ID to decode
+
+        Returns:
+            Decoded text for this token
+        """
+        if not self._available:
+            raise RuntimeError("MistralTokenizer not available")
+
+        return self.tokenizer.decode([token_id])
+
+
+def create_mistral_tokenizer(model_name: str) -> Optional[MistralTokenizerWrapper]:
+    """
+    Factory function to create a MistralTokenizerWrapper.
+
+    Returns None if mistral-common is not available or loading fails.
+
+    Args:
+        model_name: HuggingFace model path
+
+    Returns:
+        MistralTokenizerWrapper instance or None
+    """
+    wrapper = MistralTokenizerWrapper(model_name)
+    if wrapper.is_available:
+        return wrapper
+    return None

backend/model_service.py

@@ -229,6 +229,16 @@ class ModelManager:
             self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
             self.tokenizer.pad_token = self.tokenizer.eos_token
 
+            # For Devstral, also load MistralTokenizer for correct encoding
+            self.mistral_tokenizer = None
+            if self.model_id == "devstral-small":
+                from .mistral_tokenizer import create_mistral_tokenizer
+                self.mistral_tokenizer = create_mistral_tokenizer(self.model_name)
+                if self.mistral_tokenizer:
+                    logger.info("Loaded MistralTokenizer for Devstral (correct Tekken encoding)")
+                else:
+                    logger.warning("MistralTokenizer not available - Devstral may produce garbage output")
+
             # Create model adapter for multi-model support
             from .model_adapter import create_adapter
             try:
@@ -1514,11 +1524,22 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
             temperature = model_config["recommended_temperature"]
             logger.info(f"Using model recommended temperature={temperature}")
 
-        # Tokenize and prepare
-        inputs = manager.tokenizer(formatted_prompt, return_tensors="pt").to(manager.device)
-        prompt_length = inputs["input_ids"].shape[1]
-        prompt_token_ids = inputs["input_ids"][0].tolist()
-        prompt_tokens = [manager.tokenizer.decode([tid], skip_special_tokens=False) for tid in prompt_token_ids]
+        # Tokenize and prepare - use MistralTokenizer for Devstral
+        if manager.model_id == "devstral-small" and manager.mistral_tokenizer is not None:
+            # Use MistralTokenizer for correct Tekken encoding
+            system_prompt = system_prompt_override or (model_config.get("system_prompt") if model_config else "")
+            prompt_token_ids = manager.mistral_tokenizer.encode_chat(system_prompt, prompt)
+            inputs = {"input_ids": torch.tensor([prompt_token_ids]).to(manager.device)}
+            prompt_length = len(prompt_token_ids)
+            # Decode tokens using MistralTokenizer for accuracy
+            prompt_tokens = [manager.mistral_tokenizer.decode_token(tid) for tid in prompt_token_ids]
+            logger.info(f"Used MistralTokenizer for Devstral: {prompt_length} tokens")
+        else:
+            # Standard HF tokenization for other models
+            inputs = manager.tokenizer(formatted_prompt, return_tensors="pt").to(manager.device)
+            prompt_length = inputs["input_ids"].shape[1]
+            prompt_token_ids = inputs["input_ids"][0].tolist()
+            prompt_tokens = [manager.tokenizer.decode([tid], skip_special_tokens=False) for tid in prompt_token_ids]
 
         # Storage for generation
         generated_token_ids = []

requirements.txt

@@ -3,16 +3,17 @@ fastapi==0.104.1
 uvicorn[standard]==0.24.0
 websockets==12.0
 python-multipart==0.0.6
-pydantic==2.5.0
+pydantic>=2.0.0  # Relaxed for mistral-common compatibility
 
 # Machine Learning
 # torch 2.3+ required for transformers 4.44+ (pytree API compatibility)
 torch>=2.3.0
 transformers>=4.44.0
 accelerate>=0.30.0
+mistral-common>=1.5.0  # Required for Devstral Tekken tokenizer
 
 # Utilities
-numpy==1.24.3
+numpy>=1.24.0,<2.0  # Relaxed for mistral-common compatibility
 aiofiles==23.2.1
 python-dotenv==1.0.0
 zarr==2.14.2