Upload 2 files

app.py

@@ -1,36 +1,58 @@
+import math
+import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import nltk, math, torch
+import nltk
 
-MODEL_ID = "Ilyakk/t5-summarization" 
-MAX_INPUT_LEN = 512
+
+MODEL_ID = "Ilyakk/t5-summarization"   
+MAX_INPUT_LEN = 512                    
+GEN_MAX_LEN = 64                    
+
+def ensure_nltk():
+    try:
+        nltk.data.find("tokenizers/punkt")
+    except LookupError:
+        nltk.download("punkt")
+    try:
+        nltk.data.find("tokenizers/punkt_tab")
+    except LookupError:
+        nltk.download("punkt_tab")
+
+ensure_nltk()
 
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
-
-try:
-    nltk.data.find("tokenizers/punkt")
-except LookupError:
-    nltk.download("punkt")
+def _first_sentence(text: str) -> str:
+    text = (text or "").strip()
+    if not text:
+        return ""
+    try:
+        sents = nltk.sent_tokenize(text)
+        return sents[0].strip() if sents else text
+    except Exception:
+        for sep in [".", "!", "?"]:
+            if sep in text:
+                return text.split(sep)[0].strip()
+        return text
 
 def generate_titles(text: str, num_titles: int = 3, temperature: float = 0.7):
     text = (text or "").strip()
     if not text:
         return ["Введите текст статьи выше."]
 
-    
-    enc = tokenizer(["summarize: " + text], return_tensors="pt")
-    ids  = enc["input_ids"][0]
+    enc = tokenizer(["summarize: " + text], return_tensors="pt", truncation=False)
+    ids = enc["input_ids"][0]
     mask = enc["attention_mask"][0]
-
-   
     num_tokens = len(ids)
+
     num_spans = max(1, math.ceil(num_tokens / MAX_INPUT_LEN))
     overlap = math.ceil((num_spans * MAX_INPUT_LEN - num_tokens) / max(num_spans - 1, 1)) if num_spans > 1 else 0
 
-   
     spans = []
     start = 0
     for i in range(num_spans):
@@ -39,24 +61,26 @@ def generate_titles(text: str, num_titles: int = 3, temperature: float = 0.7):
         spans.append([max(0, b0), min(num_tokens, b1)])
         start -= overlap
 
-  
     chosen = [spans[i % len(spans)] for i in range(num_titles)]
 
-    batch_ids  = [ids[b0:b1]  for (b0, b1) in chosen]
+    batch_ids = [ids[b0:b1] for (b0, b1) in chosen]
     batch_mask = [mask[b0:b1] for (b0, b1) in chosen]
-    batch = {"input_ids": torch.stack(batch_ids), "attention_mask": torch.stack(batch_mask)}
+    batch = {
+        "input_ids": torch.stack(batch_ids).to(device),
+        "attention_mask": torch.stack(batch_mask).to(device),
+    }
 
     with torch.no_grad():
         outputs = model.generate(
             **batch,
             do_sample=True,
             temperature=float(temperature),
-            max_length=64,          
-            num_beams=1             
+            max_length=GEN_MAX_LEN,
+            num_beams=1
         )
 
     decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    titles = [ (nltk.sent_tokenize(d.strip())[0] if d.strip() else "").strip() for d in decoded ]
+    titles = [_first_sentence(d) for d in decoded]
     return titles
 
 demo = gr.Interface(
@@ -72,4 +96,5 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
+    
     demo.launch()

requirements.txt

@@ -1,5 +1,5 @@
 transformers
 torch
-nltk
 gradio
 sentencepiece
+nltk>=3.8.1