model.py

import os
import re
import math
import random
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
from tqdm import tqdm
import torchvision.transforms.functional as TF
MAX_SEQ_LEN = 512
NUM_NESTED_LEVELS = 10
MAX_IDENTIFIER_LEN = 15
PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN = '<PAD>', '<SOS>', '<EOS>', '<UNK>'

D_MODEL = 256
D_FF = 1024
NUM_HEADS = 8
GROWTH_RATE = 24
NUM_ENCODER_LAYERS = 16
NUM_DECODER_LAYERS = 3
BOTTLENECK = True
DROPOUT_RATE = 0.3


TOKEN_REGEX = re.compile(
    r"(\\[a-zA-Z]+(?:\*)?)|(\\\{|\\\}|\\.)|([0-9]+(?:\.[0-9]+)?(?:[eE][+-]?[0-9]+)?)|([A-Za-z]+)|([+\-*/=<>!~^_&|%])|([{}()\[\],.;:?'])|(\s+)|(\S)"
)

def tokenize_latex(s: str) -> list[str]:
    tokens_grouped = TOKEN_REGEX.findall(s)
    tokens = []
    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

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 __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 cache_key in self._cache:
            return self._cache[cache_key]

        T = len(tokens)
        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]

                # Обработка команд с одним аргументом: ^, _, \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

                # Обработка \left, \right (они не добавляют уровень, но влияют на разметку)
                # В данной реализации мы их просто пропускаем, как и другие группирующие символы
                # Более сложная логика могла бы их учитывать, но это выходит за рамки статьи

                i += 1

        parse_recursive(0, T, 'M')

        # Заменяем префиксы обратно на полные строки
        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

        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

        # Если картинка уже в нужных рамках, ничего не делаем
        return img
        
class ScaleAugmentation:
    def __init__(self, lo: float, hi: float) -> None:
        assert lo <= hi
        self.lo = lo
        self.hi = hi

    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

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

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

        # --- Инициализируем наши классы трансформаций ---
        H_MIN, H_MAX = 32, 256
        W_MIN, W_MAX = 32, 512
        
        # Эти классы будут вызываться вручную
        self.scale_augmenter = None
        if self.is_train:
            # Эта трансформация работает с NUMPY
            self.scale_augmenter = ScaleAugmentation(0.8, 1.2)
        
        # Эта трансформация тоже работает с 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))
            
        # --- НАШ РУЧНОЙ ПАЙПЛАЙН ---

        # Шаг А: Применяем ScaleAugmentation (numpy -> numpy)
        if self.scale_augmenter:
            img_np = self.scale_augmenter(img_np)
        
        # Шаг Б: Применяем ScaleToLimitRange (numpy -> numpy)
        # Он применяется всегда, и для train, и для val
        img_np = self.size_controller(img_np)

        # Шаг В: Конвертируем в PIL только в самом конце, чтобы отдать в collate_fn
        final_img_pil = Image.fromarray(img_np)
        
        # --- Конец пайплайна ---

        # Остальной код без изменений
        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)

    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)

        return rearrange(dropped_feature, "b h w d -> b (h w) d"), rearrange(mask, "b h w -> b (h w)")

# =================================================================================
#                                  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 IAC(nn.Module):
    def __init__(self, d_model, num_heads, dropout, vocab):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        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)

        # ИСПРАВЛЕНО: Расширенный список структурных токенов
        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)

class EnhancedDecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, d_ff, dropout, cross_attention_module):
        super().__init__()
        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)

    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

# =================================================================================
#                                  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.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

        # УЛУЧШЕНИЕ: Добавляем Label Smoothing в модель
        self.loss_fn_rec = LabelSmoothingCrossEntropy(smoothing=0.05)

        self._init_weights()
        self.rel_pos_map = {0: 'M', 1: 'L', 2: 'R'}

    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]