Upload tokenizer implementation

tokenizer.py

@@ -0,0 +1,287 @@
+import numpy as np
+import logging
+from typing import Dict, List, Tuple, Optional
+from dataclasses import dataclass
+from transformers import PreTrainedTokenizer
+import os
+import json
+from huggingface_hub import Repository
+from huggingface_hub import HfApi
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+WAVELET_TOKENIZER_CONFIG = {
+    "model_type": "wavelet",
+    "tokenizer_class": "WaveletTokenizer",
+    "auto_map": {
+        "AutoTokenizer": ["tokenizer.WaveletTokenizer", None]
+    }
+}
+
+@dataclass
+class WaveletTokenizerConfig:
+    vocab_size: int = 256
+    padding_idx: int = 0
+    eeg_channels: int = 74     # Source modality (EEG)
+    mu: float = 255.0         # Static μ value for μ-law compression
+    verbose: bool = True       # Control logging
+
+class WaveletTokenizer(PreTrainedTokenizer):
+    model_input_names = ["input_ids", "attention_mask", "position_ids"]
+    
+    def __init__(
+        self,
+        vocab_size: int = 256,
+        mu: float = 255.0,
+        verbose: bool = True,
+        **kwargs
+    ):
+        self.auto_map = {
+            "AutoTokenizer": ["tokenizer.WaveletTokenizer", None]
+        }
+        
+        # Set vocab size first
+        self._vocab_size = vocab_size
+        self.mu = mu
+        self.verbose = verbose
+        
+        # Store normalization state
+        self.channel_mins = None
+        self.channel_maxs = None
+        
+        # Initialize parent class after setting vocab_size
+        super().__init__(**kwargs)
+        
+        if self.verbose:
+            logger.info(f"Initialized WaveletTokenizer with μ={self.mu:.2f}")
+    
+    @property
+    def vocab_size(self) -> int:
+        """Returns the size of vocabulary (number of possible quantization levels)."""
+        return self._vocab_size
+    
+    @vocab_size.setter
+    def vocab_size(self, size: int):
+        self._vocab_size = size
+    
+    def save_pretrained(
+        self, 
+        save_directory: str,
+        legacy_format: bool = True,
+        filename_prefix: Optional[str] = None,
+        push_to_hub: bool = False,
+        **kwargs
+    ) -> Tuple[str, ...]:
+        """Save tokenizer configuration to a directory."""
+        if not os.path.exists(save_directory):
+            os.makedirs(save_directory)
+            
+        # Save tokenizer config
+        config = {
+            **WAVELET_TOKENIZER_CONFIG,
+            "vocab_size": self.vocab_size,
+            "mu": self.mu,
+            "verbose": self.verbose
+        }
+        
+        config_file = os.path.join(
+            save_directory, 
+            (filename_prefix + "-" if filename_prefix else "") + "tokenizer_config.json"
+        )
+        
+        with open(config_file, "w") as f:
+            json.dump(config, f, indent=2)
+            
+        # Save vocabulary
+        vocab_files = self.save_vocabulary(save_directory, filename_prefix=filename_prefix)
+        
+        if push_to_hub:
+            # Upload files using HTTP
+            api = HfApi()
+            api.upload_file(
+                path_or_fileobj=config_file,
+                path_in_repo="tokenizer_config.json",
+                repo_id=save_directory,
+                commit_message=kwargs.get("commit_message", "Upload tokenizer config")
+            )
+            
+            # Upload vocabulary file
+            vocab_file = vocab_files[0]
+            api.upload_file(
+                path_or_fileobj=vocab_file,
+                path_in_repo=os.path.basename(vocab_file),
+                repo_id=save_directory,
+                commit_message=kwargs.get("commit_message", "Upload tokenizer vocabulary")
+            )
+            
+        return vocab_files + (config_file,)
+    
+    @classmethod
+    def from_pretrained(
+        cls, 
+        pretrained_model_name_or_path: str, 
+        **kwargs
+    ) -> "WaveletTokenizer":
+        """Load tokenizer from HuggingFace Hub."""
+        # Load config first
+        config_file = os.path.join(pretrained_model_name_or_path, "tokenizer_config.json")
+        if os.path.exists(config_file):
+            with open(config_file, "r") as f:
+                config = json.load(f)
+            # Update with any passed kwargs
+            config.update(kwargs)
+        else:
+            config = kwargs
+            
+        return cls(**config)
+    
+    def get_vocab(self) -> Dict[str, int]:
+        """Returns vocab as a dict mapping token strings to ids."""
+        # Create a minimal vocabulary with quantization levels
+        return {str(i): i for i in range(self.vocab_size)}
+    
+    def _convert_token_to_id(self, token: str) -> int:
+        """Converts a token string to its ID."""
+        try:
+            return int(token)
+        except ValueError:
+            return 0  # Return 0 for unknown tokens
+    
+    def _convert_id_to_token(self, index: int) -> str:
+        """Converts an ID back to its token string."""
+        return str(index)
+    
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """Converts a sequence of tokens to a single string."""
+        return " ".join(tokens)
+    
+    def _tokenize(self, text: str) -> List[str]:
+        """Basic tokenization for compatibility."""
+        if isinstance(text, str):
+            return [text]
+        return [str(t) for t in text]
+    
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str, ...]:
+        """Save the vocabulary to a directory."""
+        vocab_file = os.path.join(
+            save_directory, 
+            (filename_prefix + "-" if filename_prefix else "") + "vocab.json"
+        )
+        
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            json.dump(self.get_vocab(), f, ensure_ascii=False)
+            
+        return (vocab_file,)
+    
+    def __call__(
+        self,
+        eeg_data: np.ndarray,
+        **kwargs
+    ) -> Dict[str, np.ndarray]:
+        """
+        Main entry point for tokenization. Handles numpy array input.
+        
+        Args:
+            eeg_data: Raw EEG array of shape (n_channels, time_points)
+            
+        Returns:
+            Dictionary containing:
+                - input_ids: Tokenized signal values
+                - attention_mask: Binary mask (all ones since we don't pad)
+                - position_ids: Sequential position indices
+        """
+        # Process through tokenization pipeline
+        input_ids = self.encode(eeg_data)
+        
+        # Create attention mask (all ones since we're not padding)
+        attention_mask = np.ones_like(input_ids)
+        
+        # Create position IDs
+        n_channels, time_points = eeg_data.shape
+        position_ids = np.tile(np.arange(time_points), (n_channels, 1))
+        
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids
+        }
+    
+    def encode(self, eeg_data: np.ndarray) -> np.ndarray:
+        """Convert EEG data to token IDs."""
+        # 1. Normalize to [0, 1]
+        normalized = self.normalize(eeg_data)
+        
+        # 2. Convert to [-1, 1] for μ-law compression
+        centered = 2 * normalized - 1
+        
+        # 3. Apply μ-law compression
+        compressed = self.mu_law_encode(centered)
+        
+        # 4. Quantize to tokens
+        input_values = (compressed + 1) / 2  # to [0, 1]
+        token_ids = (input_values * (self.vocab_size - 1)).astype(np.int64)
+        
+        return token_ids
+    
+    def normalize(self, x: np.ndarray) -> np.ndarray:
+        """
+        Apply static normalization per channel and store min/max values.
+        Input shape: (n_channels, time_points)
+        """
+        # Compute min/max per channel and expand dimensions to match input
+        self.channel_mins = x.min(axis=1)[:, np.newaxis]  # Shape: (n_channels, 1)
+        self.channel_maxs = x.max(axis=1)[:, np.newaxis]  # Shape: (n_channels, 1)
+        
+        normalized = (x - self.channel_mins) / (self.channel_maxs - self.channel_mins + 1e-8)
+        
+        if self.verbose:
+            logger.info(f"Min-max normalization: input range [{x.min():.3f}, {x.max():.3f}] → [{normalized.min():.3f}, {normalized.max():.3f}]")
+        return normalized
+    
+    def mu_law_encode(self, x: np.ndarray) -> np.ndarray:
+        """
+        Apply μ-law compression.
+        Expects input in [-1, 1] range.
+        """
+        assert np.all(x >= -1.0) and np.all(x <= 1.0), f"Input must be in [-1, 1] range, got min={x.min():.3f}, max={x.max():.3f}"
+        compressed = np.sign(x) * np.log1p(self.mu * np.abs(x)) / np.log1p(self.mu)
+        
+        if self.verbose:
+            logger.info(f"μ-law compression (μ={self.mu:.2f}): variance before={np.var(x):.3f}, after={np.var(compressed):.3f}")
+        return compressed
+    
+    def mu_law_decode(self, x: np.ndarray) -> np.ndarray:
+        """
+        Inverse μ-law compression.
+        Expects input in [-1, 1] range.
+        """
+        assert np.all(x >= -1.0) and np.all(x <= 1.0), f"Input must be in [-1, 1] range, got min={x.min():.3f}, max={x.max():.3f}"
+        return np.sign(x) * (1/self.mu) * (np.power(1 + self.mu, np.abs(x)) - 1.0)
+    
+    def decode(self, token_ids: np.ndarray) -> np.ndarray:
+        """
+        Decode token IDs back to EEG signal.
+        
+        Args:
+            token_ids: Array of token IDs of shape (n_channels, time_points)
+            
+        Returns:
+            Array of shape (n_channels, time_points)
+        """
+        # Convert to continuous values in [-1, 1]
+        values = token_ids.astype(np.float32) / (self.vocab_size - 1)  # [0, 1]
+        values = 2 * values - 1  # [-1, 1]
+        
+        # Apply inverse μ-law compression
+        values = self.mu_law_decode(values)
+        
+        # Convert back to [0, 1]
+        values = (values + 1) / 2
+        
+        # Denormalize to original scale
+        if self.channel_mins is not None and self.channel_maxs is not None:
+            values = values * (self.channel_maxs - self.channel_mins) + self.channel_mins
+        
+        return values