core/unified_model.py

"""
Unified Intelligent Tokenizer Model v6.1.2
Compression-First Learning with Adaptive Splitting
- 64 byte chunks for aggressive compression
- 50 epoch checkpoints with automatic splitting
- Group relation learning for reconstruction
- Boundary adjustment for semantic units
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Dict, List, Optional, Tuple, Union


class PositionalEncoding(nn.Module):
    """
    Sinusoidal Positional Encoding (Transformer 원본 방식)
    학습 가능한 위치 임베딩 대신 고정된 sin/cos 패턴 사용
    """
    
    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
        super().__init__()
        self.dropout = nn.Dropout(dropout)
        
        # Create sinusoidal position encodings
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, 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)  # Even dimensions
        pe[:, 1::2] = torch.cos(position * div_term)  # Odd dimensions
        
        # Register as buffer (not trainable)
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Add positional encoding to input
        Args:
            x: (batch_size, seq_len, d_model)
        """
        x = x + self.pe[:, :x.size(1)]
        return self.dropout(x)


class ByteTokenizer:
    """
    Pure byte-level tokenizer - no language rules
    """
    
    def __init__(self, max_seq_len: int = 64):  # v6.1.2: 64 bytes for compression-first approach
        self.max_seq_len = max_seq_len
        self.PAD = 256
        self.BOS = 257
        self.EOS = 258
        self.MASK = 259
    
    def encode(self, text: str, add_special_tokens: bool = True) -> Dict[str, torch.Tensor]:
        # Convert to UTF-8 bytes
        byte_seq = list(text.encode('utf-8'))
        
        # Truncate if needed
        if len(byte_seq) > self.max_seq_len - 2:
            byte_seq = byte_seq[:self.max_seq_len - 2]
        
        # Add special tokens
        if add_special_tokens:
            byte_seq = [self.BOS] + byte_seq + [self.EOS]
        
        input_ids = torch.tensor(byte_seq, dtype=torch.long)
        attention_mask = torch.ones_like(input_ids)
        
        return {
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'length': len(input_ids)
        }
    
    def encode_batch(self, texts: List[str]) -> Dict[str, torch.Tensor]:
        encoded = [self.encode(text) for text in texts]
        max_len = min(max(e['length'] for e in encoded), self.max_seq_len)
        
        batch_size = len(texts)
        input_ids = torch.full((batch_size, max_len), self.PAD, dtype=torch.long)
        attention_mask = torch.zeros((batch_size, max_len), dtype=torch.float32)
        
        for i, enc in enumerate(encoded):
            seq_len = min(enc['length'], max_len)
            input_ids[i, :seq_len] = enc['input_ids'][:seq_len]
            attention_mask[i, :seq_len] = 1.0
        
        return {
            'input_ids': input_ids,
            'attention_mask': attention_mask
        }
    
    def decode(self, input_ids: torch.Tensor, skip_special_tokens: bool = True) -> str:
        if isinstance(input_ids, torch.Tensor):
            input_ids = input_ids.cpu().numpy().tolist()
        
        if skip_special_tokens:
            input_ids = [b for b in input_ids if b < 256]
        
        try:
            byte_array = bytes([min(b, 255) for b in input_ids if b != self.PAD])
            return byte_array.decode('utf-8', errors='replace')
        except:
            return "".join([chr(b) if b < 128 else '?' for b in input_ids if b < 256])


