import streamlit as st
from transformers import pipeline
import time

# Настройки для экономии памяти
MAX_LENGTH = 200
MODEL_NAME = "sberbank-ai/rugpt3small_based_on_gpt2"

st.set_page_config(page_title="Магистратура Помощник")
st.title("🎓 Помощник по вопросам магистратуры")
st.write("Задайте вопрос о поступлении или обучении в магистратуре")

# Кешируем загрузку модели
@st.cache_resource(show_spinner=False)
def load_model():
    with st.spinner("Загружаем модель... Это займет около 30 секунд"):
        return pipeline(
            "text-generation",
            model=MODEL_NAME,
            device_map="auto"
        )

generator = load_model()

# История диалога
if "history" not in st.session_state:
    st.session_state.history = []

# Функция генерации ответа
def generate_answer(question):
    prompt = f"Вопрос о магистратуре: {question}\nОтвет:"
    try:
        result = generator(
            prompt,
            max_length=MAX_LENGTH,
            num_return_sequences=1,
            temperature=0.7,
            repetition_penalty=1.2,
            pad_token_id=50256
        )
        return result[0]['generated_text'].split("Ответ:")[-1].strip()
    except Exception as e:
        return f"Ошибка: {str(e)}"

# Форма ввода
with st.form("question_form"):
    user_input = st.text_input("Ваш вопрос:", placeholder="Что нужно для поступления?")
    submit_button = st.form_submit_button("Спросить")

# Обработка вопроса
if submit_button and user_input:
    with st.spinner("Генерируем ответ..."):
        start_time = time.time()
        response = generate_answer(user_input)
        st.session_state.history.append((user_input, response))
        
        # Очищаем поле ввода после отправки
        st.rerun()

# Показываем историю диалога
for i, (question, answer) in enumerate(st.session_state.history[::-1]):
    st.divider()
    st.subheader(f"❓ {question}")
    st.write(f"💡 {answer}")

# Информация о модели
st.sidebar.markdown("""
### О приложении
Использует русскоязычную модель [rugpt3small](https://huggingface.co/sberbank-ai/rugpt3small_based_on_gpt2) для ответов на вопросы о магистратуре.

**Примеры вопросов:**
- Какие документы нужны для поступления?
- Сколько длится обучение?
- Чем магистратура отличается от бакалавриата?
- Какие есть направления?
- Есть ли бюджетные места?
""")