import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# โหลด tokenizer และโมเดล
model_name = "chanyaphas/summarize"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def main():
    st.title("Text Summarization App")

    # สร้าง text area สำหรับป้อนข้อความ
    text = st.text_area('Enter the text to summarize:', '')

    # คลิกปุ่ม "Summarize" เพื่อสรุปข้อความ
    if st.button('Summarize'):
        if text:
            # ใช้โมเดลเพื่อสร้างสรุปข้อความ
            input_ids = tokenizer(text, return_tensors="pt", padding=True, truncation=True).input_ids
            with torch.no_grad():
                output_ids = model.generate(input_ids)

            # แปลงผลลัพธ์กลับเป็นข้อความ
            summary = tokenizer.decode(output_ids[0], skip_special_tokens=True)

            # แสดงสรุปข้อความ
            st.subheader("Summary:")
            st.write(summary)
        else:
            st.warning("Please enter some text to summarize.")

if __name__ == "__main__":
    main()