import streamlit as st
from model.model import Seq2SeqModel, predict, model_tokenize
from model.model import Seq2SeqModel, predict, model_tokenize
from transformers import AutoTokenizer

model_args = {
    "max_seq_length": 512,
    "max_length": 32,
    "manual_seed": 42
}

model = Seq2SeqModel(
    encoder_decoder_type="bartpho",
    encoder_decoder_name="result",  # Checkpoint của model
    args=model_args,
)

tokenizer = AutoTokenizer.from_pretrained("vinai/bartpho-word-base")



# Giao diện Streamlit
st.title("Đánh giá cảm xúc theo khía cạnh")
st.write("Nhập văn bản tiếng Việt để mô hình dự đoán.")

# Input từ người dùng
input_text = st.text_area("Nhập văn bản tại đây:", placeholder="Ví dụ: nhà hàng này ăn ngon ...")

if st.button("Dự đoán"):
    if input_text.strip():
        # Gọi hàm predict và hiển thị kết quả
        with st.spinner("Đang xử lý..."):
            try:
                result = predict(model.model, input_text, tokenizer, model_tokenize, processed=False, printout=True)
                st.success("Dự đoán hoàn tất!")
                st.write("Kết quả dự đoán:")
                st.write(result)
            except RuntimeError as e:
                st.error(f"Lỗi khi chạy mô hình: {e}")
        st.spinner("Dự đoán hoàn tất!")
    else:
        st.error("Vui lòng nhập văn bản để dự đoán!")