Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import re
 import gc
-from typing import List, Dict, Tuple, Iterable, Optional, Any  # <-- FIX: добавили Any
+from typing import List, Dict, Tuple, Iterable, Optional, Any
 
 import gradio as gr
 import numpy as np
@@ -53,7 +53,8 @@ def _load_encoder(name: str):
 
 
 def _load_seq2seq(name: str):
-    tok = AutoTokenizer.from_pretrained(name, use_fast=False)  # важно для T5/MT5
+    # критично для T5/MT5: use_fast=False
+    tok = AutoTokenizer.from_pretrained(name, use_fast=False)
     model = AutoModelForSeq2SeqLM.from_pretrained(
         name,
         torch_dtype=DTYPE,
@@ -67,6 +68,7 @@ print("Loading models...")
 emb_tok, emb_model = _load_encoder(EMB_MODEL_NAME)
 gen_tok, gen_model = _load_seq2seq(PRIMARY_GEN_MODEL)
 
+# fallback лениво (экономия памяти)
 fb_tok = None
 fb_model = None
 print("Models loaded.")
@@ -224,7 +226,7 @@ def retrieve_topk(query: str, embeddings_f16: np.ndarray, top_k: int = 4) -> Tup
 
 
 # =======================
-# ГЕНЕРАЦИЯ + САНИТИЗАЦИЯ
+# ГЕНЕРАЦИЯ: защита от OOR
 # =======================
 BANNED = [
     "контекст", "вопрос:", "ответ:", "правила", "требования",
@@ -265,26 +267,64 @@ def looks_bad(text: str) -> bool:
     return False
 
 
+def _ensure_embeddings_size(tokenizer, model, required_size: int):
+    emb = model.get_input_embeddings()
+    cur = int(emb.num_embeddings)
+    if required_size > cur:
+        model.resize_token_embeddings(required_size)
+        # на всякий случай вернём на нужное устройство после ресайза
+        model.to(DEVICE)
+
+
 @torch.inference_mode()
 def seq2seq_generate(tokenizer, model, prompt: str, max_new_tokens: int = 220, max_input_tokens: int = 512) -> str:
     prompt = (prompt or "").strip()
     if not prompt:
         return ""
+
     batch = tokenizer(
         prompt,
         return_tensors="pt",
         truncation=True,
         max_length=max_input_tokens,
     )
+
+    # ВАЖНО: before .to(DEVICE) можно посчитать max_id на CPU
+    input_ids = batch["input_ids"]
+    max_id = int(input_ids.max().item()) if input_ids.numel() else 0
+    needed = max(int(len(tokenizer)), max_id + 1)
+
+    _ensure_embeddings_size(tokenizer, model, needed)
+
+    # После возможного resize — переносим на устройство
     batch = {k: v.to(DEVICE) for k, v in batch.items()}
-    out_ids = model.generate(
-        **batch,
-        max_new_tokens=max_new_tokens,
-        num_beams=4,
-        do_sample=False,
-        no_repeat_ngram_size=3,
-        early_stopping=True,
-    )
+
+    # Доп. страховка: если по какой-то причине всё ещё OOR — зажмём
+    emb_size = int(model.get_input_embeddings().num_embeddings)
+    if int(batch["input_ids"].max().item()) >= emb_size:
+        batch["input_ids"] = batch["input_ids"].clamp_max(emb_size - 1)
+
+    try:
+        out_ids = model.generate(
+            **batch,
+            max_new_tokens=max_new_tokens,
+            num_beams=4,
+            do_sample=False,
+            no_repeat_ngram_size=3,
+            early_stopping=True,
+        )
+    except IndexError:
+        # retry: синхронизируем по len(tokenizer) и повторяем
+        _ensure_embeddings_size(tokenizer, model, int(len(tokenizer)))
+        out_ids = model.generate(
+            **batch,
+            max_new_tokens=max_new_tokens,
+            num_beams=4,
+            do_sample=False,
+            no_repeat_ngram_size=3,
+            early_stopping=True,
+        )
+
     return tokenizer.decode(out_ids[0], skip_special_tokens=True).strip()
 
 
@@ -302,7 +342,7 @@ def generate_clean(primary_prompt: str, fallback_prompt: str) -> str:
 
 
 # =======================
-# КОНСПЕКТ ПО ЧАСТИ
+# КОНСПЕКТ
 # =======================
 def summarize_part(part_1based: int, state: Dict[str, Any]) -> Tuple[str, str, str]:
     chunks: List[str] = state.get("chunks", [])
@@ -601,4 +641,5 @@ with gr.Blocks(title="EduMultiSpace") as demo:
 
 
 if __name__ == "__main__":
-    demo.launch()
+    # чтобы не было параллельных генераций, которые могут раздувать память на Spaces
+    demo.queue(concurrency_count=1, max_size=16).launch()