import streamlit as st
from transformers import pipeline, AutoTokenizer, AutoModelForMaskedLM
import torch

# تنظیمات صفحه
st.set_page_config(page_title="دستیار هوش مصنوعی (پر کردن ماسک)", layout="wide")

# تابع بارگذاری مدل و pipeline
@st.cache_resource
def load_model_and_pipeline():
    # بارگذاری pipeline
    pipe = pipeline("fill-mask", model="HooshvareLab/bert-fa-base-uncased")

    # بارگذاری مدل و توکنایزر به صورت مستقیم
    tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-fa-base-uncased")
    model = AutoModelForMaskedLM.from_pretrained("HooshvareLab/bert-fa-base-uncased")

    return pipe, tokenizer, model

pipe, tokenizer, model = load_model_and_pipeline()

# تابع پیش‌بینی با pipeline
def predict_with_pipeline(text):
    results = pipe(text)
    return [{"word": res["token_str"], "score": res["score"]} for res in results]

# تابع پیش‌بینی با مدل مستقیم
def predict_with_model(text):
    inputs = tokenizer(text, return_tensors="pt")
    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]

    # اجرای مدل
    outputs = model(**inputs)
    logits = outputs.logits

    # پیدا کردن توکن‌های برتر
    mask_token_logits = logits[0, mask_token_index, :]
    top_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()

    results = [{"word": tokenizer.decode([token])} for token in top_tokens]
    return results

# رابط کاربری
def main():
    st.title("دستیار هوش مصنوعی (پر کردن ماسک)")

    # ورودی کاربر
    st.subheader("ورودی متن:")
    text = st.text_input(
        "متن خود را وارد کنید (از [MASK] برای نشان دادن کلمه حذف شده استفاده کنید):",
        value="ما در هوشواره معتقدیم [MASK] دانش و آگاهی می‌تواند جامعه را تغییر دهد.",
    )

    if st.button("پیش‌بینی با pipeline"):
        if "[MASK]" not in text:
            st.error("لطفاً یک متن شامل [MASK] وارد کنید.")
        else:
            st.subheader("نتایج پیش‌بینی با pipeline:")
            predictions = predict_with_pipeline(text)
            for pred in predictions:
                st.write(f"کلمه: {pred['word']} - احتمال: {pred['score']:.2f}")

    if st.button("پیش‌بینی با مدل مستقیم"):
        if "[MASK]" not in text:
            st.error("لطفاً یک متن شامل [MASK] وارد کنید.")
        else:
            st.subheader("نتایج پیش‌بینی با مدل مستقیم:")
            predictions = predict_with_model(text)
            for pred in predictions:
                st.write(f"کلمه: {pred['word']}")

if __name__ == "__main__":
    main()