import streamlit as st
import torch
from transformers import T5Tokenizer, T5ForConditionalGeneration


model_name = "sarahai/ruT5-base-summarizer"  
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

def summarize(text):


    try:
        
        input_ids = tokenizer(text, return_tensors="pt", padding="max_length").input_ids

        
        outputs = model.generate(
            input_ids,
            max_length=250,
            min_length=150,
            length_penalty=2.0,
            num_beams=4,
            early_stopping=True,
        )
        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)

        return summary

    except Exception as e:
        return f"Ошибка: {str(e)}"

st.title("Суммаризатор")
st.write("Введите текст :")

user_input = st.text_area("", height=200)

if st.button("Суммаризовать"):
    if user_input:
        summary = summarize(user_input)
        st.success(f"Результат:\n{summary}")
    else:
        st.warning("Пожалуйста, введите хотя бы маленький текст")