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 = '', '', '', '' 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 = ['', '', '', 'M', 'L', 'R'] self.stoi = {s: i for i, s in enumerate(self.itos)} self.pad_id = self.stoi[''] self.sos_id = self.stoi[''] # ================================================================================= # 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 ['', 'M', '']] 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['']] 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['']] 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['']] 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]