import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification  # Thay bằng mô hình của bạn
import torch

# Load mô hình (thay thế bằng mô hình sửa lỗi chính tả của bạn)
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("Diezu/bat_pho_bo")  # Thay bằng tên mô hình của bạn
    model = AutoModelForSequenceClassification.from_pretrained("Diezu/bat_pho_bo")
    return tokenizer, model

# Hàm phát hiện lỗi chính tả
def detect_errors(text, tokenizer, model):
    errors = []
    words = text.split()
    for word in words:
        inputs = tokenizer(word, return_tensors="pt", padding=True, truncation=True)
        with torch.no_grad():
            outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=-1)
        predicted_class = torch.argmax(probabilities, dim=-1).item()
        if predicted_class == 1:  # Nếu từ sai
            errors.append(word)
    return errors

# Streamlit App
st.title("Công cụ phát hiện lỗi chính tả")
# CSS tùy chỉnh
st.markdown(
    """
    <style>
    .stTextArea textarea {
        font-size: 18px; /* Thay đổi kích thước font */
        font-weight: bold; /* Làm đậm chữ */
        weight : 30%;
    }
    </style>
    """,
    unsafe_allow_html=True
)
# Tải mô hình
tokenizer, model = load_model()

# Tạo bố cục 2 cột
col1, col2 = st.columns([1, 1])  # Hai cột có kích thước bằng nhau

# Cột bên trái: Nhập văn bản
with col1:
    st.header("Nhập văn bản")
    input_text = st.text_area(
        "Nhập văn bản:", 
        height=300, 
        max_chars=5000, 
        placeholder="Nhập văn bản của bạn ở đây..."
    )

# Cột bên phải: Hiển thị kết quả
with col2:
    st.header("Kết quả phát hiện lỗi")
    if st.button("Phát hiện lỗi", use_container_width=True):
        if input_text.strip():
            errors = detect_errors(input_text, tokenizer, model)
            if errors:
                st.text_area(
                    "Các từ phát hiện lỗi:",
                    value=", ".join(errors),
                    height=300,
                    disabled=True,
                )
            else:
                st.success("Không phát hiện lỗi nào trong văn bản.")
        else:
            st.warning("Vui lòng nhập văn bản để kiểm tra.")