import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# بارگذاری مدل
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("HooshvareLab/bert-fa-base-uncased")
    model = AutoModelForSequenceClassification.from_pretrained("HooshvareLab/bert-fa-base-uncased")
    return tokenizer, model

tokenizer, model = load_model()

# تولید پاسخ
def generate_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("دستیار هوش مصنوعی (BERT فارسی)")

    # مدیریت پیام‌ها با Session State
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # نمایش پیام‌ها
    for message in st.session_state.messages:
        if message["role"] == "user":
            st.write(f"👤 کاربر: {message['content']}")
        else:
            st.write(f"🤖 دستیار: {message['content']}")

    # ورودی کاربر
    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 __name__ == "__main__":
    main()