Update model.py

model.py

@@ -1,798 +1,708 @@
 import os
 import re
 import math
-import pickle
 import random
-from collections import defaultdict, namedtuple
-
+from collections import defaultdict
+from einops import rearrange
 import numpy as np
+import cv2
 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 torchvision import transforms
 from tqdm import tqdm
+import torchvision.transforms.functional as TF
 
-# ================ 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
+    r"(\\[a-zA-Z]+(?:\*)?)|(\\\{|\\\}|\\.)|([0-9]+(?:\.[0-9]+)?(?:[eE][+-]?[0-9]+)?)|([A-Za-z]+)|([+\-*/=<>!~^_&|%])|([{}()\[\],.;:?'])|(\s+)|(\S)"
 )
 
-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)  # Удаление переносов
-
+def tokenize_latex(s: str) -> list[str]:
+    tokens_grouped = TOKEN_REGEX.findall(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)
-
+    for group in tokens_grouped:
+        non_empty_token = next(filter(None, group), '')
+        if not non_empty_token.isspace():
+             tokens.append(non_empty_token)
     return tokens
 
-# Position Forest Implementation
-STRUCTURE_CMDS = {'^', '_', '\\sqrt', '\\frac', '\\sum', '\\int', '\\lim'}
-
+class Vocab:
+    def __init__(self, expressions=None, min_freq=1):
+        self.itos = [PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN]
+        if expressions:
+            freq = defaultdict(int)
+            for expr in expressions:
+                for t in tokenize_latex(expr):
+                    freq[t] += 1
+            self.itos.extend([tok for tok, count in sorted(freq.items(), key=lambda x: -x[1]) if count >= min_freq])
+        self.stoi = {tok: i for i, tok in enumerate(self.itos)}
+        self.pad_id = self.stoi[PAD_TOKEN]
+        self.sos_id = self.stoi[SOS_TOKEN]
+        self.eos_id = self.stoi[EOS_TOKEN]
+        self.unk_id = self.stoi[UNK_TOKEN]
+    def encode(self, tokens: list[str]) -> list[int]:
+        return [self.stoi.get(t, self.unk_id) for t in tokens]
+    def decode(self, ids: list[int]) -> str:
+        tokens = [self.itos[i] for i in ids if i not in {self.pad_id, self.sos_id, self.eos_id}]
+        return "".join(tokens)
+
+class PosVocab:
+    def __init__(self):
+        self.itos = ['<PAD>', '<SOS>', '<EOS>', 'M', 'L', 'R']
+        self.stoi = {s: i for i, s in enumerate(self.itos)}
+        self.pad_id = self.stoi['<PAD>']
+        self.sos_id = self.stoi['<SOS>']
+
+# =================================================================================
+#                                  2. POSITION FOREST (ИСПРАВЛЕНО)
+# =================================================================================
 class PositionForestEncoder:
-    def position_forest_ids(self, tokens):
+    def __init__(self, pos_vocab, max_len=MAX_SEQ_LEN, max_identifier_len=MAX_IDENTIFIER_LEN):
+        self.pos_vocab = pos_vocab
+        self.max_len = max_len
+        self.max_identifier_len = max_identifier_len
+        self._cache = {}
+
+    def _get_pos_strings(self, tokens):
+        """
+        ИСПРАВЛЕНО: Полностью рекурсивный парсер, который точнее следует логике
+        построения дерева позиций для вложенных структур.
+        """
         cache_key = tuple(tokens)
-        if hasattr(self, '_cache'):
-            if cache_key in self._cache:
-                return self._cache[cache_key]
-        else:
-            self._cache = {}
+        if cache_key in self._cache:
+            return self._cache[cache_key]
+
         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:
+        ids = ['M'] * T
+
+        def find_matching_brace(start_index):
+            depth = 1
+            for i in range(start_index + 1, T):
+                if tokens[i] == '{':
+                    depth += 1
+                elif tokens[i] == '}':
+                    depth -= 1
+                    if depth == 0:
+                        return i
+            return T - 1
+
+        def parse_recursive(start, end, current_pos_prefix):
+            i = start
+            while i < end:
                 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)
+                # Обработка команд с одним аргументом: ^, _, \sqrt
+                if tok in ('^', '_', '\\sqrt') and i + 1 < end:
+                    label = 'L' if tok in ('^', '\\sqrt') else 'R'
+                    if tokens[i+1] == '{':
+                        brace_end = find_matching_brace(i + 1)
+                        for k in range(i + 2, brace_end):
+                            ids[k] = current_pos_prefix + label
+                        parse_recursive(i + 2, brace_end, current_pos_prefix + label)
+                        i = brace_end
+                    else: # Аргумент - один токен
+                        ids[i+1] = current_pos_prefix + label
+                        i += 1
+
+                # Обработка команд с двумя аргументами: \frac
+                elif tok == '\\frac' and i + 1 < end and tokens[i+1] == '{':
+                    num_end = find_matching_brace(i + 1)
+                    if num_end + 1 < end and tokens[num_end + 1] == '{':
+                        den_end = find_matching_brace(num_end + 1)
+                        # Числитель
+                        for k in range(i + 2, num_end):
+                            ids[k] = current_pos_prefix + 'L'
+                        parse_recursive(i + 2, num_end, current_pos_prefix + 'L')
+                        # Знаменатель
+                        for k in range(num_end + 2, den_end):
+                            ids[k] = current_pos_prefix + 'R'
+                        parse_recursive(num_end + 2, den_end, current_pos_prefix + 'R')
+                        i = den_end
+                    else:
+                        i = num_end
 
