tokenizer.py

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