Update model.py

model.py

@@ -717,18 +717,17 @@ class ImprovedImageEncoder(nn.Module):
 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. Получаем признаки из энкодера один раз
+    # 2. Получаем признаки из энкодера один раз.
+    #    image_tensor уже должен иметь batch_size=1, поэтому убираем .unsqueeze(0)
     with torch.no_grad():
-        # enc_output всегда имеет batch_size=1
-        enc_output = model.enc(image_tensor.unsqueeze(0))
+        enc_output = model.enc(image_tensor)
 
     # 3. Инициализация лучей.
-    # beams - это список из кортежей (последовательность_тензор, score)
     beams = [(torch.tensor([sos_id], dtype=torch.long, device=device), 0.0)]
 
     # 4. Пошаговая генерация
@@ -742,73 +741,58 @@ def generate_beam_search(model, image_tensor, beam_size, max_len, sos_id, eos_id
             
             with torch.no_grad():
                 input_seq = seq.unsqueeze(0)
-                
-
                 outputs = model.forward_inference(enc_output, input_seq)
-                
-                logits = outputs['logits'][:, -1, :] # Логиты для последнего токена
+                logits = outputs['logits'][:, -1, :]
             
             log_probs = F.log_softmax(logits, dim=-1)
             top_log_probs, top_ids = torch.topk(log_probs, beam_size, dim=-1)
 
-            # Создаем новых кандидатов
             for i in range(beam_size):
                 next_id = top_ids[0, i]
                 log_prob = top_log_probs[0, i].item()
-                
-                # Создаем новую последовательность, добавляя новый токен
                 new_seq = torch.cat([seq, next_id.view(1)])
                 new_score = score + log_prob
                 all_candidates.append((new_seq, new_score))
 
-        # 5. Сортируем всех кандидатов и выбираем `beam_size` лучших
         ordered = sorted(all_candidates, key=lambda x: x[1], reverse=True)
         beams = ordered[:beam_size]
         
-        # 6. Условие остановки: если все лучшие лучи закончились на EOS
         if all(b[0][-1] == eos_id for b in beams):
             break
             
-    # 7. Выбираем лучший луч, нормализуя на длину, чтобы не штрафовать длинные последовательности
     best_beam = sorted(beams, key=lambda x: x[1] / len(x[0]), reverse=True)[0]
     best_seq = best_beam[0]
     
     return best_seq
 
-# ======================================================================
-# 2. Удобная обертка для предсказания
-# ======================================================================
 
-def predict_with_beam_search(model, image_path, vocab, beam_size, device):
+def predict_with_beam_search(model, image_tensor, vocab, beam_size, device):
     """
-    Загружает изображение, запускает Beam Search и возвращает LaTeX строку.
+    ЗАПУСКАЕТ BEAM SEARCH НА ОСНОВЕ ГОТОВОГО ТЕНЗОРА ИЗОБРАЖЕНИЯ.
     """
-    # 1. Загрузка и трансформация изображения
-    transform = transforms.Compose([
-        transforms.Resize((IMG_HEIGHT, IMG_WIDTH)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
-    ])
-    image = Image.open(image_path).convert("RGB")
-    image_tensor = transform(image).to(device)
+    # <<< МЫ УДАЛИЛИ ВСЮ ЛОГИКУ ЗАГРУЗКИ И ТРАНСФОРМАЦИИ ИЗОБРАЖЕНИЯ ОТСЮДА >>>
+    # Теперь функция просто принимает image_tensor
 
     # 2. Получение ID спецтокенов
     sos_id = vocab.stoi[SOS_TOKEN]
     eos_id = vocab.stoi[EOS_TOKEN]
     pad_id = vocab.stoi[PAD_TOKEN]
 
-    # 3. Генерация
+    # 3. Генерация (теперь передаем тензор напрямую)
+    #    Убедимся, что тензор на правильном устройстве
     predicted_ids = generate_beam_search(
-        model, image_tensor, beam_size, MAX_SEQ_LEN, sos_id, eos_id, pad_id, device
+        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