-    def reset_state(self):
-        """Сб��ос состояния IAC - важно для начала новой последовательности"""
-        self.accumulated_attention = None
+                # Обработка \left, \right (они не добавляют уровень, но влияют на разметку)
+                # В данной реализации мы их просто пропускаем, как и другие группирующие символы
+                # Более сложная логика могла бы их учитывать, но это выходит за рамки статьи
 
-    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
+                i += 1
 
-    def _compute_phi(self, accumulated_attention):
-        """
-        Вычисляет φ(A^k) - функцию коррекции на основе накопленного внимания.
-        ИСПРАВЛЕННАЯ ВЕРСИЯ для обработки 5D тензора.
+        parse_recursive(0, T, 'M')
 
-        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)
+        # Заменяем префиксы обратно на полные строки
+        final_ids = []
+        for i in range(T):
+            if ids[i] == 'M':
+                final_ids.append('M')
+            else:
+                final_ids.append(ids[i])
+
+        self._cache[cache_key] = final_ids
+        return final_ids
+
+    def process_formula(self, latex_str: str):
+        # Эта часть остается без изменений
+        tokens = tokenize_latex(latex_str)
+        tokens_gt = [SOS_TOKEN] + tokens[:self.max_len - 2] + [EOS_TOKEN]
+        pos_strings_raw = self._get_pos_strings(tokens[:self.max_len - 2])
+        nested_depth_gt = [0] + [min(len(pid) - 1, NUM_NESTED_LEVELS) for pid in pos_strings_raw] + [0]
+        rel_pos_gt = [0] + [1 if pid.endswith('L') else 2 if pid.endswith('R') else 0 for pid in pos_strings_raw] + [0]
+        pos_identifiers_ids = []
+        sos_pos_ids = [self.pos_vocab.stoi[c] for c in ['<SOS>', 'M', '<EOS>']]
+        pos_identifiers_ids.append(sos_pos_ids)
+        for pos_str in pos_strings_raw:
+            ids = [self.pos_vocab.stoi.get(c, 0) for c in pos_str]
+            ids = [self.pos_vocab.sos_id] + ids + [self.pos_vocab.stoi['<EOS>']]
+            pos_identifiers_ids.append(ids)
+        pos_identifiers_ids.append(sos_pos_ids)
+        final_len = len(tokens_gt)
+        tokens_gt_padded = tokens_gt + [PAD_TOKEN] * (self.max_len - final_len)
+        nested_depth_gt_padded = nested_depth_gt + [0] * (self.max_len - final_len)
+        rel_pos_gt_padded = rel_pos_gt + [0] * (self.max_len - final_len)
+        for i in range(len(pos_identifiers_ids)):
+            seq = pos_identifiers_ids[i][:self.max_identifier_len]
+            pos_identifiers_ids[i] = seq + [self.pos_vocab.pad_id] * (self.max_identifier_len - len(seq))
+        empty_pos_id_seq = [self.pos_vocab.pad_id] * self.max_identifier_len
+        padded_pos_ids = pos_identifiers_ids + [empty_pos_id_seq] * (self.max_len - final_len)
+        return {"tokens_gt": tokens_gt_padded, "pos_matrix": torch.tensor(padded_pos_ids, dtype=torch.long),
+                "nested_gt": torch.tensor(nested_depth_gt_padded, dtype=torch.long), "rel_pos_gt": torch.tensor(rel_pos_gt_padded, dtype=torch.long)}
+
+# =================================================================================
+#                                  3. DATASET & PREPROCESSING (без изменений)
+# =================================================================================
+class ResizeWithPadding:
+    def __init__(self, target_height, max_target_width, padding_value=255):
+        self.target_height, self.max_target_width, self.padding_value = target_height, max_target_width, padding_value
+    def __call__(self, img):
+        w, h = img.size; new_w = int(w * (self.target_height / h)); new_w = min(new_w, self.max_target_width)
+        img_resized = img.resize((new_w, self.target_height), Image.LANCZOS)
+        new_img = Image.new(img.mode, (self.max_target_width, self.target_height), self.padding_value)
+        new_img.paste(img_resized, (0, 0))
+        mask = torch.ones((self.target_height, self.max_target_width), dtype=torch.bool)
+        mask[:, :new_w] = False
+        return new_img, mask
+
+class ScaleToLimitRange:
+    def __init__(self, w_lo: int, w_hi: int, h_lo: int, h_hi: int) -> None:
+        assert w_lo <= w_hi and h_lo <= h_hi
+        self.w_lo = w_lo
+        self.w_hi = w_hi
+        self.h_lo = h_lo
+        self.h_hi = h_hi
+
+    def __call__(self, img: np.ndarray) -> np.ndarray:
+        h, w = img.shape[:2]
+        scale_r = min(self.h_hi / h, self.w_hi / w)
+        if scale_r < 1.0:
+            # Картинка слишком большая, сжимаем ее пропорционально
+            img = cv2.resize(
+                img, None, fx=scale_r, fy=scale_r, interpolation=cv2.INTER_LINEAR
+            )
+            return img
 
-        # 3. Применяем conv2d для извлечения локальных паттернов покрытия
-        phi_conv_out = self.phi_conv(x)  # Shape: [B * Lq, H, Hp, Wp]
+        scale_r = max(self.h_lo / h, self.w_lo / w)
+        if scale_r > 1.0:
+            # Картинка слишком маленькая, увеличиваем ее пропорционально
+            img = cv2.resize(
+                img, None, fx=scale_r, fy=scale_r, interpolation=cv2.INTER_LINEAR
+            )
+            return img
 
-        # 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)
+        # Если картинка уже в нужных рамках, ничего не делаем
+        return img
+        
+class ScaleAugmentation:
+    def __init__(self, lo: float, hi: float) -> None:
+        assert lo <= hi
+        self.lo = lo
+        self.hi = hi
 
-        # 5. Проецируем в d_model пространство и нормализуем
-        phi_features = self.phi_linear(phi_flat)  # Shape: [B * Lq, Hp*Wp, d_model]
-        phi_features = self.phi_norm(phi_features)
+    def __call__(self, img: np.ndarray) -> np.ndarray:
+        k = np.random.uniform(self.lo, self.hi)
+        img = cv2.resize(img, None, fx=k, fy=k, interpolation=cv2.INTER_LINEAR)
+        return img
 
-        # 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)
+class CROHMEDataset(Dataset):
+    def __init__(self, base_dir, caption_file, vocab, pos_vocab, is_train=True):
+        self.vocab = vocab
+        self.pfe = PositionForestEncoder(pos_vocab)
+        self.items = []
+        self.is_train = is_train
 
-        return phi_out
+        img_dir = os.path.join(base_dir, 'img')
+        formulas_path = os.path.join(base_dir, caption_file)
+        with open(formulas_path, 'r', encoding='utf-8') as f:
+            for line in f:
+                fid, latex = line.strip().split('\t')
+                self.items.append((os.path.join(img_dir, f"{fid}.bmp"), latex))
 
-    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
+        # --- Инициализируем наши классы трансформаций ---
+        H_MIN, H_MAX = 32, 256
+        W_MIN, W_MAX = 32, 512
         
-        # Создаем маску за одну быструю операцию на GPU
-        # Сравниваем каждый ID в батче со списком ID структурных символов
-        is_structure_mask = (current_symbols_ids.unsqueeze(-1) == self.structure_ids.view(1, 1, -1)).any(dim=-1)
+        # Эти классы будут вызываться вручную
+        self.scale_augmenter = None
+        if self.is_train:
+            # Эта трансформация работает с NUMPY
+            self.scale_augmenter = ScaleAugmentation(0.8, 1.2)
         
-        # Инвертируем маску (1.0 для обычных символов, 0.0 для структурных)
-        # и приводим к нужному виду для умножения
-        indicator = (~is_structure_mask).float().view(B, 1, Lq, 1)
+        # Эта трансформация тоже работает с NUMPY
+        self.size_controller = ScaleToLimitRange(h_lo=H_MIN, h_hi=H_MAX, w_lo=W_MIN, w_hi=W_MAX)
+
+    def __len__(self):
+        return len(self.items)
+
+    def __getitem__(self, idx):
+        path, latex = self.items[idx]
+        try:
+            # 1. ЧИТАЕМ КАРТИНКУ СРАЗУ В NUMPY МАССИВ. БОЛЬШЕ НИКАКИХ PIL.OPEN
+            img_np = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
+            if img_np is None: 
+                raise FileNotFoundError()
+        except Exception:
+            # Если файл битый, берем следующий
+            return self.__getitem__((idx + 1) % len(self))
+            
+        # --- НАШ РУЧНОЙ ПАЙПЛАЙН ---
 
-        # Дальнейшая логика остается той же, но теперь она работает на GPU без тормозов
-        masked_attn = attn_weights * indicator
+        # Шаг А: Применяем ScaleAugmentation (numpy -> numpy)
+        if self.scale_augmenter:
+            img_np = self.scale_augmenter(img_np)
         
-        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)
+        # Шаг Б: Применяем ScaleToLimitRange (numpy -> numpy)
+        # Он применяется всегда, и для train, и для val
+        img_np = self.size_controller(img_np)
 
+        # Шаг В: Конвертируем в PIL только в самом конце, чтобы отдать в collate_fn
+        final_img_pil = Image.fromarray(img_np)
+        
+        # --- Конец пайплайна ---
 
-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):
+        # Остальной код без изменений
+        data = self.pfe.process_formula(latex)
+        token_ids = torch.tensor(self.vocab.encode(data["tokens_gt"]), dtype=torch.long)
+        
+        return (final_img_pil, token_ids, data["pos_matrix"], data["nested_gt"], data["rel_pos_gt"], latex)
+# =================================================================================
+#                                  4. ENCODER (без изменений)
+# =================================================================================
+class ImgPosEnc(nn.Module):
+    def __init__(self, d_model: int, temperature: float = 10000.0, normalize: bool = True, scale: float = None):
         super().__init__()
+        if d_model % 4 != 0: raise ValueError(f"d_model ({d_model}) должен быть кратен 4.")
         self.d_model = d_model
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is None: scale = 2 * math.pi
+        self.scale = scale
+
+    def forward(self, x: torch.Tensor, mask: torch.BoolTensor) -> torch.Tensor:
+        not_mask = ~mask
+        y_embed = not_mask.cumsum(1, dtype=torch.float32)
+        x_embed = not_mask.cumsum(2, dtype=torch.float32)
+        
+        if self.normalize:
+            eps = 1e-6
+            y_embed = (y_embed / (y_embed[:, -1:, :] + eps)) * self.scale
+            x_embed = (x_embed / (x_embed[:, :, -1:] + eps)) * self.scale
+
+        dim_t_half = self.d_model // 2
+        dim_t = torch.arange(dim_t_half, dtype=torch.float32, device=x.device)
+        dim_t = self.temperature ** (2 * (dim_t // 2) / dim_t_half)
+        
+        pos_x = x_embed[:, :, :, None] / dim_t
+        pos_y = y_embed[:, :, :, None] / dim_t
+        
+        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
+        
+        pos = torch.cat((pos_y, pos_x), dim=3)
+        return x + pos
+
+class _Bottleneck(nn.Module):
+    def __init__(self, n_channels: int, growth_rate: int, use_dropout: bool):
+        super(_Bottleneck, self).__init__()
+        interChannels = 4 * growth_rate
+        self.conv1 = nn.Conv2d(n_channels, interChannels, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(interChannels)
+        self.conv2 = nn.Conv2d(interChannels, growth_rate, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(growth_rate)
+        self.use_dropout = use_dropout
+        self.dropout = nn.Dropout(p=0.2)
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        if self.use_dropout: out = self.dropout(out)
+        out = F.relu(self.bn2(self.conv2(out)), inplace=True)
+        if self.use_dropout: out = self.dropout(out)
+        out = torch.cat((x, out), 1)
+        return out
+class _Transition(nn.Module):
+    def __init__(self, n_channels: int, n_out_channels: int, use_dropout: bool):
+        super(_Transition, self).__init__()
+        self.conv1 = nn.Conv2d(n_channels, n_out_channels, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(n_out_channels)
+        self.use_dropout = use_dropout
+        self.dropout = nn.Dropout(p=0.2)
+    def forward(self, x):
+        out = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        if self.use_dropout: out = self.dropout(out)
+        out = F.avg_pool2d(out, 2, ceil_mode=True)
+        return out
+class DenseNet(nn.Module):
+    def __init__(self, growth_rate: int, num_layers: int, reduction: float = 0.5, bottleneck: bool = True, use_dropout: bool = True):
+        super(DenseNet, self).__init__()
+        n_dense_blocks = num_layers
+        n_channels = 2 * growth_rate
+        self.conv1 = nn.Conv2d(1, n_channels, kernel_size=7, padding=3, stride=2, bias=False)
+        self.norm1 = nn.BatchNorm2d(n_channels)
+        self.dense1 = self._make_dense(n_channels, growth_rate, n_dense_blocks, bottleneck, use_dropout)
+        n_channels += n_dense_blocks * growth_rate
+        n_out_channels = int(math.floor(n_channels * reduction))
+        self.trans1 = _Transition(n_channels, n_out_channels, use_dropout)
+        n_channels = n_out_channels
+        self.dense2 = self._make_dense(n_channels, growth_rate, n_dense_blocks, bottleneck, use_dropout)
+        n_channels += n_dense_blocks * growth_rate
+        n_out_channels = int(math.floor(n_channels * reduction))
+        self.trans2 = _Transition(n_channels, n_out_channels, use_dropout)
+        n_channels = n_out_channels
+        self.dense3 = self._make_dense(n_channels, growth_rate, n_dense_blocks, bottleneck, use_dropout)
+        self.out_channels = n_channels + n_dense_blocks * growth_rate
+        self.post_norm = nn.BatchNorm2d(self.out_channels)
+    @staticmethod
+    def _make_dense(n_channels, growth_rate, n_dense_blocks, bottleneck, use_dropout):
+        layers = []
+        for _ in range(int(n_dense_blocks)):
+            if bottleneck:
+                layers.append(_Bottleneck(n_channels, growth_rate, use_dropout))
+            n_channels += growth_rate
+        return nn.Sequential(*layers)
+    def forward(self, x, x_mask):
+        out = self.conv1(x)
+        out = self.norm1(out)
+        out_mask = x_mask[:, ::2, ::2]
+        out = F.relu(out, inplace=True)
+        out = F.max_pool2d(out, 2, ceil_mode=True)
+        out_mask = out_mask[:, ::2, ::2]
+        out = self.dense1(out)
+        out = self.trans1(out)
+        out_mask = out_mask[:, ::2, ::2]
+        out = self.dense2(out)
+        out = self.trans2(out)
+        out_mask = out_mask[:, ::2, ::2]
+        out = self.dense3(out)
+        out = self.post_norm(out)
+        return out, out_mask
+class Encoder(nn.Module):
+    def __init__(self, d_model, growth_rate, num_layers, dropout):
+        super().__init__()
+        self.densenet = DenseNet(growth_rate, num_layers)
+        self.feature_proj = nn.Conv2d(self.densenet.out_channels, d_model, 1)
+        self.pos_enc_2d = ImgPosEnc(d_model, normalize=True)
+        self.norm = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(p=dropout)
 
-        # 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)
+    def forward(self, img, img_mask):
+        feature, mask = self.densenet(img, img_mask)
+        feature = self.feature_proj(feature)
+        
+        feature_permuted = feature.permute(0, 2, 3, 1)
+        pos_encoded_feature = self.pos_enc_2d(feature_permuted, mask)
+        
+        normed_feature = self.norm(pos_encoded_feature)
+        dropped_feature = self.dropout(normed_feature)
 
-        # Feed-forward
-        residual = x
-        x = self.norm3(x)
-        ffn_output = self.ffn(x)
-        x = residual + self.dropout(ffn_output)
+        return rearrange(dropped_feature, "b h w d -> b (h w) d"), rearrange(mask, "b h w -> b (h w)")
 
-        return x
+# =================================================================================
+#                                  5. DECODER & IAC (ИСПРАВЛЕНО)
+# =================================================================================
+class StandardCrossAttention(nn.Module):
+    """Обертка для стандартного MHA, чтобы интерфейс был как у IAC."""
+    def __init__(self, d_model, num_heads, dropout):
+        super().__init__()
+        self.mha = nn.MultiheadAttention(d_model, num_heads, dropout=dropout, batch_first=True)
+    def forward(self, q, k, v, ids, pad_mask): # `ids` не используется, но нужен для совместимости
+        return self.mha(q, k, v, key_padding_mask=pad_mask)[0]
 
-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,
-    ):
+class IAC(nn.Module):
+    def __init__(self, d_model, num_heads, dropout, vocab):
         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
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
 
-        # 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)
+        self.wq, self.wk, self.wv, self.wo = [nn.Linear(d_model, d_model, bias=False) for _ in range(4)]
+        self.phi_conv = nn.Conv2d(num_heads, num_heads, kernel_size=3, padding=1, groups=num_heads)
+        self.phi_linear = nn.Linear(self.num_heads, self.head_dim)
+        self.dropout = nn.Dropout(dropout)
 
-        return {
-            'token_ids': token_logits,
-            'nested': nested_logits,
-            'relpos': relpos_logits,
-            'features': x
+        # ИСПРАВЛЕНО: Расширенный список структурных токенов
+        structure_symbols = {
+            '^', '_', '{', '}', '\\frac', '\\sqrt', '\\left', '\\right',
+            '\\big', '\\Big', '\\bigg', '\\Bigg', # Размеры скобок
+            '&', '\\\\', # Элементы матриц и таблиц
+            # Служебные токены также считаются структурными
+            PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN
         }
+        s_ids = [vocab.stoi[s] for s in structure_symbols if s in vocab.stoi]
+        self.register_buffer('structure_ids', torch.tensor(s_ids, dtype=torch.long))
+
+    def _compute_phi(self, accum_attn):
+        B, H, Lq, Lk = accum_attn.shape
+        # Проверка, что Lk является идеальным квадратом, чтобы избежать ошибок с sqrt
+        s_dim_float = math.sqrt(Lk)
+        if s_dim_float != int(s_dim_float): return 0 # Не можем сформировать квадратную карту внимания
+        s_dim = int(s_dim_float)
+
+        phi_in = rearrange(accum_attn, 'b h lq (s1 s2) -> (b lq) h s1 s2', s1=s_dim, s2=s_dim)
+        conv_out = self.phi_conv(phi_in)
+        phi_permuted = rearrange(conv_out, '(b lq) h s1 s2 -> b lq (s1 s2) h', b=B)
+        phi_features = self.phi_linear(phi_permuted)
+        return phi_features.unsqueeze(1)
+
+    def forward(self, q, k, v, ids, pad_mask):
+        B, Lq, _ = q.shape; Lk = k.shape[1]
+        Q = self.wq(q).view(B, Lq, self.num_heads, self.head_dim).transpose(1, 2)
+        K = self.wk(k).view(B, Lk, self.num_heads, self.head_dim).transpose(1, 2)
+        V = self.wv(v).view(B, Lk, self.num_heads, self.head_dim).transpose(1, 2)
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
+
+        temp_attn_weights = F.softmax(scores, dim=-1).detach()
+        indicator = (~(ids.unsqueeze(-1) == self.structure_ids).any(-1)).float().view(B, 1, Lq, 1)
+        masked_attn = temp_attn_weights * indicator
+        shifted = torch.zeros_like(masked_attn)
+        if Lq > 1: shifted[:, :, 1:, :] = masked_attn[:, :, :-1, :]
+        accumulated_attention = torch.cumsum(shifted, dim=2)
+
+        phi = self._compute_phi(accumulated_attention)
+        correction = (Q.unsqueeze(3) * phi).sum(dim=-1) if isinstance(phi, torch.Tensor) else 0
+        corrected_scores = scores - correction
+
+        if pad_mask is not None:
+            corrected_scores = corrected_scores.masked_fill(pad_mask.unsqueeze(1).unsqueeze(2), float('-inf'))
+        final_attn_weights = F.softmax(corrected_scores, dim=-1)
+        context = torch.matmul(self.dropout(final_attn_weights), V).transpose(1, 2).reshape(B, Lq, self.d_model)
+        return self.wo(context)
 
-    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):
+class EnhancedDecoderLayer(nn.Module):
+    def __init__(self, d_model, num_heads, d_ff, dropout, cross_attention_module):
         super().__init__()
-        # Загрузка предобученной DenseNet-121
-        base = models.densenet121(pretrained=True)
+        self.self_attn = nn.MultiheadAttention(d_model, num_heads, dropout=dropout, batch_first=True)
+        # ИСПРАВЛЕНО: Используем переданный модуль (Standard MHA или IAC)
+        self.cross_attn = cross_attention_module
+        self.ffn = nn.Sequential(nn.Linear(d_model, d_ff), nn.GELU(), nn.Linear(d_ff, d_model))
+        self.norm1, self.norm2, self.norm3 = [nn.LayerNorm(d_model) for _ in range(3)]
+        self.dropout = nn.Dropout(dropout)
 
-        # Используем все слои до последнего пулинга
-        self.features = base.features
+    def forward(self, x, enc_out, ids, tgt_mask, tgt_pad_mask, mem_pad_mask):
+        x = x + self.dropout(self.self_attn(self.norm1(x), self.norm1(x), self.norm1(x), attn_mask=tgt_mask, key_padding_mask=tgt_pad_mask)[0])
+        x = x + self.dropout(self.cross_attn(self.norm2(x), enc_out, enc_out, ids, mem_pad_mask))
+        x = x + self.dropout(self.ffn(self.norm3(x)))
+        return x
 
-        # Проекция в d_model
-        self.spatial_proj = nn.Conv2d(1024, d_model, 1)
+# =================================================================================
+#                                  6. FULL POSFORMER MODEL (ИСПРАВЛЕНО)
+# =================================================================================
+class LabelSmoothingCrossEntropy(nn.Module):
+    """УЛУЧШЕНИЕ: Добавлено для борьбы с переобучением."""
+    def __init__(self, smoothing=0.0):
+        super(LabelSmoothingCrossEntropy, self).__init__()
+        self.smoothing = smoothing
+    def forward(self, x, target, ignore_index=-100):
+        confidence = 1. - self.smoothing
+        logprobs = F.log_softmax(x, dim=-1)
+        nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1)).squeeze(1)
+        smooth_loss = -logprobs.mean(dim=-1)
+        loss = confidence * nll_loss + self.smoothing * smooth_loss
+        mask = (target != ignore_index)
+        return (loss * mask.float()).sum() / mask.float().sum()
+
+class PosFormer(nn.Module):
+    def __init__(self, main_vocab, pos_vocab):
+        super().__init__()
+        self.vocab, self.pos_vocab = main_vocab, pos_vocab
+        self.pad_id, self.sos_id, self.eos_id = main_vocab.pad_id, main_vocab.sos_id, main_vocab.eos_id
+        self.encoder = Encoder(D_MODEL, GROWTH_RATE, NUM_ENCODER_LAYERS, DROPOUT_RATE)
+        self.emb_tok = nn.Embedding(len(main_vocab.itos), D_MODEL)
+        self.pos_id_emb = nn.Embedding(len(pos_vocab.itos), D_MODEL)
+        self.xi_proj = nn.Sequential(nn.Linear(D_MODEL * MAX_IDENTIFIER_LEN, D_MODEL), nn.GELU(), nn.LayerNorm(D_MODEL))
+
+        # ИСПРАВЛЕНО: Создаем декодеры с разным типом внимания
+        decoder_layers = []
+        for i in range(NUM_DECODER_LAYERS):
+            use_iac = i > 0 # IAC на 2м и 3м слое (индексы 1 и 2)
+            cross_attn_module = IAC(D_MODEL, NUM_HEADS, DROPOUT_RATE, main_vocab) if use_iac \
+                else StandardCrossAttention(D_MODEL, NUM_HEADS, DROPOUT_RATE)
+            decoder_layers.append(
+                EnhancedDecoderLayer(D_MODEL, NUM_HEADS, D_FF, DROPOUT_RATE, cross_attn_module)
+            )
+        self.decoders = nn.ModuleList(decoder_layers)
 
-        # Адаптивный пулинг для фиксированного размера
-        self.adaptive_pool = nn.AdaptiveAvgPool2d((7, 7))  # Фиксируем размер 7x7
+        self.dec_norm = nn.LayerNorm(D_MODEL)
+        self.head_tok = nn.Linear(D_MODEL, len(main_vocab.itos))
+        self.head_nested = nn.Linear(D_MODEL, NUM_NESTED_LEVELS + 1)
+        self.head_rel_pos = nn.Linear(D_MODEL, 3) # 0: M, 1: L, 2: R
 
-        # Позиционные эмбеддинги
-        self.pos_embed = nn.Parameter(torch.randn(1, 49, d_model) * 0.1)
-        self.norm = nn.LayerNorm(d_model)
+        # УЛУЧШЕНИЕ: Добавляем Label Smoothing в модель
+        self.loss_fn_rec = LabelSmoothingCrossEntropy(smoothing=0.05)
 
-        # Dropout для регуляризации
-        self.dropout = nn.Dropout(0.1)
+        self._init_weights()
+        self.rel_pos_map = {0: 'M', 1: 'L', 2: 'R'}
 
-    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. Получаем признаки из энкодера один раз.
-    #    image_tensor уже должен иметь batch_size=1, поэтому убираем .unsqueeze(0)
-    with torch.no_grad():
-        enc_output = model.enc(image_tensor)
-
-    # 3. Инициализация лучей.
-    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))
-
-        ordered = sorted(all_candidates, key=lambda x: x[1], reverse=True)
-        beams = ordered[:beam_size]
-        
-        if all(b[0][-1] == eos_id for b in beams):
-            break
-            
-    best_beam = sorted(beams, key=lambda x: x[1] / len(x[0]), reverse=True)[0]
-    best_seq = best_beam[0]
-    
-    return best_seq
-
-
-def predict_with_beam_search(model, image_tensor, vocab, beam_size, device):
-    """
-    ЗАПУСКАЕТ BEAM SEARCH НА ОСНОВЕ ГОТОВОГО ТЕНЗОРА ИЗОБРАЖЕНИЯ.
-    """
-    # <<< МЫ УДАЛИЛИ ВСЮ ЛОГИКУ ЗАГРУЗКИ И ТРАНСФОРМАЦИИ ИЗОБРАЖЕНИЯ ОТСЮДА >>>
-    # Теперь функция просто принимает image_tensor
-
-    # 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.to(device), beam_size, MAX_SEQ_LEN, sos_id, eos_id, pad_id, device
-    )
-
-    # 4. Декодирование
-    predicted_ids = predicted_ids.cpu().numpy()
-    try:
-        # Убираем все после EOS токена
-        eos_index = list(predicted_ids).index(eos_id)
-        predicted_ids = predicted_ids[:eos_index]
-    except ValueError:
-        pass # EOS не был сгенерирован
-
-    # Декодируем, пропуская SOS токен в начале
-    latex_string = vocab.decode(predicted_ids)
-    return latex_string
+    def _init_weights(self):
+        for p in self.parameters():
+            if p.dim() > 1: nn.init.xavier_uniform_(p)
+
+    def forward(self, imgs, img_mask, pos_matrix, token_ids):
+        vis_feats, mem_pad_mask = self.encoder(imgs, img_mask)
+        pos_input = self.xi_proj(rearrange(self.pos_id_emb(pos_matrix[:, :-1, :]), 'b l d e -> b l (d e)'))
+        token_input_ids = token_ids[:, :-1]
+        token_embeds = self.emb_tok(token_input_ids)
+        tgt = pos_input + token_embeds
+        tgt_mask = torch.triu(torch.ones(tgt.size(1), tgt.size(1), device=tgt.device), 1).bool()
+        tgt_pad_mask = (token_input_ids == self.pad_id)
+        dec_out = self.dec_norm(self._decode(tgt, vis_feats, token_input_ids, tgt_mask, tgt_pad_mask, mem_pad_mask))
+        return self.head_tok(dec_out), self.head_nested(dec_out), self.head_rel_pos(dec_out)
+
+    def _decode(self, tgt, mem, ids, tgt_mask, tgt_pad_mask, mem_pad_mask):
+        for layer in self.decoders:
+            tgt = layer(tgt, mem, ids, tgt_mask, tgt_pad_mask, mem_pad_mask)
+        return tgt
+
+    def compute_loss(self, logits, targets):
+        token_logits, nested_logits, rel_pos_logits = logits
+        token_ids_gt, nested_gt, rel_pos_gt = targets
+        token_targets = token_ids_gt[:, 1:]
+        nested_targets = nested_gt[:, 1:]
+        rel_pos_targets = rel_pos_gt[:, 1:]
+
+        # УЛУЧШЕНИЕ: Используем Label Smoothing для рекогници-потерь
+        loss_rec = self.loss_fn_rec(
+            token_logits.reshape(-1, token_logits.size(-1)),
+            token_targets.reshape(-1),
+            ignore_index=self.pad_id
+        )
+
+        mask = (token_targets != self.pad_id).flatten()
+        if not mask.any():
+            return {'total': loss_rec, 'rec': loss_rec, 'pos': torch.tensor(0.0, device=token_logits.device)}
+
+        loss_nested = F.cross_entropy(nested_logits.reshape(-1, nested_logits.size(-1))[mask],
+                                      nested_targets.reshape(-1)[mask])
+        loss_rel = F.cross_entropy(rel_pos_logits.reshape(-1, rel_pos_logits.size(-1))[mask],
+                                   rel_pos_targets.reshape(-1)[mask])
+        loss_pos = loss_nested + loss_rel
+        total_loss = loss_rec + loss_pos
+
+        return {'total': total_loss, 'rec': loss_rec, 'pos': loss_pos}
+
+    # ... (Все методы генерации generate, generate_beam_search, _construct_next_pos_string остаются без изменений) ...
+    # Они будут работать лучше, так как модель обучается на более качественных данных
+    def _construct_next_pos_string(self, prev_pos_str: str, pred_nested_level: int, pred_rel_pos: int) -> str:
+        prev_level = len(prev_pos_str) - 1
+        new_pos_str = prev_pos_str[:pred_nested_level + 1]
+        if pred_nested_level > prev_level:
+            if len(new_pos_str) == prev_level + 1:
+                 new_pos_str += self.rel_pos_map.get(pred_rel_pos, 'M')
+        return new_pos_str
+
+    @torch.no_grad()
+    def generate(self, imgs, max_gen_len=150):
+        self.eval()
+        B = imgs.shape[0]; device = imgs.device
+        img_mask = torch.zeros_like(imgs[:, 0, :, :], dtype=torch.bool)
+        vis_feats, mem_pad_mask = self.encoder(imgs, img_mask)
+        generated_ids = torch.full((B, 1), self.sos_id, dtype=torch.long, device=device)
+        pos_strings_T = [['M'] for _ in range(B)]
+        is_finished = torch.zeros(B, dtype=torch.bool, device=device)
+        for t in range(max_gen_len - 1):
+            token_embeds = self.emb_tok(generated_ids)
+            pos_ids_list = []
+            max_len_in_batch = max(len(p_list) for p_list in pos_strings_T)
+            for i in range(B):
+                batch_pos_ids = []
+                for p_str in pos_strings_T[i]:
+                    ids = [self.pos_vocab.sos_id] + [self.pos_vocab.stoi.get(c,0) for c in p_str] + [self.pos_vocab.stoi['<EOS>']]
+                    padded_ids = ids[:MAX_IDENTIFIER_LEN] + [self.pos_vocab.pad_id] * (MAX_IDENTIFIER_LEN - len(ids))
+                    batch_pos_ids.append(torch.tensor(padded_ids, device=device))
+                while len(batch_pos_ids) < max_len_in_batch:
+                    batch_pos_ids.append(torch.full((MAX_IDENTIFIER_LEN,), self.pos_vocab.pad_id, device=device))
+                pos_ids_list.append(torch.stack(batch_pos_ids))
+            pos_matrix = torch.stack(pos_ids_list)
+            pos_embeds = self.xi_proj(rearrange(self.pos_id_emb(pos_matrix), 'b l d e -> b l (d e)'))
+            tgt = token_embeds + pos_embeds
+            tgt_mask = torch.triu(torch.ones(tgt.size(1), tgt.size(1), device=device), 1).bool()
+            dec_out = self.dec_norm(self._decode(tgt, vis_feats, generated_ids, tgt_mask, None, mem_pad_mask))
+            last_step_out = dec_out[:, -1, :]
+            token_logits, nested_logits, rel_pos_logits = self.head_tok(last_step_out), self.head_nested(last_step_out), self.head_rel_pos(last_step_out)
+            next_token_id = torch.argmax(token_logits, dim=-1).unsqueeze(1)
+            pred_nested_level, pred_rel_pos = torch.argmax(nested_logits, dim=-1), torch.argmax(rel_pos_logits, dim=-1)
+            generated_ids = torch.cat([generated_ids, next_token_id], dim=1)
+            for i in range(B):
+                if not is_finished[i]:
+                    new_pos_str = self._construct_next_pos_string(pos_strings_T[i][-1], pred_nested_level[i].item(), pred_rel_pos[i].item())
+                    pos_strings_T[i].append(new_pos_str)
+            is_finished |= (next_token_id.squeeze(-1) == self.eos_id)
+            if is_finished.all(): break
+        return generated_ids
+
+    @torch.no_grad()
+    def generate_beam_search(self, imgs, beam_size=5, max_gen_len=100):
+
+        self.eval(); B = imgs.shape[0]; device = imgs.device
+        img_mask = torch.zeros_like(imgs[:, 0, :, :], dtype=torch.bool)
+        vis_feats, mem_pad_mask = self.encoder(imgs, img_mask)
+        vis_feats = vis_feats.repeat_interleave(beam_size, dim=0)
+        if mem_pad_mask is not None: mem_pad_mask = mem_pad_mask.repeat_interleave(beam_size, dim=0)
+        effective_batch_size = B * beam_size
+        generated_ids = torch.full((effective_batch_size, 1), self.sos_id, dtype=torch.long, device=device)
+        pos_strings_T = [['M'] for _ in range(effective_batch_size)]
+        log_scores = torch.zeros(effective_batch_size, device=device)
+        is_finished = torch.zeros(effective_batch_size, dtype=torch.bool, device=device)
+        for t in range(max_gen_len - 1):
+            if is_finished.all(): break
+            token_embeds = self.emb_tok(generated_ids)
+            pos_ids_list = []
+            for i in range(effective_batch_size):
+                batch_pos_ids = []
+                for p_str in pos_strings_T[i]:
+                    ids = [self.pos_vocab.sos_id] + [self.pos_vocab.stoi.get(c, 0) for c in p_str] + [self.pos_vocab.stoi['<EOS>']]
+                    padded_ids = ids[:MAX_IDENTIFIER_LEN] + [self.pos_vocab.pad_id] * (MAX_IDENTIFIER_LEN - len(ids))
+                    batch_pos_ids.append(torch.tensor(padded_ids, device=device))
+                pos_ids_list.append(torch.stack(batch_pos_ids))
+            pos_matrix = torch.stack(pos_ids_list)
+            pos_embeds = self.xi_proj(rearrange(self.pos_id_emb(pos_matrix), 'b l d e -> b l (d e)'))
+            tgt = token_embeds + pos_embeds
+            tgt_mask = torch.triu(torch.ones(tgt.size(1), tgt.size(1), device=device), 1).bool()
+            dec_out = self.dec_norm(self._decode(tgt, vis_feats, generated_ids, tgt_mask, None, mem_pad_mask))
+            last_step_out = dec_out[:, -1, :]
+            token_logits, nested_logits, rel_pos_logits = self.head_tok(last_step_out), self.head_nested(last_step_out), self.head_rel_pos(last_step_out)
+            log_probs = F.log_softmax(token_logits, dim=-1)
+            if t > 0:
+                log_probs[is_finished] = -float('inf')
+                log_probs[is_finished, self.pad_id] = 0
+            total_scores = log_probs + log_scores.unsqueeze(1)
+            total_scores = total_scores.view(B, -1)
+            top_scores, top_indices = torch.topk(total_scores, beam_size, dim=1)
+            beam_indices = top_indices // len(self.vocab.itos)
+            token_indices = top_indices % len(self.vocab.itos)
+            batch_indices = torch.arange(B, device=device).view(-1, 1).repeat(1, beam_size)
+            beam_indices_abs = beam_indices + (batch_indices * beam_size)
+            generated_ids = generated_ids[beam_indices_abs.view(-1)]
+            pos_strings_T = [pos_strings_T[i] for i in beam_indices_abs.view(-1).tolist()]
+            generated_ids = torch.cat([generated_ids, token_indices.view(-1, 1)], dim=1)
+            pred_nested_levels = torch.argmax(nested_logits[beam_indices_abs.view(-1)], dim=-1)
+            pred_rel_poses = torch.argmax(rel_pos_logits[beam_indices_abs.view(-1)], dim=-1)
+            for i in range(effective_batch_size):
+                new_pos_str = self._construct_next_pos_string(pos_strings_T[i][-1], pred_nested_levels[i].item(), pred_rel_poses[i].item())
+                pos_strings_T[i].append(new_pos_str)
+            log_scores = top_scores.view(-1)
+            is_finished = is_finished[beam_indices_abs.view(-1)] | (token_indices.view(-1) == self.eos_id)
+        seq_lengths = (generated_ids != self.pad_id).sum(dim=1).float()
+        seq_lengths = torch.max(seq_lengths, torch.ones_like(seq_lengths))
+        normalized_scores = (log_scores / seq_lengths).view(B, beam_size)
+        best_beam_indices = torch.argmax(normalized_scores, dim=1)
+        final_indices = best_beam_indices + torch.arange(B, device=device) * beam_size
+        return generated_ids[final_indices]