modeling_emg.py

"""
HuggingFace Integration for EMG Model and MorPiece Tokenizer
This file makes your custom model and tokenizer compatible with HuggingFace and lm_eval
"""

import json
import os
from typing import List, Optional, Union, Dict, Any
import torch
import torch.nn as nn
from transformers import (
    PreTrainedModel, 
    PretrainedConfig, 
    PreTrainedTokenizer,
    AutoConfig,
    AutoModel,
    AutoTokenizer,
    AutoModelForCausalLM,
    GenerationMixin,  # Add this import
)
from transformers.modeling_outputs import CausalLMOutputWithPast

# Import your existing classes
from model_eMG_simplified import EMGLanguageModel, EMGConfig, OptimizedEMG, OptimizedEMGCell
from tokenizer_MorPiece import MorPiece


class MorPieceTokenizer(PreTrainedTokenizer):
    """
    HuggingFace compatible wrapper for MorPiece tokenizer
    """
    
    def __init__(self, 
                 vocab_file=None,
                 model_file=None,
                 unk_token="<unk>",
                 pad_token="<pad>", 
                 bos_token="<s>",
                 eos_token="</s>",
                 **kwargs):
        
        # Initialize the MorPiece tokenizer
        self.morpiece = MorPiece()
        
        # Load from file if provided
        if vocab_file or model_file:
            model_path = vocab_file or model_file
            if os.path.isdir(model_path):
                self.morpiece.from_pretrained(model_path)
            else:
                # Load from JSON file
                with open(model_path, 'r') as f:
                    data = json.load(f)
                self.morpiece.roots = data.get('roots', data)
                if 'vocab' in data:
                    self.morpiece.vocab_to_id = data['vocab']
                else:
                    self.morpiece.build_vocab_lookup()
        
        # Get vocabulary
        self.vocab = self.morpiece.get_vocab()
        
        # Set special tokens
        super().__init__(
            unk_token=unk_token,
            pad_token=pad_token,
            bos_token=bos_token,
            eos_token=eos_token,
            **kwargs
        )
        
    @property
    def vocab_size(self):
        return len(self.vocab)
    
    def get_vocab(self):
        return self.vocab.copy()
    
    def _tokenize(self, text: str) -> List[str]:
        """Tokenize text into tokens"""
        # For HuggingFace compatibility, we need to return string tokens
        token_ids = self.morpiece.encode(text)
        tokens = self.morpiece.decode(token_ids)
        return tokens
    
    def _convert_token_to_id(self, token: str) -> int:
        """Convert token to ID"""
        return self.vocab.get(token, self.vocab.get(self.unk_token, 0))
    
    def _convert_id_to_token(self, index: int) -> str:
        """Convert ID to token"""
        for token, idx in self.vocab.items():
            if idx == index:
                return token
        return self.unk_token
    
    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        """Convert tokens back to string"""
        # Handle special tokens
        text = "".join(tokens)
        # Clean up special tokens for display
        for special_token in [self.pad_token, self.bos_token, self.eos_token]:
            if special_token:
                text = text.replace(special_token, "")
        return text.strip()
    
    def encode(self, text: str, add_special_tokens: bool = True, **kwargs) -> List[int]:
        """Encode text to token IDs"""
        if add_special_tokens and self.bos_token:
            text = f"{self.bos_token} {text}"
        if add_special_tokens and self.eos_token:
            text = f"{text} {self.eos_token}"
            
        return self.morpiece.encode(text)
    
    def decode(self, token_ids: List[int], skip_special_tokens: bool = True, **kwargs) -> str:
        """Decode token IDs to text"""
        tokens = []
        for token_id in token_ids:
            token = self._convert_id_to_token(token_id)
            if skip_special_tokens and token in [self.pad_token, self.bos_token, self.eos_token, self.unk_token]:
                continue
            tokens.append(token)
        return self.convert_tokens_to_string(tokens)
    
    def save_pretrained(self, save_directory: str, **kwargs):
        """Save tokenizer"""
        os.makedirs(save_directory, exist_ok=True)
        
        # Save MorPiece data
        tokenizer_file = os.path.join(save_directory, "tokenizer.json")
        self.morpiece.save(tokenizer_file)
        
        # Save tokenizer config
        config = {
            "tokenizer_class": "MorPieceTokenizer",
            "unk_token": self.unk_token,
            "pad_token": self.pad_token,
            "bos_token": self.bos_token,
            "eos_token": self.eos_token,
        }
        
        config_file = os.path.join(save_directory, "tokenizer_config.json")
        with open(config_file, 'w') as f:
            json.dump(config, f, indent=2)
    
    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
        """Load tokenizer from pretrained"""
        return cls(vocab_file=pretrained_model_name_or_path, **kwargs)


