Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,35 +1,54 @@
 import streamlit as st
 from transformers import pipeline
+import torch
 
-# 1. Настройка внешнего вида сайта
-st.set_page_config(page_title="HiperAI", page_icon="🤖")
-st.title("HiperAI — Моя личная нейросеть")
-st.markdown("---")
+# Настройка страницы
+st.set_page_config(page_title="HiperAI Pro", page_icon="🧠", layout="centered")
 
-# 2. Загрузка модели (кешируем, чтобы не грузить при каждом клике)
+# Кастомный CSS для красоты
+st.markdown("""
+    <style>
+    .main { background-color: #0e1117; }
+    .stChatMessage { border-radius: 15px; padding: 10px; margin-bottom: 10px; }
+    </style>
+    """, unsafe_allow_input_with_html=True)
+
+st.title("🚀 HiperAI: Твой личный ИИ")
+
+# Загрузка модели (Qwen2.5 — одна из лучших маленьких моделей для русского языка)
 @st.cache_resource
 def load_model():
-    # Используем TinyLlama — она идеально подходит для бесплатного хостинга
-    return pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device_map="auto")
-
-with st.spinner('Загрузка мозгов ИИ...'):
-    generator = load_model()
-
-# 3. Поле ввода
-user_input = st.text_input("Напиши сообщение своему ИИ (лучше на английском):", placeholder="Hello, who are you?")
-
-if user_input:
-    with st.spinner('HiperAI думает...'):
-        # Настройка генерации: max_new_tokens — длина ответа, temperature — креативность
-        result = generator(user_input, max_new_tokens=100, do_sample=True, temperature=0.7)
-        
-        # Вырезаем ответ (убираем сам вопрос из начала текста)
-        full_text = result[0]['generated_text']
-        answer = full_text.replace(user_input, "").strip()
-        
-        # Вывод на экран
-        st.write("### Ответ:")
-        st.info(answer if answer else "ИИ задумался и ничего не выдал. Попробуй еще раз!")
-
-st.markdown("---")
-st.caption("Работает локально на Hugging Face Spaces без использования внешних API.")
+    model_id = "Qwen/Qwen2.5-1.5B-Instruct"
+    return pipeline("text-generation", model=model_id, device_map="auto", torch_dtype="auto")
+
+generator = load_model()
+
+# Инициализация истории чата
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Отображение истории чата
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Поле ввода
+if prompt := st.chat_input("Спроси меня о чем угодно..."):
+    # Добавляем сообщение пользователя в историю
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    # Ответ ИИ
+    with st.chat_message("assistant"):
+        with st.spinner("Думаю..."):
+            # Формируем промпт для диалога
+            messages = [{"role": "user", "content": prompt}]
+            
+            # Генерация
+            output = generator(messages, max_new_tokens=512, do_sample=True, temperature=0.7)
+            response = output[0]['generated_text'][-1]['content']
+            
+            st.markdown(response)
+            st.session_state.messages.append({"role": "assistant", "content": response})
+