Update app.py

app.py

@@ -1,76 +1,44 @@
 import streamlit as st
-from transformers import pipeline
-import time
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Настройки для экономии памяти
-MAX_LENGTH = 200
-MODEL_NAME = "sberbank-ai/rugpt3small_based_on_gpt2"
+st.set_page_config(page_title="ИТМО Магистратура Чат-бот", page_icon="🎓")
+st.title("🎓 Чат-бот про магистратуру ИТМО на русском")
 
-st.set_page_config(page_title="Магистратура Помощник")
-st.title("🎓 Помощник по вопросам магистратуры")
-st.write("Задайте вопрос о поступлении или обучении в магистратуре")
+MODEL_NAME = "Grossmend/rudialogpt3_medium_based_on_gpt2"
 
-# Кешируем загрузку модели
-@st.cache_resource(show_spinner=False)
+@st.cache_resource
 def load_model():
-    with st.spinner("Загружаем модель... Это займет около 30 секунд"):
-        return pipeline(
-            "text-generation",
-            model=MODEL_NAME,
-            device_map="auto"
-        )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    if torch.cuda.is_available():
+        model = model.to('cuda')
+    return tokenizer, model
 
-generator = load_model()
+tokenizer, model = 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)}"
+SYSTEM_PROMPT = """Вы являетесь виртуальным помощником для абитуриентов магистратуры Университета ИТМО.
+Отвечаете на вопросы о магистерских программах ИТМО.
+"""
 
-# Форма ввода
-with st.form("question_form"):
-    user_input = st.text_input("Ваш вопрос:", placeholder="Что нужно для поступления?")
-    submit_button = st.form_submit_button("Спросить")
+user_input = st.text_input("Введите ваш вопрос про магистратуру ИТМО:")
 
-# Обработка вопроса
-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()
+if user_input:
+    input_text = SYSTEM_PROMPT + "\n" + user_input
+    inputs = tokenizer(input_text, return_tensors="pt")
 
-# Показываем историю диалога
-for i, (question, answer) in enumerate(st.session_state.history[::-1]):
-    st.divider()
-    st.subheader(f"❓ {question}")
-    st.write(f"💡 {answer}")
+    if torch.cuda.is_available():
+        inputs = {k: v.to('cuda') for k, v in inputs.items()}
 
-# Информация о модели
-st.sidebar.markdown("""
-### О приложении
-Использует русскоязычную модель [rugpt3small](https://huggingface.co/sberbank-ai/rugpt3small_based_on_gpt2) для ответов на вопросы о магистратуре.
-**Примеры вопросов:**
-- Какие документы нужны для поступления?
-- Сколько длится обучение?
-- Чем магистратура отличается от бакалавриата?
-- Какие есть направления?
-- Есть ли бюджетные места?
-""")
+    outputs = model.generate(**inputs, max_length=500, do_sample=True, temperature=0.7, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    reply = response[len(input_text):].strip()
+
+    st.session_state.history.append((user_input, reply))
+
+for i, (q, a) in enumerate(st.session_state.history):
+    st.markdown(f"**Вы:** {q}")
+    st.markdown(f"**Бот:** {a}")