Create modeling.py

modeling.py

@@ -0,0 +1,335 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import Seq2SeqLMOutput
+from typing import Optional, Tuple, Union
+import math
+
+
+class PositionalEncoding(nn.Module):
+    """Positional encoding for transformer"""
+    def __init__(self, d_model, max_length=5000):
+        super().__init__()
+        pe = torch.zeros(max_length, d_model)
+        position = torch.arange(0, max_length, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1)]
+
+
+class TranslationTransformerConfig(PretrainedConfig):
+    """Configuration class for TranslationTransformer"""
+    model_type = "translation_transformer"
+    
+    def __init__(
+        self,
+        vocab_size=32000,
+        d_model=512,
+        nhead=8,
+        num_encoder_layers=6,
+        num_decoder_layers=6,
+        dim_feedforward=2048,
+        dropout=0.1,
+        pad_token_id=0,
+        bos_token_id=2,
+        eos_token_id=3,
+        max_length=512,
+        **kwargs
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            **kwargs
+        )
+        
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.nhead = nhead
+        self.num_encoder_layers = num_encoder_layers
+        self.num_decoder_layers = num_decoder_layers
+        self.dim_feedforward = dim_feedforward
+        self.dropout = dropout
+        self.max_length = max_length
+        
+        # Required for HuggingFace compatibility
+        self.is_encoder_decoder = True
+        self.decoder_start_token_id = bos_token_id
+
+
+class TranslationTransformerModel(PreTrainedModel):
+    """
+    Encoder-Decoder Transformer for Translation
+    Compatible with HuggingFace Hub
+    """
+    config_class = TranslationTransformerConfig
+    base_model_prefix = "translation_transformer"
+    supports_gradient_checkpointing = True
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        
+        # Embeddings
+        self.embedding = nn.Embedding(
+            config.vocab_size, 
+            config.d_model, 
+            padding_idx=config.pad_token_id
+        )
+        self.pos_encoder = PositionalEncoding(config.d_model, config.max_length)
+        self.pos_decoder = PositionalEncoding(config.d_model, config.max_length)
+        
+        # Transformer
+        self.transformer = nn.Transformer(
+            d_model=config.d_model,
+            nhead=config.nhead,
+            num_encoder_layers=config.num_encoder_layers,
+            num_decoder_layers=config.num_decoder_layers,
+            dim_feedforward=config.dim_feedforward,
+            dropout=config.dropout,
+            batch_first=True
+        )
+        
+        # Output layer
+        self.fc_out = nn.Linear(config.d_model, config.vocab_size)
+        
+        # Initialize weights
+        self.post_init()
+    
+    def _init_weights(self, module):
+        """Initialize weights"""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+    
+    def get_encoder(self):
+        """Return encoder for compatibility"""
+        return self.transformer.encoder
+    
+    def get_decoder(self):
+        """Return decoder for compatibility"""
+        return self.transformer.decoder
+    
+    def generate_square_subsequent_mask(self, sz, device):
+        """Generate causal mask for decoder"""
+        mask = torch.triu(torch.ones(sz, sz, device=device), diagonal=1)
+        mask = mask.masked_fill(mask == 1, float('-inf'))
+        return mask
+    
+    def create_padding_mask(self, seq, pad_token_id):
+        """Create padding mask"""
+        return (seq == pad_token_id)
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        decoder_attention_mask: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs
+    ) -> Union[Tuple, Seq2SeqLMOutput]:
+        """
+        Forward pass
+        
+        Args:
+            input_ids: Source sequence tokens [batch_size, src_seq_len]
+            attention_mask: Source attention mask [batch_size, src_seq_len]
+            decoder_input_ids: Target sequence tokens [batch_size, tgt_seq_len]
+            decoder_attention_mask: Target attention mask [batch_size, tgt_seq_len]
+            labels: Labels for loss calculation [batch_size, tgt_seq_len]
+            output_attentions: Whether to output attentions
+            output_hidden_states: Whether to output hidden states
+            return_dict: Whether to return ModelOutput
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        device = input_ids.device
+        
+        # If labels provided but no decoder_input_ids, shift labels to create decoder_input_ids
+        if labels is not None and decoder_input_ids is None:
+            # Replace -100 with pad_token_id for embedding
+            labels_shifted = labels.clone()
+            labels_shifted[labels_shifted == -100] = self.config.pad_token_id
+            
+            # Shift right: [BOS, token1, token2, ...] from [token1, token2, ..., EOS]
+            decoder_input_ids = torch.cat([
+                torch.full((labels.shape[0], 1), self.config.bos_token_id, dtype=torch.long, device=device),
+                labels_shifted[:, :-1]
+            ], dim=1)
+        
+        # Embeddings with scaling
+        src_emb = self.embedding(input_ids) * math.sqrt(self.config.d_model)
+        src_emb = self.pos_encoder(src_emb)
+        
+        tgt_emb = self.embedding(decoder_input_ids) * math.sqrt(self.config.d_model)
+        tgt_emb = self.pos_decoder(tgt_emb)
+        
+        # Create masks
+        tgt_seq_len = decoder_input_ids.size(1)
+        tgt_mask = self.generate_square_subsequent_mask(tgt_seq_len, device)
+        
+        src_key_padding_mask = self.create_padding_mask(input_ids, self.config.pad_token_id)
+        tgt_key_padding_mask = self.create_padding_mask(decoder_input_ids, self.config.pad_token_id)
+        
+        # Transformer forward pass
+        output = self.transformer(
+            src_emb,
+            tgt_emb,
+            tgt_mask=tgt_mask,
+            src_key_padding_mask=src_key_padding_mask,
+            tgt_key_padding_mask=tgt_key_padding_mask,
+            memory_key_padding_mask=src_key_padding_mask
+        )
+        
+        # Output projection
+        logits = self.fc_out(output)
+        
+        # Calculate loss if labels provided
+        loss = None
+        if labels is not None:
+            # Use -100 as ignore_index (standard for HuggingFace)
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+        
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+        
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=logits,
+        )
+    
+    def prepare_inputs_for_generation(
+        self,
+        decoder_input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        use_cache=None,
+        encoder_outputs=None,
+        **kwargs
+    ):
+        """Prepare inputs for generation (required for HuggingFace generate)"""
+        return {
+            "input_ids": kwargs.get("input_ids"),
+            "decoder_input_ids": decoder_input_ids,
+            "attention_mask": attention_mask,
+        }
+    
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        """Reorder cache for beam search (placeholder)"""
+        return past_key_values
+    
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        max_length: int = 128,
+        num_beams: int = 1,
+        temperature: float = 1.0,
+        do_sample: bool = False,
+        top_k: int = 50,
+        top_p: float = 1.0,
+        **kwargs
+    ) -> torch.LongTensor:
+        """
+        Generate translations
+        
+        Args:
+            input_ids: Source sequence [batch_size, src_seq_len]
+            attention_mask: Source attention mask
+            max_length: Maximum generation length
+            num_beams: Number of beams for beam search
+            temperature: Sampling temperature
+            do_sample: Whether to use sampling
+            top_k: Top-k sampling parameter
+            top_p: Nucleus sampling parameter
+        
+        Returns:
+            Generated sequences [batch_size, tgt_seq_len]
+        """
+        device = input_ids.device
+        batch_size = input_ids.size(0)
+        
+        # Start with BOS token
+        decoder_input_ids = torch.full(
+            (batch_size, 1),
+            self.config.bos_token_id,
+            dtype=torch.long,
+            device=device
+        )
+        
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        
+        # Generate tokens one by one
+        for _ in range(max_length - 1):
+            # Forward pass
+            outputs = self.forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                decoder_input_ids=decoder_input_ids,
+                return_dict=True
+            )
+            
+            # Get next token logits
+            next_token_logits = outputs.logits[:, -1, :] / temperature
+            
+            if do_sample:
+                # Apply top-k filtering
+                if top_k > 0:
+                    indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                    next_token_logits[indices_to_remove] = float('-inf')
+                
+                # Apply top-p (nucleus) filtering
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+                    next_token_logits[indices_to_remove] = float('-inf')
+                
+                # Sample
+                probs = torch.softmax(next_token_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                # Greedy selection
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+            
+            # Mark finished sequences (those that generated EOS)
+            finished = finished | (next_token.squeeze(-1) == self.config.eos_token_id)
+            
+            # Replace tokens in finished sequences with PAD
+            next_token[finished] = self.config.pad_token_id
+            
+            # Append to decoder input
+            decoder_input_ids = torch.cat([decoder_input_ids, next_token], dim=1)
+            
+            # Stop if all sequences are finished
+            if finished.all():
+                break
+        
+        return decoder_input_ids
+
+
+# Register the model in the AutoModel registry
+from transformers import AutoConfig, AutoModel, AutoModelForSeq2SeqLM
+
+AutoConfig.register("translation_transformer", TranslationTransformerConfig)
+AutoModel.register(TranslationTransformerConfig, TranslationTransformerModel)
+AutoModelForSeq2SeqLM.register(TranslationTransformerConfig, TranslationTransformerModel)