import gradio as gr
from transformers import LEDTokenizer, LEDForConditionalGeneration
import torch

# โหลดโมเดลและ tokenizer
model_name = "allenai/led-large-16384-arxiv"
tokenizer = LEDTokenizer.from_pretrained(model_name)
model = LEDForConditionalGeneration.from_pretrained(model_name)

# ฟังก์ชันสรุปข้อความ
def summarize_text(text, min_len=100, max_len=300):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=16384)
    with torch.no_grad():
        summary_ids = model.generate(
            **inputs,
            max_length=max_len,
            min_length=min_len,
            num_beams=4,
            length_penalty=2.0
        )
    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
    return summary

# UI Gradio
app = gr.Interface(
    fn=summarize_text,
    inputs=[
        gr.Textbox(lines=15, placeholder="วางข้อความยาว ๆ ที่นี่...", label="Input Text"),
        gr.Slider(50, 500, value=100, label="Minimum Summary Length"),
        gr.Slider(100, 1000, value=300, label="Maximum Summary Length"),
    ],
    outputs=gr.Textbox(label="Summarized Text"),
    title="📄 Long Document Summarizer (LED 16k)",
    description="ใช้โมเดล allenai/led-large-16384-arxiv สำหรับสรุปข้อความยาวมาก เช่น บทความ งานวิจัย หรือเอกสารหลายหน้า"
)

app.launch()