Update app.py

app.py

@@ -1,49 +1,73 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import pipeline, AutoTokenizer, AutoModelForMaskedLM
 import torch
 
-# بارگذاری مدل
+# تنظیمات صفحه
+st.set_page_config(page_title="دستیار هوش مصنوعی (پر کردن ماسک)", layout="wide")
+
+# تابع بارگذاری مدل و pipeline
 @st.cache_resource
-def load_model():
+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 = AutoModelForSequenceClassification.from_pretrained("HooshvareLab/bert-fa-base-uncased")
-    return tokenizer, model
+    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]
 
-tokenizer, model = load_model()
+# تابع پیش‌بینی با مدل مستقیم
+def predict_with_model(text):
+    inputs = tokenizer(text, return_tensors="pt")
+    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
 
-# تولید پاسخ
-def generate_response(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
+    # اجرای مدل
+    outputs = model(**inputs)
     logits = outputs.logits
-    predicted_class = torch.argmax(logits, dim=1).item()
-    response = f"پاسخ مدل: دسته‌بندی {predicted_class}"
-    return response
+
+    # پیدا کردن توکن‌های برتر
+    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.set_page_config(page_title="دستیار هوش مصنوعی", layout="wide")
-    st.title("دستیار هوش مصنوعی (BERT فارسی)")
+    st.title("دستیار هوش مصنوعی (پر کردن ماسک)")
 
-    # مدیریت پیام‌ها با Session State
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
+    # ورودی کاربر
+    st.subheader("ورودی متن:")
+    text = st.text_input(
+        "متن خود را وارد کنید (از [MASK] برای نشان دادن کلمه حذف شده استفاده کنید):",
+        value="ما در هوشواره معتقدیم [MASK] دانش و آگاهی می‌تواند جامعه را تغییر دهد.",
+    )
 
-    # نمایش پیام‌ها
-    for message in st.session_state.messages:
-        if message["role"] == "user":
-            st.write(f"👤 کاربر: {message['content']}")
+    if st.button("پیش‌بینی با pipeline"):
+        if "[MASK]" not in text:
+            st.error("لطفاً یک متن شامل [MASK] وارد کنید.")
         else:
-            st.write(f"🤖 دستیار: {message['content']}")
+            st.subheader("نتایج پیش‌بینی با pipeline:")
+            predictions = predict_with_pipeline(text)
+            for pred in predictions:
+                st.write(f"کلمه: {pred['word']} - احتمال: {pred['score']:.2f}")
 
-    # ورودی کاربر
-    user_input = st.text_input("پیام خود را وارد کنید:")
-    if user_input:
-        st.session_state.messages.append({"role": "user", "content": user_input})
-        response = generate_response(user_input)
-        st.session_state.messages.append({"role": "assistant", "content": response})
-        st.experimental_rerun()
+    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()