Update app.py

app.py

@@ -1,7 +1,8 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 
-# بارگذاری مدل "HooshvareLab/bert-fa-base-uncased" از Hugging Face
+# بارگذاری مدل
 @st.cache_resource
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-fa-base-uncased")
@@ -10,81 +11,39 @@ def load_model():
 
 tokenizer, model = load_model()
 
-# تابع برای تولید پاسخ (فعلا فقط متن ورودی را برمی‌گرداند)
+# تولید پاسخ
 def generate_response(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-    # TODO: پردازش خروجی مدل و تولید پاسخ مناسب
-    # در حال حاضر، فقط متن ورودی را برمی‌گردانیم
-    response = text
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits
+    predicted_class = torch.argmax(logits, dim=1).item()
+    response = f"پاسخ مدل: دسته‌بندی {predicted_class}"
     return response
 
+# رابط کاربری
 def main():
     st.set_page_config(page_title="دستیار هوش مصنوعی", layout="wide")
-    st.title("دستیار هوش مصنوعی پیشرفته (Hugging Face)")
+    st.title("دستیار هوش مصنوعی (BERT فارسی)")
 
-    # مقداردهی اولیه Session State
+    # مدیریت پیام‌ها با Session State
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
-    # نمایش HTML
-    with open("index.html", "r", encoding="utf-8") as file:
-        html_content = file.read()
+    # نمایش پیام‌ها
+    for message in st.session_state.messages:
+        if message["role"] == "user":
+            st.write(f"👤 کاربر: {message['content']}")
+        else:
+            st.write(f"🤖 دستیار: {message['content']}")
 
-    # اضافه کردن کد Javascript برای ارتباط با Streamlit
-    html_content = html_content.replace("</body>", """
-        <script>
-            // تابع ارسال پیام به Streamlit
-            function sendMessage() {
-                var userInput = document.getElementById("user-input").value;
-                document.getElementById("user-input").value = ""; // پاک کردن ورودی
-                // اضافه کردن پیام کاربر به لیست پیام‌ها
-                var chatMessages = document.getElementById("chat-messages");
-                var newMessage = document.createElement("div");
-                newMessage.classList.add("message", "user-message");
-                newMessage.innerText = userInput;
-                chatMessages.appendChild(newMessage);
-                // ارسال پیام به Streamlit
-                Streamlit.setComponentValue("user-input", userInput);
-            }
-
-            // دریافت پاسخ از Streamlit و نمایش اون
-            Streamlit.on('update', function() {
-                var response = Streamlit.getComponentValue("response");
-                if (response) {
-                    // اضافه کردن پیام به لیست پیام‌ها
-                    var chatMessages = document.getElementById("chat-messages");
-                    var newMessage = document.createElement("div");
-                    newMessage.classList.add("message", "assistant-message");
-                    newMessage.innerText = response;
-                    chatMessages.appendChild(newMessage);
-                }
-            });
-
-            // ارسال پیام با فشردن Enter
-            var userInput = document.getElementById("user-input");
-            userInput.addEventListener("keyup", function(event) {
-                if (event.keyCode === 13) {
-                    event.preventDefault();
-                    document.getElementById("send-button").click();
-                }
-            });
-        </script>
-    </body>
-    """)
-
-    st.components.v1.html(html_content, height=700, scrolling=True)
-
-    # دریافت پیام از Session State
-    user_input = st.components.v1.text_input("user-input", value="", key="user_input")
+    # ورودی کاربر
+    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})
-
-    # ارسال پاسخ به Session State
-    if st.session_state.messages:
-        st.components.v1.text_area("response", value=st.session_state.messages[-1]["content"], key="response")
+        st.experimental_rerun()
 
 if __name__ == "__main__":
-    main()
+    main()