class ByteEncoderV61(nn.Module):
    """
    v6.1: 5-Layer Encoder with Layer-Specialized Architecture
    Layer 0: 768d - Byte to character (with curriculum learning)
    Layer 1: 896d - Language pattern discovery (no labels)
    Layer 2: 1024d - Eojeol/Word formation (+ eojeol PE)
    Layer 3: 1152d - Small phrase grouping (2-3 eojeols)
    Layer 4: 1280d - Final refinement (+ context PE)

    Target: 어절(eojeol) to 구(phrase) level compression (3:1 ratio)
    """

    def __init__(
        self,
        vocab_size: int = 260,
        hidden_dims: List[int] = [768, 896, 1024, 1152, 1280],  # v6.1 dimensions
        num_heads: List[int] = [12, 14, 16, 18, 20],  # v6.1: Progressive heads per layer
        dropout: float = 0.1,
        max_seq_len: int = 64  # v6.1.2: 64 chunk for compression-first
    ):
        super().__init__()
        
        # Layer 0: Byte to Character with Curriculum Learning
        self.byte_embedding = nn.Embedding(vocab_size, hidden_dims[0])

        # v6.1: Multi-level boundary predictors for hierarchical segmentation
        # Level 1: Character boundaries (UTF-8 multi-byte)
        self.char_boundary_predictor = nn.Linear(hidden_dims[0], 3)  # 0: continue, 1: start, 2: end

        # Level 2: Eojeol boundaries (space + particle analysis)
        self.eojeol_boundary_predictor = nn.Linear(hidden_dims[2], 4)  # 0: inside, 1: space, 2: particle, 3: punct

        # Level 3: Phrase boundaries (syntactic chunks)
        self.phrase_boundary_predictor = nn.Linear(hidden_dims[3], 3)  # 0: inside, 1: weak boundary, 2: strong boundary

        # v6.1: Positional encoding ONLY for Layer 0
        self.pos_encoding = PositionalEncoding(hidden_dims[0], max_seq_len, dropout)

        # v6.1: Layer 1 - Language pattern discovery (no labels!)
        self.pattern_discoverer = nn.Linear(hidden_dims[1], 256)  # Discover patterns autonomously (from 896d)
        self.lang_signal_generator = nn.Linear(hidden_dims[1], 128)  # Generate language signals (from 896d)

        # v6.1: Group-aware relative position encodings for Layer 2-4
        self.group_pe_layer2 = nn.Embedding(max_seq_len, hidden_dims[2])  # For eojeol/word units
        self.group_pe_layer3 = nn.Embedding(max_seq_len, hidden_dims[3])  # For small phrases (2-3 eojeols)
        self.group_pe_layer4 = nn.Embedding(max_seq_len, hidden_dims[4])  # For context/discourse

        # 5 Transformer layers with dimension changes
        self.layers = nn.ModuleList()
        for i in range(len(hidden_dims)):
            input_dim = hidden_dims[i-1] if i > 0 else hidden_dims[0]
            output_dim = hidden_dims[i]

            # Projection layer if dimension changes
            if input_dim != output_dim:
                proj = nn.Linear(input_dim, output_dim)
            else:
                proj = None

            # v6.1: Layer-specific head count for optimal dimension per head
            # Target: 64-80 dim per head
            layer_heads = num_heads[i] if isinstance(num_heads, list) else num_heads

            # Transformer encoder layer
            layer = nn.TransformerEncoderLayer(
                d_model=output_dim,
                nhead=layer_heads,
                dim_feedforward=output_dim * 4,
                dropout=dropout,
                activation='gelu',
                batch_first=True,
                norm_first=True
            )
            
            self.layers.append(nn.ModuleDict({
                'projection': proj,
                'transformer': layer,
                'norm': nn.LayerNorm(output_dim)
            }))
        
        self.dropout = nn.Dropout(dropout)
        
    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        boundary_labels: Optional[torch.Tensor] = None,
        epoch: int = 0
    ) -> Dict[str, torch.Tensor]:
        """
        v6.1 Forward pass with curriculum learning
        Args:
            boundary_labels: UTF-8 boundary labels for curriculum learning (training only)
            epoch: Current epoch for curriculum schedule
        """
        batch_size, seq_len = input_ids.shape

        # Layer 0: Byte embedding with curriculum learning
        x = self.byte_embedding(input_ids)

        # v6.1: Positional encoding ONLY at Layer 0
        x = self.pos_encoding(x)

        # v6.1: Predict character boundaries (Layer 0)
        char_boundaries = self.char_boundary_predictor(x)

        # v6.1: Curriculum learning for character boundaries
        # Note: boundary_labels are eojeol boundaries (4 classes), not char boundaries (3 classes)
        # So we don't mix them with char_boundaries - they serve different purposes
        char_boundary_weights = F.softmax(char_boundaries, dim=-1)
        
        # Prepare attention mask
        if attention_mask is not None:
            # Keep attention mask as is for TransformerEncoderLayer
            # It expects shape (batch_size, seq_len) and handles masking internally
            pass
        
        # v6.1: Process through 5 specialized layers
        all_hidden_states = []
        discovered_patterns = None
        eojeol_boundaries = None
        phrase_boundaries = None

        for i, layer_dict in enumerate(self.layers):
            # Project if needed (before layer-specific processing)
            if layer_dict['projection'] is not None:
                x = layer_dict['projection'](x)

            # Layer 1: Add language signals (autonomous discovery)
            if i == 1:
                # Discover language patterns WITHOUT labels (x is now 896d)
                discovered_patterns = self.pattern_discoverer(x)
                lang_signals = self.lang_signal_generator(x)

            # Layer 2: Predict eojeol boundaries and add position encoding
            elif i == 2:
                # Predict eojeol boundaries (spaces, particles, punctuation)
                eojeol_boundaries = self.eojeol_boundary_predictor(x)
                positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
                group_pe = self.group_pe_layer2(positions)
                x = x + group_pe * 0.1  # Mild addition to preserve main signal

            # Layer 3: Predict phrase boundaries and add position encoding
            elif i == 3:
                # Predict phrase boundaries (weak/strong syntactic breaks)
                phrase_boundaries = self.phrase_boundary_predictor(x)
                positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
                group_pe = self.group_pe_layer3(positions)
                x = x + group_pe * 0.1

            elif i == 4:
                positions = torch.arange(seq_len, device=x.device).unsqueeze(0).expand(batch_size, -1)
                group_pe = self.group_pe_layer4(positions)
                x = x + group_pe * 0.1

            # Transformer layer - properly handle mask
            if attention_mask is not None:
                key_padding_mask = (attention_mask == 0)
                x = layer_dict['transformer'](x, src_key_padding_mask=key_padding_mask)
            else:
                x = layer_dict['transformer'](x)
            x = layer_dict['norm'](x)
            all_hidden_states.append(x)
        
        # Pool for sequence representation
        if attention_mask is not None:
            # Masked mean pooling - attention_mask is (batch, seq)
            mask = attention_mask.unsqueeze(-1)  # (batch, seq, 1)
            pooled = (x * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
        else:
            pooled = x.mean(dim=1)
        
        return {
            'last_hidden_state': x,
            'pooled_output': pooled,
            'all_hidden_states': all_hidden_states,
            # v6.1 boundary predictions
            'char_boundaries': char_boundaries,
            'char_boundary_weights': char_boundary_weights,
            'eojeol_boundaries': eojeol_boundaries,
            'phrase_boundaries': phrase_boundaries,
            'discovered_patterns': discovered_patterns
        }


class CrossAttention(nn.Module):
    """
    Enhanced Cross-attention for relation learning between sequences
    추론 레이어 연결을 위한 강화된 관계 학습
    """
    
    def __init__(self, hidden_dim: int = 1280, num_heads: int = 20, dropout: float = 0.1):
        super().__init__()

        # v6.1: Adjusted for 1280d (64 per head with 20 heads)
        self.cross_attn = nn.MultiheadAttention(
            hidden_dim, num_heads, dropout, batch_first=True
        )
        
        # v6.1: Enhanced relation classifier with reconstruction focus
        # 0: identity (완벽한 복원), 1: similar, 2: different, 3: continuation
        # 4: translation, 5: summary, 6: expansion, 7: contradiction
        self.relation_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 8)
        )

        # v6.1: Reconstruction-specific attention (복원 전용 어텐션)
        # Use 10 heads for reconstruction (128 per head)
        self.reconstruction_attn = nn.MultiheadAttention(
            hidden_dim, 10, dropout * 0.5, batch_first=True
        )
        
        # Gating mechanism for adaptive fusion
        self.gate = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.Sigmoid()
        )
        
        self.norm1 = nn.LayerNorm(hidden_dim)
        self.norm2 = nn.LayerNorm(hidden_dim)
        
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        query_mask: Optional[torch.Tensor] = None,
        key_mask: Optional[torch.Tensor] = None
    ) -> Dict[str, torch.Tensor]:
        # Normalize inputs
        query_norm = self.norm1(query)
        key_norm = self.norm2(key)
        
        # Fix key_mask dimension if needed
        if key_mask is not None:
            # Ensure key_mask matches key sequence length
            if key_mask.dim() == 2 and key_mask.size(1) != key.size(1):
                # Create new mask with correct dimensions
                batch_size = key.size(0)
                seq_len = key.size(1)
                key_mask = torch.ones(batch_size, seq_len, dtype=key_mask.dtype, device=key_mask.device)
        
        # Cross attention
        attn_output, attn_weights = self.cross_attn(
            query_norm, key_norm, key_norm,
            key_padding_mask=(key_mask == 0) if key_mask is not None else None
        )
        
        # Residual connection
        attn_output = attn_output + query

        # v6.1: Reconstruction-focused attention (복원 최적화)
        recon_output, recon_weights = self.reconstruction_attn(
            query_norm, query_norm, query_norm,  # Self-attention for consistency
            key_padding_mask=(query_mask == 0) if query_mask is not None else None
        )

        # Combine cross and reconstruction attention
        combined_attn = attn_output * 0.7 + recon_output * 0.3

        # Adaptive gating for fusion
        gate_input = torch.cat([query.mean(dim=1), key.mean(dim=1)], dim=-1)
        gate_weights = self.gate(gate_input).unsqueeze(1)

        # Gated fusion: 적응적으로 attention 결과 조절
        fused_output = gate_weights * combined_attn + (1 - gate_weights) * query
        
        # Pool for relation classification
        query_pooled = query.mean(dim=1) if query_mask is None else \
                      (query * query_mask.unsqueeze(-1)).sum(1) / query_mask.sum(1, keepdim=True).clamp(min=1e-9)
        key_pooled = key.mean(dim=1) if key_mask is None else \
                    (key * key_mask.unsqueeze(-1)).sum(1) / key_mask.sum(1, keepdim=True).clamp(min=1e-9)
        
        # Classify relations with enhanced head
        combined = torch.cat([query_pooled, key_pooled], dim=-1)
        relation_logits = self.relation_head(combined)
        
        return {
            'cross_attention': fused_output,  # Gated fusion output
            'attention_weights': attn_weights,
            'reconstruction_weights': recon_weights,  # v6.1: 복원 어텐션 가중치
            'relation_logits': relation_logits,
            'gate_weights': gate_weights.squeeze(1),  # For analysis
            'reconstruction_score': F.softmax(relation_logits, dim=-1)[:, 0]  # identity 확률 (복원도)
        }


