import torch
import gradio as gr
from transformers import pipeline
import time


# Charger le modèle
summarizer = pipeline("summarization", model="Falconsai/text_summarization")


def generate_summary_stream(article):
    # Générer le résumé complet
    result = summarizer(
        article,
        min_length=30,
        do_sample=False
    )

    full_text = result[0]['summary_text']

    # Streamer mot par mot
    words = full_text.split()
    for i in range(len(words)):
        # Retourner les mots accumulés jusqu'à présent
        yield " ".join(words[:i + 1])
        time.sleep(0.1)  # Délai entre les mots


# Interface avec streaming
with gr.Blocks() as demo:
    gr.Markdown("# 📚  Streaming Text Summarization ")

    with gr.Row():
        input_box = gr.Textbox(
            label="Text",
            placeholder="Paste your article here...",
            lines=10
        )

    with gr.Row():
        output_box = gr.Textbox(
            label="Summary (appears progressively)",
            lines=6,
            interactive=False
        )

    generate_btn = gr.Button("Summarize", variant="primary")

    # Lier le bouton à la fonction
    generate_btn.click(
        fn=generate_summary_stream,
        inputs=input_box,
        outputs=output_box
    )

demo.launch()