import gradio as gr
from transformers import pipeline, AutoTokenizer
from optimum.onnxruntime import ORTModelForSeq2SeqLM

# Загружаем модель с явным указанием имён файлов
model = ORTModelForSeq2SeqLM.from_pretrained(
    "Genn9508/Sss",
    file_name="encoder_model_quantized.onnx",
    decoder_file_name="decoder_with_past_model_quantized.onnx",
)

# Токенизатор
tokenizer = AutoTokenizer.from_pretrained("Genn9508/Sss")

# Пайплайн
pipe = pipeline("summarization", model=model, tokenizer=tokenizer)

# Функция суммаризации
def summarize(text):
    if not text.strip():
        return "Введите текст для краткого изложения."
    
    # Генерация краткого варианта
    result = pipe(
        text,
        max_length=100,      # максимум слов в выводе
        min_length=20,       # минимум — чтобы не было слишком коротко
        do_sample=False,     # детерминированная генерация
        no_repeat_ngram_size=2,
        truncation=True      # если текст длиннее 512 токенов
    )
    return result[0]['summary_text']

# Интерфейс
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📝 ONNX T5 — Суммаризатор текста")
    gr.Markdown("Модель оптимизирована для **краткого изложения текста**. Просто введите любой текст — и получите краткую версию.")

    with gr.Row():
        with gr.Column(scale=2):
            input_text = gr.Textbox(
                label="Текст для краткого изложения",
                placeholder="Введите текст (статью, отзыв, письмо и т.д.) — модель сделает кратко.",
                lines=6
            )
            btn = gr.Button("✨ Сделать кратко", variant="primary")

        with gr.Column(scale=2):
            output_text = gr.Textbox(
                label="Краткое изложение",
                lines=6,
                interactive=False
            )

    btn.click(fn=summarize, inputs=input_text, outputs=output_text)

    gr.Examples(
        [
            ["Искусственный интеллект меняет мир. Он используется в медицине, образовании, транспорте и других сферах. Благодаря ИИ люди могут решать сложные задачи быстрее и точнее."],
            ["Компания запустила новый продукт. Он помогает автоматизировать рутинные задачи. Пользователи уже оставили положительные отзывы."],
            ["На прошлой неделе прошла конференция по искусственному интеллекту. Участники обсудили последние достижения в области генеративных моделей и этики ИИ."]
        ],
        inputs=input_text,
        label="Примеры текстов"
    )

demo.launch()