class TransformerDecoder(nn.Module):
    """
    Transformer Decoder with Positional Encoding
    """

    def __init__(
        self,
        vocab_size: int = 260,
        hidden_dim: int = 1280,  # v6.1: Match final encoder dim
        num_heads: int = 16,      # v6.1: 1280/16 = 80 per head
        num_layers: int = 8,      # v6.1 FINAL: 8 layers for better reconstruction
        dropout: float = 0.1,
        max_seq_len: int = 64    # v6.1.2: 64 chunk for compression-first
    ):
        super().__init__()
        
        # Token embedding
        self.token_embedding = nn.Embedding(vocab_size, hidden_dim)
        
        # Positional encoding
        self.pos_encoding = PositionalEncoding(hidden_dim, max_seq_len, dropout)
        
        # Transformer decoder
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=hidden_dim * 4,
            dropout=dropout,
            activation='gelu',
            batch_first=True,
            norm_first=True
        )
        
        self.transformer = nn.TransformerDecoder(decoder_layer, num_layers)
        
        # Output projection
        self.output_projection = nn.Linear(hidden_dim, vocab_size)
        
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        
    def forward(
        self,
        encoder_hidden: torch.Tensor,
        decoder_input_ids: Optional[torch.Tensor] = None,
        encoder_mask: Optional[torch.Tensor] = None,
        decoder_mask: Optional[torch.Tensor] = None
    ) -> Dict[str, torch.Tensor]:
        batch_size = encoder_hidden.size(0)
        
        # Start with BOS if no input
        if decoder_input_ids is None:
            decoder_input_ids = torch.full((batch_size, 1), 257, device=encoder_hidden.device)
        
        # Embed and add positional encoding
        dec_seq_len = decoder_input_ids.size(1)
        x = self.token_embedding(decoder_input_ids)
        x = self.pos_encoding(x)
        
        # Create causal mask
        causal_mask = torch.triu(
            torch.ones(dec_seq_len, dec_seq_len, device=x.device) * float('-inf'),
            diagonal=1
        )
        
        # Decoder forward - handle variable-length encoder outputs
        # The encoder may compress the sequence, so memory (encoder_hidden) might be shorter
        # than the decoder sequence. This is expected and correct behavior.
        enc_seq_len = encoder_hidden.size(1)
        
        # Adjust encoder mask if needed
        if encoder_mask is not None:
            if encoder_mask.size(1) != enc_seq_len:
                # Encoder compressed the sequence, create new mask for compressed length
                # All compressed positions are valid (not masked)
                memory_key_padding_mask = torch.zeros(
                    encoder_hidden.size(0), enc_seq_len,
                    dtype=torch.bool, device=encoder_hidden.device
                )
            else:
                memory_key_padding_mask = (encoder_mask == 0)
        else:
            memory_key_padding_mask = None
            
        # Decoder attends to compressed encoder states via cross-attention
        # This naturally handles different sequence lengths
        decoder_output = self.transformer(
            tgt=x,  # Decoder sequence (original length)
            memory=encoder_hidden,  # Encoder sequence (possibly compressed)
            tgt_mask=causal_mask,
            memory_key_padding_mask=memory_key_padding_mask,
            tgt_key_padding_mask=(decoder_mask == 0) if decoder_mask is not None else None
        )
        
        # Project to vocabulary
        logits = self.output_projection(decoder_output)
        
        return {
            'logits': logits,
            'hidden_states': decoder_output
        }
    
    @torch.no_grad()
    def generate(
        self,
        encoder_hidden: torch.Tensor,
        encoder_mask: Optional[torch.Tensor] = None,
        max_length: int = 128,
        temperature: float = 0.1,  # 토크나이저는 보수적 생성 (정확한 복원)
        top_k: int = 10,  # 상위 10개만 고려
        top_p: float = 0.95
    ) -> torch.Tensor:
        batch_size = encoder_hidden.size(0)
        device = encoder_hidden.device

        # Start with BOS
        decoder_input_ids = torch.full((batch_size, 1), 257, device=device)

        # Track which sequences are done
        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)

        for _ in range(max_length - 1):
            # Forward pass
            outputs = self.forward(encoder_hidden, decoder_input_ids, encoder_mask)
            next_token_logits = outputs['logits'][:, -1, :] / temperature

            # 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')

            # Top-p filtering
            if top_p < 1.0:
                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
                cumulative_probs = torch.cumsum(F.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 = F.softmax(next_token_logits, dim=-1)
            next_tokens = torch.multinomial(probs, 1)

            # For finished sequences, force PAD token
            next_tokens[finished] = 256  # PAD token

            decoder_input_ids = torch.cat([decoder_input_ids, next_tokens], dim=-1)

            # Update finished status
            finished = finished | (next_tokens.squeeze(-1) == 258)  # Mark as finished if EOS

            # Stop when all sequences are done
            if finished.all():
                break

        return decoder_input_ids


class IntelligentTokenizerModelV61(nn.Module):
    """
    Complete Intelligent Tokenizer Model v6.1
    Pure learning-based with curriculum learning
    - No language labels during training
    - Curriculum learning for boundaries
    - Group-aware position encodings
    """

    def __init__(
        self,
        vocab_size: int = 260,
        encoder_dims: List[int] = [768, 896, 1024, 1152, 1280],  # v6.1 dimensions
        encoder_heads: List[int] = [12, 14, 16, 18, 20],  # v6.1: Optimal heads per layer
        decoder_hidden: int = 1280,  # Match final encoder dim
        decoder_heads: int = 16,     # v6.1: 80 per head for decoder
        num_decoder_layers: int = 8,  # v6.1 FINAL: 8 layers for better reconstruction
        dropout: float = 0.1,
        max_seq_len: int = 64  # v6.1.2: 64 chunk for compression-first
    ):
        super().__init__()

        # v6.1 Components with optimized head counts
        self.tokenizer = ByteTokenizer(max_seq_len)
        self.encoder = ByteEncoderV61(vocab_size, encoder_dims, encoder_heads, dropout, max_seq_len)
        self.decoder = TransformerDecoder(vocab_size, decoder_hidden, decoder_heads, num_decoder_layers, dropout, max_seq_len)
        self.cross_attention = CrossAttention(encoder_dims[-1], 20, dropout)  # 20 heads for 1280d
        
    def forward(
        self,
        input_texts: Optional[List[str]] = None,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        decoder_input_ids: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        boundary_labels: Optional[torch.Tensor] = None,  # v6.1: for curriculum learning
        epoch: int = 0,  # v6.1: for curriculum schedule
        use_cross_attention: bool = True
    ) -> Dict[str, torch.Tensor]:
        # Tokenize if text input
        if input_texts is not None:
            tokenized = self.tokenizer.encode_batch(input_texts)
            input_ids = tokenized['input_ids']
            attention_mask = tokenized['attention_mask']
        
        # 시퀀스 길이 체크 및 조정
        batch_size, seq_len = input_ids.shape
        device = input_ids.device
        
        # v6.1: Encode with curriculum learning
        encoder_outputs = self.encoder(input_ids, attention_mask, boundary_labels, epoch)
        encoder_hidden = encoder_outputs['last_hidden_state']  # v6.1: [batch, seq, 1280]

        # v6.1: 차원 확인 - 최종 차원은 1280
        assert encoder_hidden.size(-1) == 1280, f"Encoder dim mismatch: {encoder_hidden.size(-1)}"

        # Prepare decoder input for teacher forcing during training
        if decoder_input_ids is None:
            if labels is not None:
                # During training, use shifted labels as decoder input (teacher forcing)
                # Add BOS at the beginning and remove last token
                bos_tokens = torch.full((batch_size, 1), self.tokenizer.BOS, device=labels.device, dtype=labels.dtype)
                decoder_input_ids = torch.cat([bos_tokens, labels[:, :-1]], dim=1)
            else:
                # For inference/test, start with BOS token
                decoder_input_ids = torch.full((batch_size, 1), self.tokenizer.BOS, device=device, dtype=torch.long)

        # Decode
        decoder_outputs = self.decoder(
            encoder_hidden,
            decoder_input_ids,
            attention_mask
        )
        decoder_hidden = decoder_outputs['hidden_states']  # [batch, seq, 768]
        
        # Cross-Attention (마지막 레이어에서 관계 학습)
        cross_attn_outputs = None
        relation_logits = None
        
        if use_cross_attention and decoder_hidden is not None:
            # 디코더 출력과 인코더 출력 간 크로스어텐션
            cross_attn_outputs = self.cross_attention(
                query=decoder_hidden,  # 디코더가 query
                key=encoder_hidden,     # 인코더가 key/value
                query_mask=None,        # decoder mask는 causal이므로 별도 처리
                key_mask=attention_mask
            )
            
            # 관계 학습 결과
            relation_logits = cross_attn_outputs['relation_logits']
            
            # Cross-attention으로 강화된 디코더 표현
            enhanced_decoder = decoder_hidden + cross_attn_outputs['cross_attention']
            
            # 최종 로짓 재계산 (cross-attention 적용 후)
            if hasattr(self.decoder, 'output_projection'):
                decoder_outputs['logits'] = self.decoder.output_projection(enhanced_decoder)
        
        # Calculate loss if labels provided
        loss = None
        if labels is not None:
            # Reconstruction loss
            loss_fct = nn.CrossEntropyLoss(ignore_index=self.tokenizer.PAD)
            recon_loss = loss_fct(
                decoder_outputs['logits'].reshape(-1, decoder_outputs['logits'].size(-1)),
                labels.reshape(-1)
            )

            # Boundary loss (if boundary labels provided)
            boundary_loss = 0
            if boundary_labels is not None and encoder_outputs.get('eojeol_boundaries') is not None:
                # Eojeol boundary loss
                eojeol_boundaries = encoder_outputs['eojeol_boundaries']  # [batch, seq, 4]
                if eojeol_boundaries.size(1) == boundary_labels.size(1):
                    # Ensure boundary labels are in valid range (0-3)
                    # Clamp to valid range to prevent CUDA errors
                    boundary_labels_clamped = torch.clamp(boundary_labels, min=0, max=3)

                    boundary_loss_fct = nn.CrossEntropyLoss(ignore_index=-1)  # Use -1 for padding
                    boundary_loss = boundary_loss_fct(
                        eojeol_boundaries.reshape(-1, 4),
                        boundary_labels_clamped.reshape(-1)
                    ) * 0.5  # Weight for boundary loss

            # Relation loss (if cross-attention used)
            relation_loss = 0
            if relation_logits is not None:
                # 자기 관계는 identity (class 0)여야 함
                batch_identity = torch.zeros(batch_size, dtype=torch.long, device=device)
                relation_loss = F.cross_entropy(relation_logits, batch_identity) * 0.1

            loss = recon_loss + boundary_loss + relation_loss
        
        return {
            'loss': loss,
            'logits': decoder_outputs['logits'],
            'decoder_logits': decoder_outputs['logits'],  # Add for compatibility
            'encoder_hidden_states': encoder_hidden,
            'decoder_hidden_states': decoder_hidden,
            'pooled_output': encoder_outputs['pooled_output'],
            'cross_attention': cross_attn_outputs['cross_attention'] if cross_attn_outputs else None,
            'relation_logits': relation_logits,
            'all_encoder_states': encoder_outputs.get('all_hidden_states', None),
            # Add boundary predictions for visualization
            'char_boundaries': encoder_outputs.get('char_boundaries'),
            'eojeol_boundaries': encoder_outputs.get('eojeol_boundaries'),
            'phrase_boundaries': encoder_outputs.get('phrase_boundaries'),
            'discovered_patterns': encoder_outputs.get('discovered_patterns')
        }
    
    def encode_text(self, text: str) -> torch.Tensor:
        """Encode single text to representation"""
        tokenized = self.tokenizer.encode(text)
        # Move to same device as model
        device = next(self.parameters()).device
        input_ids = tokenized['input_ids'].unsqueeze(0).to(device)
        attention_mask = tokenized['attention_mask'].unsqueeze(0).to(device)
        
        with torch.no_grad():
            outputs = self.encoder(input_ids, attention_mask)
        
        return outputs['pooled_output'].squeeze(0)
    
    def decode_representation(self, representation: torch.Tensor, max_length: int = 128) -> str:
        """Decode representation back to text"""
        if representation.dim() == 1:
            representation = representation.unsqueeze(0).unsqueeze(0)
        elif representation.dim() == 2:
            representation = representation.unsqueeze(1)
        
        with torch.no_grad():
            output_ids = self.decoder.generate(representation, max_length=max_length)
        
        text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
        return text
    
    def compute_relation(self, text1: str, text2: str) -> torch.Tensor:
        """Compute relation between two texts"""
        # Encode both texts
        enc1 = self.encode_text(text1).unsqueeze(0).unsqueeze(0)
        enc2 = self.encode_text(text2).unsqueeze(0).unsqueeze(0)
        
        # Compute cross-attention and relations
        with torch.no_grad():
            outputs = self.cross_attention(enc1, enc2)
        
        return F.softmax(outputs['relation_logits'], dim=-1)