Upload 2 files

app.py

@@ -1,101 +1,129 @@
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import nltk
-import math
+import os
+import re
 import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-model_name = "AGIvan/t5-base-title-generation"
-max_input_length = 512
-
-st.header("Generate candidate titles for articles")
-
-st_model_load = st.text('Loading title generator model...')
-
-@st.cache_data()
-def load_model():
-    print("Loading model...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-    nltk.download('punkt')
-    print("Model loaded!")
-    return tokenizer, model
-
-tokenizer, model = load_model()
-st.success('Model loaded!')
-st_model_load.text("")
-
-with st.sidebar:
-    st.header("Model parameters")
-    if 'num_titles' not in st.session_state:
-        st.session_state.num_titles = 5
-    def on_change_num_titles():
-        st.session_state.num_titles = num_titles
-    num_titles = st.slider("Number of titles to generate", min_value=1, max_value=10, value=1, step=1, on_change=on_change_num_titles)
-    if 'temperature' not in st.session_state:
-        st.session_state.temperature = 0.7
-    def on_change_temperatures():
-        st.session_state.temperature = temperature
-    temperature = st.slider("Temperature", min_value=0.1, max_value=1.5, value=0.6, step=0.05, on_change=on_change_temperatures)
-    st.markdown("_High temperature means that results are more random_")
-
-if 'text' not in st.session_state:
-    st.session_state.text = ""
-st_text_area = st.text_area('Text to generate the title for', value=st.session_state.text, height=500)
-
-def generate_title():
-    st.session_state.text = st_text_area
-
-    # tokenize text
-    inputs = ["summarize: " + st_text_area]
-    inputs = tokenizer(inputs, return_tensors="pt")
-
-    # compute span boundaries
-    num_tokens = len(inputs["input_ids"][0])
-    print(f"Input has {num_tokens} tokens")
-    max_input_length = 512
-    num_spans = math.ceil(num_tokens / max_input_length)
-    print(f"Input has {num_spans} spans")
-    overlap = math.ceil((num_spans * max_input_length - num_tokens) / max(num_spans - 1, 1))
-    spans_boundaries = []
-    start = 0
-    for i in range(num_spans):
-        spans_boundaries.append([start + max_input_length * i, start + max_input_length * (i + 1)])
-        start -= overlap
-    print(f"Span boundaries are {spans_boundaries}")
-    spans_boundaries_selected = []
-    j = 0
-    for _ in range(num_titles):
-        spans_boundaries_selected.append(spans_boundaries[j])
-        j += 1
-        if j == len(spans_boundaries):
-            j = 0
-    print(f"Selected span boundaries are {spans_boundaries_selected}")
-
-    # transform input with spans
-    tensor_ids = [inputs["input_ids"][0][boundary[0]:boundary[1]] for boundary in spans_boundaries_selected]
-    tensor_masks = [inputs["attention_mask"][0][boundary[0]:boundary[1]] for boundary in spans_boundaries_selected]
-
-    inputs = {
-        "input_ids": torch.stack(tensor_ids),
-        "attention_mask": torch.stack(tensor_masks)
-    }
-
-    # compute predictions
-    outputs = model.generate(**inputs, do_sample=True, temperature=temperature)
-    decoded_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    predicted_titles = [nltk.sent_tokenize(decoded_output.strip())[0] for decoded_output in decoded_outputs]
-
-    st.session_state.titles = predicted_titles
-
-# generate title button
-st_generate_button = st.button('Generate title', on_click=generate_title)
-
-# title generation labels
-if 'titles' not in st.session_state:
-    st.session_state.titles = []
-
-if len(st.session_state.titles) > 0:
-    with st.container():
-        st.subheader("Generated titles")
-        for title in st.session_state.titles:
-            st.markdown("__" + title + "__")
+# ==== настройки ====
+# Можно задать через переменную окружения MODEL_ID в Settings → Repository secrets.
+MODEL_ID = os.environ.get("MODEL_ID", "Spyspook/my-t5-medium-summarizer")
+
+MAX_INPUT_LENGTH = 512     # вход в токенах
+MAX_TARGET_LENGTH = 64     # длина заголовка
+DEFAULT_NUM_TITLES = 3
+DEFAULT_TEMPERATURE = 0.7
+DEFAULT_BEAMS = 4          # для стабильности метрик можно 4; для разнообразия ставь do_sample=True
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ==== загрузка модели/токенизатора один раз ====
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID).to(device)
+model.eval()
+
+# Простая функция выделения первой фразы без NLTK (чтобы не тянуть ресурсы в Space)
+_SENT_END_RE = re.compile(r"([.!?])\s+")
+
+def first_sentence(text: str) -> str:
+    text = text.strip()
+    if not text:
+        return text
+    parts = _SENT_END_RE.split(text, maxsplit=1)
+    # parts = [before, sep, after] или просто [text]
+    if len(parts) >= 2:
+        return (parts[0] + parts[1]).strip()
+    return text
+
+def generate_titles(article_text, num_titles, temperature, beams, do_sample):
+    if not article_text or not article_text.strip():
+        return []
+
+    # Префикс для T5
+    prefixed = "summarize: " + article_text.strip()
+
+    # Токенизация и обрезка по контексту
+    inputs = tokenizer(
+        prefixed,
+        return_tensors="pt",
+        truncation=True,
+        max_length=MAX_INPUT_LENGTH,
+    )
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+
+    gen_kwargs = dict(
+        max_length=MAX_TARGET_LENGTH,
+        num_return_sequences=int(num_titles),
+        early_stopping=True,
+    )
+
+    # Логика генерации:
+    # - Если do_sample=True → семплирование (temperature, top_p),
+    # - иначе — детерминированный beam search.
+    if do_sample:
+        gen_kwargs.update(
+            dict(
+                do_sample=True,
+                temperature=float(temperature),
+                top_p=0.95,
+                num_beams=1,
+            )
+        )
+    else:
+        gen_kwargs.update(
+            dict(
+                do_sample=False,
+                num_beams=int(beams),
+                length_penalty=1.0,
+            )
+        )
+
+    with torch.no_grad():
+        outputs = model.generate(**inputs, **gen_kwargs)
+
+    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+
+    # Берем первую фразу и убираем дубликаты, сохраняя порядок
+    seen = set()
+    titles = []
+    for t in decoded:
+        t1 = first_sentence(t)
+        if t1 and t1 not in seen:
+            seen.add(t1)
+            titles.append(t1)
+
+    # Вернем как список списков для удобной таблицы
+    return [[t] for t in titles]
+
+# ==== интерфейс Gradio ====
+with gr.Blocks() as demo:
+    gr.Markdown("## T5 Article Title Generator")
+
+    with gr.Row():
+        text_in = gr.Textbox(
+            label="Article text",
+            placeholder="Paste article text here…",
+            lines=14,
+        )
+
+    with gr.Row():
+        num_titles = gr.Slider(1, 10, value=DEFAULT_NUM_TITLES, step=1, label="Number of titles")
+        temperature = gr.Slider(0.1, 1.5, value=DEFAULT_TEMPERATURE, step=0.05, label="Temperature (sampling)")
+    with gr.Row():
+        beams = gr.Slider(1, 8, value=DEFAULT_BEAMS, step=1, label="Beams (if sampling is OFF)")
+        do_sample = gr.Checkbox(value=True, label="Use sampling (ON) / Beam search (OFF)")
+
+    generate_btn = gr.Button("Generate")
+    out_table = gr.Dataframe(headers=["Title"], row_count=(0, "dynamic"), wrap=True)
+
+    generate_btn.click(
+        fn=generate_titles,
+        inputs=[text_in, num_titles, temperature, beams, do_sample],
+        outputs=out_table,
+        api_name="generate",
+    )
+
+# Для HF Spaces достаточно экспортировать переменную приложения
+app = demo
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1,3 +1,4 @@
-nltk
 torch
-transformers
+transformers
+gradio
+sentencepiece