modeling_baseline.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
import os
from x_transformers import Encoder, Decoder
from transformers import AutoTokenizer

# --- SMART TOKENIZER SETUP ---
try:
    if os.path.exists("tokenizer_config.json"):
        tokenizer = AutoTokenizer.from_pretrained(".")
    else:
        tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-de-en")
except Exception as e:
    print(f"Warning: Tokenizer load failed: {e}")
# -----------------------------


class RoPETransformer(nn.Module):
    def __init__(self, num_encoder_layers, num_decoder_layers, num_heads, d_model, dff, vocab_size, max_length, dropout):
        super().__init__()
        self.d_model = d_model
        self.embedding = nn.Embedding(vocab_size, d_model)

        # We REMOVE self.pos_encoder (RoPE handles position internally)
        self.dropout_layer = nn.Dropout(dropout)

        # --- x-transformers Encoder ---
        self.encoder = Encoder(
            dim = d_model,
            depth = num_encoder_layers,
            heads = num_heads,
            attn_dim_head = d_model // num_heads,
            ff_mult = dff / d_model,
            rotary_pos_emb = True,
            attn_flash = True,
            attn_dropout = dropout,
            ff_dropout = dropout,
            use_rmsnorm = True
        )

        # --- x-transformers Decoder ---
        self.decoder = Decoder(
            dim = d_model,
            depth = num_decoder_layers,
            heads = num_heads,
            attn_dim_head = d_model // num_heads,
            ff_mult = dff / d_model,
            rotary_pos_emb = True,
            cross_attend = True,
            attn_flash = True,
            attn_dropout = dropout,
            ff_dropout = dropout,
            use_rmsnorm = True
        )

        self.final_linear = nn.Linear(d_model, vocab_size)
        self.final_linear.weight = self.embedding.weight

    def forward(self, src, tgt, src_padding_mask, tgt_padding_mask, memory_key_padding_mask, tgt_mask):
        # 1. Embeddings (No Absolute Positional Encoding added!)
        src_emb = self.embedding(src) * math.sqrt(self.d_model)
        src_emb = self.dropout_layer(src_emb)

        tgt_emb = self.embedding(tgt) * math.sqrt(self.d_model)
        tgt_emb = self.dropout_layer(tgt_emb)

        # 2. Mask Conversion
        # User provides True=PAD. x-transformers wants True=KEEP.
        # We invert the boolean mask using ~
        enc_mask = ~src_padding_mask if src_padding_mask is not None else None
        dec_mask = ~tgt_padding_mask if tgt_padding_mask is not None else None

        # Note: 'tgt_mask' (causal mask) is handled automatically by x-transformers Decoder!
        # We do NOT pass the square causal mask manually.

        # 3. Encoder
        # x-transformers takes embeddings directly
        memory = self.encoder(src_emb, mask=enc_mask)

        # 4. Decoder
        # context = memory (from encoder)
        # context_mask = mask for memory (encoder mask)
        decoder_output = self.decoder(
            tgt_emb,
            context=memory,
            mask=dec_mask,
            context_mask=enc_mask
        )

        return self.final_linear(decoder_output)

    # Keep your existing create_masks (used for Data Processing mostly)
    def create_masks(self, src, tgt):
        src_padding_mask = (src == tokenizer.pad_token_id)
        tgt_padding_mask = (tgt == tokenizer.pad_token_id)
        # We still generate this for compatibility, though x-transformers handles causality internally
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(
            sz=tgt.size(1), device=src.device, dtype=torch.bool
        )
        return src_padding_mask, tgt_padding_mask, src_padding_mask, tgt_mask

    @torch.no_grad()
    def generate(self, src: torch.Tensor, max_length: int, num_beams: int = 5) -> torch.Tensor:
        self.eval()
        # Create Mask (True=PAD)
        src_padding_mask = (src == tokenizer.pad_token_id)
        # Invert for x-transformers (True=KEEP)
        enc_mask = ~src_padding_mask

        # Encode
        src_emb = self.embedding(src) * math.sqrt(self.d_model)
        # No Pos Encoder
        memory = self.encoder(self.dropout_layer(src_emb), mask=enc_mask)

        batch_size = src.shape[0]
        # Expand for beams
        memory = memory.repeat_interleave(num_beams, dim=0)
        enc_mask = enc_mask.repeat_interleave(num_beams, dim=0)

        initial_token = tokenizer.pad_token_id
        beams = torch.full((batch_size * num_beams, 1), initial_token, dtype=torch.long, device=src.device)
        beam_scores = torch.zeros(batch_size * num_beams, device=src.device)
        finished_beams = torch.zeros(batch_size * num_beams, dtype=torch.bool, device=src.device)

        for _ in range(max_length - 1):
            if finished_beams.all(): break

            # Embed beams
            tgt_emb = self.embedding(beams) * math.sqrt(self.d_model)
            # No Pos Encoder

            # Decode
            # x-transformers automatically handles the causal masking for the sequence length of tgt_emb
            decoder_output = self.decoder(
                self.dropout_layer(tgt_emb),
                context=memory,
                context_mask=enc_mask
            )

            logits = self.final_linear(decoder_output[:, -1, :])
            log_probs = F.log_softmax(logits, dim=-1)

            # ... (Rest of your Beam Search Logic remains identical) ...
            log_probs[:, tokenizer.pad_token_id] = -torch.inf
            if finished_beams.any(): log_probs[finished_beams, tokenizer.eos_token_id] = 0

            total_scores = beam_scores.unsqueeze(1) + log_probs
            if _ == 0:
                total_scores = total_scores.view(batch_size, num_beams, -1)
                total_scores[:, 1:, :] = -torch.inf
                total_scores = total_scores.view(batch_size * num_beams, -1)
            else:
                total_scores = beam_scores.unsqueeze(1) + log_probs

            total_scores = total_scores.view(batch_size, -1)
            top_scores, top_indices = torch.topk(total_scores, k=num_beams, dim=1)

            beam_indices = top_indices // log_probs.shape[-1]
            token_indices = top_indices % log_probs.shape[-1]

            batch_indices = torch.arange(batch_size, device=src.device).unsqueeze(1)
            effective_indices = (batch_indices * num_beams + beam_indices).view(-1)

            beams = beams[effective_indices]
            beams = torch.cat([beams, token_indices.view(-1, 1)], dim=1)
            beam_scores = top_scores.view(-1)
            finished_beams = finished_beams | (beams[:, -1] == tokenizer.eos_token_id)

        final_beams = beams.view(batch_size, num_beams, -1)
        final_scores = beam_scores.view(batch_size, num_beams)
        normalized_scores = final_scores / (final_beams != tokenizer.pad_token_id).sum(-1).float().clamp(min=1)
        best_beams = final_beams[torch.arange(batch_size), normalized_scores.argmax(1), :]
        self.train()
        return best_beams