import streamlit as st
from transformers import pipeline, AutoModelForSeq2SeqLM
from peft import PeftModel, PeftConfig
# from pathlib import Path
st.title("Thai News Summarizer")
st.subheader("Thai News Summarizer")

text = st.text_area('News', '', height=200)
# st.write('This KHAW? -->', text)

peft_model_id = "Saltywan/mt5-base-lora-thaisum-news"
config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model_name_or_path)
model = PeftModel.from_pretrained(model, peft_model_id)
summarizer = pipeline("summarization", model=model, tokenizer=config.base_model_name_or_path)

if st.button('Summarize'):
    text = summarizer("summarize: " + text, max_length=256)
    desired_text = text[0]['summary_text']
    st.write(desired_text)