Create model.py

model.py

@@ -0,0 +1,815 @@
+import os
+import re
+import math
+import pickle
+import random
+from collections import defaultdict, namedtuple
+
+import numpy as np
+from PIL import Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, models
+from tqdm import tqdm
+
+# ================ HYPERPARAMETERS ================
+IMG_HEIGHT = 256
+IMG_WIDTH = 256
+MAX_SEQ_LEN = 512
+NUM_NESTED_LEVELS = 10
+NUM_REL_POS = 3
+PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN = '<PAD>', '<SOS>', '<EOS>', '<UNK>'
+BATCH_SIZE = 8
+NUM_EPOCHS = 300
+MAX_LEARNING_RATE = 3e-4
+WARMUP_RATIO = 0.1
+WEIGHT_DECAY = 0.01
+LAMBDA_POS = 0.2
+LABEL_SMOOTHING = 0.0
+DROPOUT_RATE = 0.3
+NUM_LAYERS = 3
+TEMPERATURE_INIT = 1.0
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+SEED = 42
+BEAM_SIZE = 10
+MAX_GEN_LEN = MAX_SEQ_LEN
+
+# ===== Fix seeds for reproducibility =====
+random.seed(SEED)
+np.random.seed(SEED)
+torch.manual_seed(SEED)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(SEED)
+
+import re
+from typing import List, Tuple, Union, Optional, Dict
+import numpy as np
+
+
+# Enhanced LaTeX tokenizer with better regex patterns
+TOKEN_REGEX = re.compile(
+    r"(\\[a-zA-Z]+(?:\*)?)"       # multi-letter commands (e.g. \frac, \sqrt, \begin*)
+    r"|(\\\{|\\\}|\\.)"           # escaped chars and single-char commands
+    r"|([0-9]+(?:\.[0-9]+)?(?:[eE][+-]?[0-9]+)?)"  # numbers (int, float, scientific)
+    r"|([A-Za-z]+)"               # identifiers/variables
+    r"|([+\-*/=<>!~^_&|%])"       # operators and structure markers
+    r"|([{}()\[\],.;:?'])"        # delimiters and punctuation
+    r"|(\s+)"                     # whitespace (to handle properly)
+    r"|(\S)"                      # any other non-space char
+)
+
+def tokenize_latex(s: str) -> List[str]:
+    """
+    Split a LaTeX string into atomic tokens.
+    Filters out whitespace tokens.
+    """
+    # Исправление обработки специальных символов
+    s = re.sub(r'\\ ', ' ', s)  # Обработка пробелов
+    s = re.sub(r'\\\n', '', s)  # Удаление переносов
+
+    tokens = []
+    for match in TOKEN_REGEX.finditer(s):
+        token = match.group(0)
+        if token.isspace():
+            continue
+        # Это необязательная, но потенциально полезная эвристика для \left{ и \right}
+        # Если она вызывает проблемы, можно убрать.
+        if token in {'{', '}'} and tokens and tokens[-1] in {'\\left', '\\right'}:
+            tokens[-1] += token
+        else:
+            tokens.append(token)
+
+    return tokens
+
+# Position Forest Implementation
+STRUCTURE_CMDS = {'^', '_', '\\sqrt', '\\frac', '\\sum', '\\int', '\\lim'}
+
+class PositionForestEncoder:
+    def position_forest_ids(self, tokens):
+        cache_key = tuple(tokens)
+        if hasattr(self, '_cache'):
+            if cache_key in self._cache:
+                return self._cache[cache_key]
+        else:
+            self._cache = {}
+        T = len(tokens)
+        ids = ['M']*T
+        def find_close(i):
+            depth = 0
+            for j in range(i, T):
+                if tokens[j]=='{': depth+=1
+                elif tokens[j]=='}':
+                    depth-=1
+                    if depth==0: return j
+            return T-1
+        def mark(l,r,label):
+            for k in range(max(0,l), min(T,r+1)):
+                ids[k]+=label
+        def rec(l,r):
+            i = l
+            while i<=r:
+                tok = tokens[i]
+                # handle ^, _, ...
+                if tok in ('^','_') and i+1<=r:
+                    if tokens[i+1]=='{':
+                        e=find_close(i+1)
+                        lbl='L' if tok=='^' else 'R'
+                        mark(i+2,e-1,lbl)
+                        rec(i+2,e-1)
+                        i=e+1
+                    else:
+                        lbl='L' if tok=='^' else 'R'
+                        mark(i+1,i+1,lbl)
+                        i+=2
+                elif tok=='\\sqrt' and i+1<=r and tokens[i+1]=='{':
+                    e=find_close(i+1)
+                    mark(i+2,e-1,'L'); rec(i+2,e-1)
+                    i=e+1
+                elif tok=='\\frac' and i+1<=r and tokens[i+1]=='{':
+                    n_end=find_close(i+1)
+                    if n_end+1<=r and tokens[n_end+1]=='{':
+                        d_end=find_close(n_end+1)
+                        mark(i+2,n_end-1,'L'); rec(i+2,n_end-1)
+                        mark(n_end+2,d_end-1,'R'); rec(n_end+2,d_end-1)
+                        i=d_end+1
+                    else: i+=1
+                else: i+=1
+        rec(0,T-1)
+        self._cache[cache_key] = ids
+        return ids
+
+    def forest_MLR(self, latex: str):
+        tokens = tokenize_latex(latex)
+        if not tokens: tokens=['']
+        pos_ids = self.position_forest_ids(tokens)
+        nested_depth = [len(pid)-1 for pid in pos_ids]
+        rel_pos = [1 if pid.endswith('L') else 2 if pid.endswith('R') else 0 for pid in pos_ids]
+        # add special tokens
+        tokens = [SOS_TOKEN]+tokens+[EOS_TOKEN]
+        nested_depth = [0]+nested_depth+[0]
+        rel_pos = [0]+rel_pos+[0]
+        # pad/truncate
+        curr=len(tokens)
+        if curr>MAX_SEQ_LEN:
+            tokens=tokens[:MAX_SEQ_LEN]; nested_depth=nested_depth[:MAX_SEQ_LEN]; rel_pos=rel_pos[:MAX_SEQ_LEN]; curr=MAX_SEQ_LEN
+        pad_len=MAX_SEQ_LEN-curr
+        tokens+= [PAD_TOKEN]*pad_len
+        nested_depth+= [0]*pad_len; rel_pos+=[0]*pad_len
+        attention_mask = [1]*curr + [0]*pad_len
+        return tokens, np.array(nested_depth), np.array(rel_pos), np.array(attention_mask)
+
+# ===== Vocabulary =====
+class Vocab:
+    def __init__(self, expressions):
+        freq=defaultdict(int)
+        for expr in expressions:
+            for t in tokenize_latex(expr): freq[t]+=1
+        self.itos=[PAD_TOKEN,SOS_TOKEN,EOS_TOKEN,UNK_TOKEN]+sorted(freq.keys(), key=lambda x:-freq[x])
+        self.stoi={tok:i for i,tok in enumerate(self.itos)}
+    def encode(self, tokens):
+        return [self.stoi.get(t,self.stoi[UNK_TOKEN]) for t in tokens]
+    def decode(self, ids):
+        return ''.join(self.itos[i] for i in ids if self.itos[i] not in (PAD_TOKEN,SOS_TOKEN,EOS_TOKEN))
+
+# ===== Dataset =====
+class CROHMEDataset(Dataset):
+    def __init__(self, base_dir, caption_file, vocab, pfe):
+        self.vocab=vocab; self.pfe=pfe; self.items=[]
+        img_dir=os.path.join(base_dir,'img')
+        with open(os.path.join(base_dir,caption_file),'r',encoding='utf-8') as f:
+            for line in f:
+                fid,lt=line.strip().split('\t',1)
+                self.items.append((os.path.join(img_dir,fid+'.bmp'),lt))
+        self.transform = transforms.Compose([
+            transforms.Resize((IMG_HEIGHT, IMG_WIDTH)),
+            transforms.RandomAffine(degrees=5, translate=(0.05, 0.05), scale=(0.95, 1.05), shear=5),
+            transforms.ColorJitter(brightness=0.3, contrast=0.3),            
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+    def __len__(self): return len(self.items)
+    def __getitem__(self,idx):
+        path,latex=self.items[idx]
+        img=Image.open(path).convert('RGB'); img=self.transform(img)
+        tokens,nested,rel,mask=self.pfe.forest_MLR(latex)
+        token_ids=self.vocab.encode(tokens)
+        return img, {
+            'token_ids': torch.tensor(token_ids,dtype=torch.long),
+            'nested': torch.tensor(nested,dtype=torch.long),
+            'relpos': torch.tensor(rel,dtype=torch.long),
+            'attention_mask': torch.tensor(mask,dtype=torch.long)
+        }, latex
+
+
+# ===== Improved Implicit Attention Correction =====
+class IAC(nn.Module):
+    """
+    Implicit Attention Correction модуль согласно статье
+    Исправленная версия с правильной логикой накопления внимания
+    """
+    def __init__(self, d_model=256, num_heads=8, dropout=DROPOUT_RATE, structure_symbols_set=None,vocab=None):
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+
+        # Структурные символы - те, которые НЕ имеют визуального представления
+        self.structure_symbols = structure_symbols_set or {
+            '^', '{', '}', '_', EOS_TOKEN,UNK_TOKEN,SOS_TOKEN,PAD_TOKEN
+        }
+
+        # Linear projections для Q, K, V
+        self.wq = nn.Linear(d_model, d_model, bias=False)
+        self.wk = nn.Linear(d_model, d_model, bias=False)
+        self.wv = nn.Linear(d_model, d_model, bias=False)
+        self.wo = nn.Linear(d_model, d_model)
+
+        # φ функция для обработки accumulated attention
+        # Используем простую архитектуру: conv + linear
+        self.phi_conv = nn.Conv2d(num_heads, num_heads, kernel_size=3, padding=1, groups=num_heads)
+        self.phi_linear = nn.Linear(num_heads, d_model)
+        self.phi_norm = nn.LayerNorm(d_model)
+
+        self.attn_dropout = nn.Dropout(dropout)
+        self.proj_dropout = nn.Dropout(dropout)
+
+        # Накопленное внимание для коррекции
+        self.register_buffer('accumulated_attention', None, persistent=False)
+        structure_ids_list = [vocab.stoi.get(s, -1) for s in structure_symbols_set]
+        self.register_buffer('structure_ids', 
+                             torch.tensor([sid for sid in structure_ids_list if sid != -1], dtype=torch.long))
+        self._init_weights()
+        
+
+    def _init_weights(self):
+        for m in [self.wq, self.wk, self.wv, self.wo, self.phi_linear]:
+            nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+        nn.init.xavier_uniform_(self.phi_conv.weight)
+        nn.init.constant_(self.phi_conv.bias, 0)
+
+    def reset_state(self):
+        """Сброс состояния IAC - важно для начала новой последовательности"""
+        self.accumulated_attention = None
+
+    def _is_structure_symbol(self, symbol):
+        """
+        Проверяет, является ли символ структурным
+        Структурные символы не имеют визуального представления в изображении
+        """
+        if isinstance(symbol, (list, tuple)):
+            return [sym in self.structure_symbols for sym in symbol]
+        return symbol in self.structure_symbols
+
+    def _compute_phi(self, accumulated_attention):
+        """
+        Вычисляет φ(A^k) - функцию коррекции на основе накопленного внимания.
+        ИСПРАВЛЕННАЯ ВЕРСИЯ для обработки 5D тензора.
+
+        Args:
+            accumulated_attention: [B, H, Lq, Hp, Wp] - история накопленного внимания
+        """
+        # 1. Получаем 5 измерений
+        B, H, Lq, Hp, Wp = accumulated_attention.shape
+
+        # 2. Объединяем B и Lq в одно измерение, чтобы подать в Conv2d.
+        # Conv2d ожидает на вход (N, C_in, H_in, W_in).
+        # Наша C_in - это H (количество голов).
+        # (B, H, Lq, Hp, Wp) -> (B, Lq, H, Hp, Wp) -> (B * Lq, H, Hp, Wp)
+        x = accumulated_attention.permute(0, 2, 1, 3, 4).contiguous()
+        x = x.view(B * Lq, H, Hp, Wp)
+
+        # 3. Применяем conv2d для извлечения локальных паттернов покрытия
+        phi_conv_out = self.phi_conv(x)  # Shape: [B * Lq, H, Hp, Wp]
+
+        # 4. Преобразуем для подачи в Linear слой.
+        # (B*Lq, H, Hp, Wp) -> (B*Lq, Hp, Wp, H) -> (B*Lq, Hp*Wp, H)
+        phi_permuted = phi_conv_out.permute(0, 2, 3, 1)
+        phi_flat = phi_permuted.contiguous().view(B * Lq, Hp * Wp, H)
+
+        # 5. Проецируем в d_model пространство и нормализуем
+        phi_features = self.phi_linear(phi_flat)  # Shape: [B * Lq, Hp*Wp, d_model]
+        phi_features = self.phi_norm(phi_features)
+
+        # 6. Возвращаем обратно в форму, совместимую с forward pass.
+        # Разделяем B и Lq обратно.
+        # (B*Lq, Hp*Wp, d_model) -> (B, Lq, Hp*Wp, d_model)
+        phi_out = phi_features.view(B, Lq, Hp * Wp, self.d_model)
+
+        return phi_out
+
+    def forward(self, q, k, v, current_symbols_ids=None, key_padding_mask=None):
+        """
+        Forward pass с применением IAC коррекции (батчево для всей последовательности)
+        Args:
+            q: [B, Lq, d_model]
+            k: [B, Lk, d_model]
+            v: [B, Lv, d_model]
+            current_symbols: [B, Lq] (для обучения) или List[str] (для инференса)
+            key_padding_mask: [B, Lk]
+        """
+        B, Lq, _ = q.size()
+        _, Lk, _ = k.size()
+        D = self.head_dim
+        spatial_size = int(math.sqrt(Lk))
+        assert spatial_size * spatial_size == Lk, f"Feature map должна быть квадратной, получено {Lk}"
+
+        Q = self.wq(q).view(B, Lq, self.num_heads, D).transpose(1, 2)  # [B, H, Lq, D]
+        K = self.wk(k).view(B, Lk, self.num_heads, D).transpose(1, 2)  # [B, H, Lk, D]
+        V = self.wv(v).view(B, Lk, self.num_heads, D).transpose(1, 2)  # [B, H, Lk, D]
+
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(D)  # [B, H, Lq, Lk]
+        if key_padding_mask is not None:
+            mask = key_padding_mask.unsqueeze(1).unsqueeze(2)
+            scores = scores.masked_fill(mask, float('-inf'))
+        attn_weights = F.softmax(scores, dim=-1)  # [B, H, Lq, Lk]
+
+        # Батчевая коррекция внимания для всего Lq (только при обучении)
+        if current_symbols_ids is not None and self.accumulated_attention is not None:
+            # Проверяем совместимость размеров
+            accum_B, accum_H, accum_Lq, accum_Hp, accum_Wp = self.accumulated_attention.shape
+
+            # Если размеры не совпадают, расширяем/обрезаем накопленное внимание
+            if accum_Lq < Lq:
+                # Дополняем нулями
+                padding = torch.zeros(
+                    accum_B, accum_H, Lq - accum_Lq, accum_Hp, accum_Wp,
+                    device=self.accumulated_attention.device
+                )
+                accumulated_attention = torch.cat([self.accumulated_attention, padding], dim=2)
+            elif accum_Lq > Lq:
+                # Берем только первые Lq элементов
+                accumulated_attention = self.accumulated_attention[:, :, :Lq, :, :]
+            else:
+                accumulated_attention = self.accumulated_attention
+
+            phi_correction = self._compute_phi(accumulated_attention)
+            # Преобразуем phi_correction к attention space
+            phi_attn = phi_correction.view(B, Lq, spatial_size * spatial_size, self.d_model)
+            phi_attn = phi_attn.view(B, Lq, spatial_size * spatial_size, self.num_heads, self.head_dim).permute(0, 3, 1, 4, 2)  # [B, H, Lq, D, Lk]
+            # Q: [B, H, Lq, D], phi_attn: [B, H, Lq, D, Lk]
+            phi_scores = torch.matmul(Q.unsqueeze(-2), phi_attn).squeeze(-2)  # [B, H, Lq, Lk]
+            corrected_scores = scores - phi_scores
+            attn_weights = F.softmax(corrected_scores, dim=-1)
+
+        # Обновляем накопленное внимание ПОСЛЕ коррекции
+        if current_symbols_ids is not None:
+            self._update_accumulated_attention(attn_weights, current_symbols_ids, spatial_size)
+
+        attn_weights = self.attn_dropout(attn_weights)
+        context = torch.matmul(attn_weights, V)  # [B, H, Lq, D]
+        context = context.transpose(1, 2).reshape(B, Lq, self.d_model)
+        output = self.wo(context)
+        output = self.proj_dropout(output)
+        return output, attn_weights
+
+
+    def _update_accumulated_attention(self, attn_weights, current_symbols_ids, spatial_size):
+        """
+        Векторизованная версия без медленных циклов Python.
+        """
+        B, H, Lq, Lk = attn_weights.shape
+        
+        # Создаем маску за одну быструю операцию на GPU
+        # Сравниваем каждый ID в батче со списком ID структурных символов
+        is_structure_mask = (current_symbols_ids.unsqueeze(-1) == self.structure_ids.view(1, 1, -1)).any(dim=-1)
+        
+        # Инвертируем маску (1.0 для обычных символов, 0.0 для структурных)
+        # и приводим к нужному виду для умножения
+        indicator = (~is_structure_mask).float().view(B, 1, Lq, 1)
+
+        # Дальнейшая логика остается той же, но теперь она работает на GPU без тормозов
+        masked_attn = attn_weights * indicator
+        
+        if Lq > 1:
+            masked_attn_shifted = torch.zeros_like(masked_attn)
+            masked_attn_shifted[:, :, 1:] = masked_attn[:, :, :-1]
+            accumulated = torch.cumsum(masked_attn_shifted, dim=2)
+        else:
+            accumulated = torch.zeros_like(masked_attn)
+            
+        self.accumulated_attention = accumulated.view(B, H, Lq, spatial_size, spatial_size)
+
+
+class EnhancedDecoderLayer(nn.Module):
+    """
+    Decoder layer с поддержкой IAC - исправленная версия
+    """
+    def __init__(self, d_model=256, num_heads=8, d_ff=1024, dropout=DROPOUT_RATE, structure_symbols_set=None,vocab=None):
+        super().__init__()
+        self.d_model = d_model
+
+        # Self-attention
+        self.norm1 = nn.LayerNorm(d_model)
+        self.self_attn = nn.MultiheadAttention(d_model, num_heads, dropout=dropout, batch_first=True)
+
+        # Cross-attention with IAC
+        self.norm2 = nn.LayerNorm(d_model)
+        self.iac = IAC(d_model, num_heads, dropout, structure_symbols_set,vocab)
+
+        # Feed-forward
+        self.norm3 = nn.LayerNorm(d_model)
+        self.ffn = nn.Sequential(
+            nn.Linear(d_model, d_ff),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(d_ff, d_model),
+            nn.Dropout(dropout)
+        )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_iac_state(self):
+        """Сброс состояния IAC"""
+        self.iac.reset_state()
+
+    def forward(self, x, encoder_output, current_symbols_ids=None,
+                tgt_mask=None, tgt_key_padding_mask=None, memory_key_padding_mask=None):
+        """
+        Forward pass decoder layer
+
+        Args:
+            x: [B, Lq, d_model] - target embeddings
+            encoder_output: [B, Lk, d_model] - encoder output (visual features)
+            current_symbols: str или List[str] - текущие символы для IAC
+            tgt_mask: causal mask для self-attention
+            tgt_key_padding_mask: padding mask для target
+            memory_key_padding_mask: padding mask для encoder output
+        """
+        batch_size = x.size(0)
+
+        # Self-attention
+        residual = x
+        x = self.norm1(x)
+        self_attn_output, self_attn_weights = self.self_attn(
+            x, x, x,
+            attn_mask=tgt_mask,
+            key_padding_mask=tgt_key_padding_mask
+        )
+        x = residual + self.dropout(self_attn_output)
+
+        # Cross-attention с IAC
+        residual = x
+        x = self.norm2(x)
+
+        # IAC применяется только к cross-attention
+        cross_attn_output, cross_attn_weights = self.iac(
+            x, encoder_output, encoder_output,
+             current_symbols_ids=current_symbols_ids,
+            key_padding_mask=memory_key_padding_mask
+        )
+        x = residual + self.dropout(cross_attn_output)
+
+        # Feed-forward
+        residual = x
+        x = self.norm3(x)
+        ffn_output = self.ffn(x)
+        x = residual + self.dropout(ffn_output)
+
+        return x
+
+class PosFormerImprovedWithIAC(nn.Module):
+    """
+    PosFormer с исправленной поддержкой IAC и правильной реализацией согласно статье
+    """
+    def __init__(
+        self,
+        vocab_size,
+        vocab,
+        pad_token_id,
+        structure_symbols_set,
+        id_to_token_map,
+        d_model=256,
+        num_heads=8,
+        num_layers=NUM_LAYERS,
+        d_ff=1024,
+        dropout=0.1,
+        max_seq_len=MAX_SEQ_LEN,
+        max_nested_levels=NUM_NESTED_LEVELS,
+        identifier_vocab_size=None,
+        identifier_max_len=10,
+    ):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.vocab = vocab
+        self.pad_id = pad_token_id
+        self.structure_symbols = structure_symbols_set
+        self.id_to_token = id_to_token_map
+        self.max_seq_len = max_seq_len
+        self.max_nested_levels = max_nested_levels
+        self.identifier_max_len = identifier_max_len
+        self.d_model = d_model
+
+        # Encoder
+        self.enc = ImprovedImageEncoder(d_model)
+
+        # Token embeddings + pos
+        self.emb_tok = nn.Embedding(vocab_size, d_model)
+        self.pos_enc = nn.Parameter(torch.randn(1, max_seq_len, d_model) * 0.02)
+        self.norm_in = nn.LayerNorm(d_model)
+
+        # Identifier embeddings (ξ function)
+        if identifier_vocab_size is not None:
+            self.identifier_emb = nn.Embedding(identifier_vocab_size, d_model)
+            self.xi_function = nn.Sequential(
+                nn.Linear(d_model, d_model),
+                nn.GELU(),
+                nn.LayerNorm(d_model)
+            )
+        else:
+            self.identifier_emb = None
+            self.xi_function = None
+        self.identifier_pos_enc = nn.Parameter(torch.randn(1, max_seq_len, d_model) * 0.02)
+
+        # Decoder layers с IAC
+        self.decoders = nn.ModuleList([
+            EnhancedDecoderLayer(d_model, num_heads, d_ff, dropout, structure_symbols_set,vocab=vocab)
+            for _ in range(num_layers)
+        ])
+        self.norm_out = nn.LayerNorm(d_model)
+
+        # Prediction heads
+        self.W_n = nn.Linear(d_model, max_nested_levels + 1)
+        self.W_r = nn.Linear(d_model, 3)  # M, L, R
+        self.head_tok = nn.Linear(d_model, vocab_size)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """Правильная инициализация весов"""
+        # Embedding layers
+        nn.init.normal_(self.emb_tok.weight, std=0.02)
+        if self.identifier_emb is not None:
+            nn.init.normal_(self.identifier_emb.weight, std=0.02)
+
+        # Position encodings
+        nn.init.normal_(self.pos_enc, std=0.02)
+        nn.init.normal_(self.identifier_pos_enc, std=0.02)
+
+        # Prediction heads
+        nn.init.xavier_uniform_(self.W_n.weight)
+        nn.init.constant_(self.W_n.bias, 0.0)
+        nn.init.xavier_uniform_(self.W_r.weight)
+        nn.init.constant_(self.W_r.bias, 0.0)
+        nn.init.xavier_uniform_(self.head_tok.weight)
+        nn.init.constant_(self.head_tok.bias, 0.0)
+
+    def reset_iac_state(self):
+        """Сброс состояния IAC для всех decoder layers"""
+        for decoder in self.decoders:
+            decoder.reset_iac_state()
+
+    def process_identifiers(self, identifiers):
+        """
+        Обработка identifier embeddings согласно формуле (1)
+        Q_emb = [ξ(Q_1); ξ(Q_2); ...; ξ(Q_L)] + Q_pos
+        """
+        if self.identifier_emb is None:
+            return None
+
+        B, T, U = identifiers.size()
+
+        # Embedding lookup
+        emb = self.identifier_emb(identifiers)  # [B, T, U, d_model]
+
+        # Маска для padding
+        pad_mask = identifiers == 0
+        mask = ~pad_mask.unsqueeze(-1)  # [B, T, U, 1]
+
+        # Применяем маску и усредняем
+        emb = emb * mask.float()
+        lengths = mask.sum(dim=2).float().clamp(min=1)  # [B, T, 1]
+        avg_emb = emb.sum(dim=2) / lengths  # [B, T, d_model]
+
+        # Применяем ξ function
+        processed = self.xi_function(avg_emb)
+
+        # Добавляем позиционное кодирование
+        return processed + self.identifier_pos_enc[:, :T, :]
+
+    def get_symbol_from_id(self, tid):
+        """Получение символа по ID"""
+        if isinstance(tid, torch.Tensor):
+            tid = tid.item()
+        return self.id_to_token.get(tid, f"<unk_{tid}>")
+
+    def forward(self, imgs, feat=None, mode='train'):
+        """Основная функция forward"""
+        if mode == 'train':
+            return self._forward_train(imgs, feat)
+        else:
+            return self._forward_inference(imgs)
+
+    def _forward_train(self, imgs, decoder_input):
+        """
+        Принимает сдвинутый decoder_input и возвращает логиты для всех голов.
+        """
+        B, T = decoder_input.size()
+
+        self.reset_iac_state()
+        enc = self.enc(imgs)
+
+        x = self.emb_tok(decoder_input)
+        x = x + self.pos_enc[:, :T, :]
+        x = self.norm_in(x)
+
+        tgt_mask = torch.triu(torch.ones(T, T, device=imgs.device), diagonal=1).bool()
+        pad_mask = decoder_input == self.pad_id
+
+
+        for decoder in self.decoders:
+            x = decoder(
+                x, enc,
+                current_symbols_ids=decoder_input, # Передаем ID напрямую
+                tgt_mask=tgt_mask,
+                tgt_key_padding_mask=pad_mask
+            )
+
+        x = self.norm_out(x)
+
+        # Получаем логиты от всех трех "голов"
+        nested_logits = self.W_n(x)
+        relpos_logits = self.W_r(x)
+        token_logits = self.head_tok(x)
+
+        return {
+            'token_ids': token_logits,
+            'nested': nested_logits,
+            'relpos': relpos_logits,
+            'features': x
+        }
+
+    def forward_inference(self, enc_output, tokens): # Убираем current_symbols
+            """
+            Инференс с заданными энкодером и токенами.
+            Теперь передает ID токенов напрямую в декодер.
+            """
+            B, T = tokens.size()
+        
+            # Эмбеддинги и позиционное кодирование
+            x = self.emb_tok(tokens) + self.pos_enc[:, :T, :]
+            x = self.norm_in(x)
+        
+            # Прямой проход через декодер
+            for decoder in self.decoders:
+                x = decoder( # Возвращаемые веса нам здесь не нужны
+                    x,
+                    enc_output,
+                    current_symbols_ids=tokens, # Передаем тензор с ID
+                    tgt_mask=None # При инференсе каузальная маска не нужна
+                )
+        
+            x = self.norm_out(x)
+            token_logits = self.head_tok(x)
+        
+            return {'logits': token_logits}
+
+# ===== Image Encoder =====
+class ImprovedImageEncoder(nn.Module):
+    def __init__(self, d_model=256):
+        super().__init__()
+        # Загрузка предобученной DenseNet-121
+        base = models.densenet121(pretrained=True)
+
+        # Используем все слои до последнего пулинга
+        self.features = base.features
+
+        # Проекция в d_model
+        self.spatial_proj = nn.Conv2d(1024, d_model, 1)
+
+        # Адаптивный пулинг для фиксированного размера
+        self.adaptive_pool = nn.AdaptiveAvgPool2d((7, 7))  # Фиксируем размер 7x7
+
+        # Позиционные эмбеддинги
+        self.pos_embed = nn.Parameter(torch.randn(1, 49, d_model) * 0.1)
+        self.norm = nn.LayerNorm(d_model)
+
+        # Dropout для регуляризации
+        self.dropout = nn.Dropout(0.1)
+
+    def forward(self, x):
+        # Прямой проход через DenseNet
+        f = self.features(x)  # [B, 1024, H, W]
+
+        # Проекция в пространство d_model
+        s = self.spatial_proj(f)  # [B, d_model, H, W]
+
+        # Адаптивный пулинг для фиксированного размера
+        s = self.adaptive_pool(s)  # [B, d_model, 7, 7]
+
+        # Преобразование: [B, C, H, W] -> [B, H*W, C]
+        B, C, H, W = s.shape
+        s = s.flatten(2).transpose(1, 2)  # [B, 49, d_model]
+
+        # Добавление позиционных эмбеддингов
+        s = s + self.pos_embed
+
+        # Нормализация и dropout
+        s = self.norm(s)
+        s = self.dropout(s)
+
+        return s
+def generate_beam_search(model, image_tensor, beam_size, max_len, sos_id, eos_id, pad_id, device):
+    """
+    Генерирует последовательность токенов с использованием Beam Search.
+    """
+    model.eval()
+
+    model.reset_iac_state()
+
+    # 2. Получаем признаки из энкодера один раз
+    with torch.no_grad():
+        # enc_output всегда имеет batch_size=1
+        enc_output = model.enc(image_tensor.unsqueeze(0))
+
+    # 3. Инициализация лучей.
+    # beams - это список из кортежей (последовательность_тензор, score)
+    beams = [(torch.tensor([sos_id], dtype=torch.long, device=device), 0.0)]
+
+    # 4. Пошаговая генерация
+    for t in range(1, max_len):
+        all_candidates = []
+        
+        for seq, score in beams:
+            if seq[-1] == eos_id:
+                all_candidates.append((seq, score))
+                continue
+            
+            with torch.no_grad():
+                input_seq = seq.unsqueeze(0)
+                
+
+                outputs = model.forward_inference(enc_output, input_seq)
+                
+                logits = outputs['logits'][:, -1, :] # Логиты для последнего токена
+            
+            log_probs = F.log_softmax(logits, dim=-1)
+            top_log_probs, top_ids = torch.topk(log_probs, beam_size, dim=-1)
+
+            # Создаем новых кандидатов
+            for i in range(beam_size):
+                next_id = top_ids[0, i]
+                log_prob = top_log_probs[0, i].item()
+                
+                # Создаем новую последовательность, добавляя новый токен
+                new_seq = torch.cat([seq, next_id.view(1)])
+                new_score = score + log_prob
+                all_candidates.append((new_seq, new_score))
+
+        # 5. Сортируем всех кандидатов и выбираем `beam_size` лучших
+        ordered = sorted(all_candidates, key=lambda x: x[1], reverse=True)
+        beams = ordered[:beam_size]
+        
+        # 6. Условие остановки: если все лучшие лучи закончились на EOS
+        if all(b[0][-1] == eos_id for b in beams):
+            break
+            
+    # 7. Выбираем лучший луч, нормализуя на длину, чтобы не штрафовать длинные последовательности
+    best_beam = sorted(beams, key=lambda x: x[1] / len(x[0]), reverse=True)[0]
+    best_seq = best_beam[0]
+    
+    return best_seq
+
+# ======================================================================
+# 2. Удобная обертка для предсказания
+# ======================================================================
+
+def predict_with_beam_search(model, image_path, vocab, beam_size, device):
+    """
+    Загружает изображение, запускает Beam Search и возвращает LaTeX строку.
+    """
+    # 1. Загрузка и трансформация изображения
+    transform = transforms.Compose([
+        transforms.Resize((IMG_HEIGHT, IMG_WIDTH)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ])
+    image = Image.open(image_path).convert("RGB")
+    image = ImageOps.invert(image)
+    image_tensor = transform(image).to(device)
+
+    # 2. Получение ID спецтокенов
+    sos_id = vocab.stoi[SOS_TOKEN]
+    eos_id = vocab.stoi[EOS_TOKEN]
+    pad_id = vocab.stoi[PAD_TOKEN]
+
+    # 3. Генерация
+    predicted_ids = generate_beam_search(
+        model, image_tensor, beam_size, MAX_SEQ_LEN, sos_id, eos_id, pad_id, device
+    )
+
+    # 4. Декодирование
+    predicted_ids = predicted_ids.cpu().numpy()
+    try:
+        eos_index = list(predicted_ids).index(eos_id)
+        predicted_ids = predicted_ids[:eos_index]
+    except ValueError:
+        pass # EOS не был сгенерирован
+
+    latex_string = vocab.decode(predicted_ids)
+    return latex_string