class EMGForCausalLM(EMGLanguageModel, GenerationMixin):
    """
    Enhanced EMG model with better HuggingFace compatibility for lm_eval
    Inherits from GenerationMixin to fix the warning
    """
    
    def __init__(self, config):
        # Initialize EMGLanguageModel first
        EMGLanguageModel.__init__(self, config)
        # Then initialize GenerationMixin
        GenerationMixin.__init__(self)
        self.config = config
    
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        past_key_values: Optional[tuple] = None,
        use_cache: Optional[bool] = None,
        **kwargs
    ) -> CausalLMOutputWithPast:
        """
        Forward pass with HuggingFace compatible output format
        """
        # Get embeddings
        embedded = self.embedding(input_ids)
        
        # Pass through EMG layers
        output, hidden = self.emg(embedded, past_key_values)
        
        # Get logits
        logits = self.output_projection(output)
        
        loss = None
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            
            # Flatten the tokens
            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), 
                          shift_labels.view(-1))
        
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=hidden if use_cache else None,
            hidden_states=output,
        )
    
    def prepare_inputs_for_generation(
        self, 
        input_ids: torch.Tensor, 
        past_key_values=None, 
        attention_mask=None, 
        **kwargs
    ):
        """Prepare inputs for generation"""
        return {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "attention_mask": attention_mask,
        }
    
    def _reorder_cache(self, past_key_values, beam_idx):
        """Reorder cache for beam search"""
        if past_key_values is None:
            return None
        
        reordered_cache = []
        for layer_cache in past_key_values:
            if isinstance(layer_cache, tuple):
                reordered_cache.append(tuple(
                    cache.index_select(0, beam_idx) for cache in layer_cache
                ))
            else:
                reordered_cache.append(layer_cache.index_select(0, beam_idx))
        return tuple(reordered_cache)


# Register the custom classes with transformers
def register_emg_model():
    """Register EMG model and tokenizer with transformers"""
    
    # Register config
    AutoConfig.register("emg", EMGConfig)
    
    # Register model
    AutoModel.register(EMGConfig, EMGLanguageModel)
    AutoModelForCausalLM.register(EMGConfig, EMGForCausalLM)
    
    # Register tokenizer
    AutoTokenizer.register(EMGConfig, MorPieceTokenizer)
    
    print("EMG model and MorPiece tokenizer registered with transformers!")


def load_emg_model_and_tokenizer(model_path: str):
    """
    Load EMG model and MorPiece tokenizer from saved directory
    
    Args:
        model_path: Path to the saved model directory
        
    Returns:
        tuple: (model, tokenizer)
    """
    # Register classes first
    register_emg_model()
    
    # Load model
    config = EMGConfig.from_pretrained(model_path)
    model = EMGForCausalLM.from_pretrained(model_path, config=config)
    
    # Load tokenizer
    tokenizer = MorPieceTokenizer.from_pretrained(model_path)
    
    # Set pad token id in model config if not set
    if not hasattr(config, 'pad_token_id') or config.pad_token_id is None:
        config.pad_token_id = tokenizer.pad_token_id
        model.config.pad_token_id = tokenizer.pad_token_id
    
    return model, tokenizer


def test_model_and_tokenizer(model_path: str):
    """Test the loaded model and tokenizer"""
    model, tokenizer = load_emg_model_and_tokenizer(model_path)
    
    # Test encoding/decoding
    test_text = "Hello world, this is a test."
    print(f"Original text: {test_text}")
    
    # Encode
    encoded = tokenizer.encode(test_text)
    print(f"Encoded: {encoded}")
    
    # Decode
    decoded = tokenizer.decode(encoded, skip_special_tokens=True)
    print(f"Decoded: {decoded}")
    
    # Test model forward pass
    input_ids = torch.tensor([encoded])
    with torch.no_grad():
        outputs = model(input_ids)
        print(f"Model output shape: {outputs.logits.shape}")
        print(f"Model output type: {type(outputs)}")
    
    print("Model and tokenizer are working correctly!")
    return model, tokenizer


if __name__ == "__main__":
    # Example usage
    model_path = "path/to/your/saved/model"  # Replace with your model path
    
    # Register the classes
    register_emg_model()
    
    # Test loading
    try:
        model, tokenizer = test_model_and_tokenizer(model_path)
        print("✅ Model and tokenizer loaded successfully!")
    except Exception as e:
        print(f"❌ Error loading model: {e}")