# خط ۱: وارد کردن کتابخانه اصلی استریملیت
import streamlit as st
import torch # <--- این خط را اضافه کنید
# خط ۲: وارد کردن توابع مورد نیاز از مسیرهای صحیح
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from transformers.pipelines import pipeline 
# توجه: این مدل‌های AutoModel باید در سطح بالای پکیج باشند،
# و pipeline را از زیرماژول آن وارد می کنیم.

# 1. تنظیمات اولیه Streamlit
st.set_page_config(
    page_title="مدل خلاصه‌سازی انتزاعی mT5-FineTuned",
    layout="wide",
    initial_sidebar_state="expanded"
)

# مسیر مدل آموزش داده شده شما
# مدل Fine-Tuned شما در این مسیر ذخیره شده است.
FINETUNED_MODEL_PATH = "nafisehNik/mt5-persian-summary"

# 2. بارگذاری مدل (از Caching برای جلوگیری از بارگذاری مجدد استفاده می‌شود)


@st.cache_resource
def load_summarizer_pipeline():
    """
    پایپ‌لاین خلاصه‌سازی ترنسفورمر را بارگذاری می‌کند.
    از مدل آموزش‌دیده (Fine-Tuned) استفاده می‌شود.
    """

    st.subheader("💡 در حال بارگذاری مدل...")

    try:
        # ساخت pipeline خلاصه‌سازی
        summarizer = pipeline(
            "summarization",
           model=FINETUNED_MODEL_PATH,
        # تعیین دستگاه: استفاده از GPU (0) یا CPU (-1)
        # استفاده از متغیر device_id که قبلاً در کد Streamlit تعریف شده است
        device=-1 # با توجه به استفاده از CPU Basic، بهتر است -1 را اینجا سخت‌کد کنیم
        )
        st.success("✅ مدل mT5 با موفقیت بارگذاری شد.")
        return summarizer
    except Exception as e:
        st.error(
            f"❌ خطای بارگذاری مدل: مطمئن شوید پوشه '{FINETUNED_MODEL_PATH}' وجود دارد و شامل تمام فایل‌های مدل است.")
        st.code(e)
        return None


# بارگذاری مدل در شروع برنامه
summarizer = load_summarizer_pipeline()

# --- شروع واسط کاربری (UI) ---

st.title("💡 پروژه خلاصه‌سازی انتزاعی متون (mT5)")
st.markdown("این ابزار از مدل **mT5 Fine-Tuned** برای تولید خلاصه‌هایی از متون چندزبانه (مانند فارسی) استفاده می‌کند.")

if summarizer:
    # --- بخش ورودی ---

    with st.expander("📝 وارد کردن متن و تنظیمات", expanded=True):
        input_text = st.text_area(
            "متن اصلی را برای خلاصه سازی وارد کنید:",
            height=300,
            value="متن طولانی خود را در اینجا وارد کنید. این متن می‌تواند یک مقاله خبری، یک سند پژوهشی یا یک متن طولانی دیگر باشد. مدل ما معنای متن را درک کرده و با جملات جدید، آن را خلاصه خواهد کرد."
        )

        col1, col2 = st.columns(2)

        with col1:
            max_len = st.slider("حداکثر طول خلاصه (Max Length)",
                                min_value=50, max_value=300, value=150, step=10)

        with col2:
            min_len = st.slider("حداقل طول خلاصه (Min Length)",
                                min_value=20, max_value=150, value=30, step=10)

    # --- بخش پردازش ---

    if st.button("🚀 خلاصه کن", type="primary"):

        if len(input_text.split()) < 20:
            st.warning(
                "حداقل ۲۰ کلمه برای یک خلاصه‌سازی معنادار نیاز است. متن شما بسیار کوتاه است.")
        else:
            # نمایش Spinner تا زمانی که خلاصه آماده شود
            with st.spinner("درحال پردازش... مدل ترنسفورمر در حال تولید خلاصه انتزاعی است."):

                # اجرای مدل خلاصه‌سازی
                summary = summarizer(
              input_text,
                max_length=max_len,
                min_length=min_len,
                num_beams=4,
                no_repeat_ngram_size=2
                   
                )

                # 5. نمایش نتیجه
                st.subheader("✅ خلاصه تولید شده")
                st.markdown(f"***{summary[0]['summary_text']}***")

                # نمایش آمار فنی در نوار کناری (Sidebar)
                st.sidebar.markdown("---")
                st.sidebar.markdown("### آمار")
                st.sidebar.info(
                    f"طول متن ورودی: **{len(input_text.split())} کلمه**")
                st.sidebar.info(
                    f"طول خلاصه خروجی: **{len(summary[0]['summary_text'].split())} کلمه**")