Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,53 +3,69 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 from threading import Thread
 import torch
 
-# Настройка интерфейса
-st.set_page_config(page_title="HiperAI Ultra", page_icon="🚀")
-st.title("🚀 HiperAI Ultra")
-st.caption("Режим: Максимальная скорость + Стриминг")
+# 1. Настройка страницы
+st.set_page_config(page_title="HiperAI Turbo Pro", page_icon="🏎️", layout="centered")
 
-# 1. Загрузка модели с оптимизацией под CPU
+# Кастомный стиль для чата
+st.markdown("""
+    <style>
+    .stChatMessage { background-color: #1e2129; border-radius: 10px; }
+    </style>
+    """, unsafe_allow_input_with_html=True)
+
+st.title("🏎️ HiperAI Ultra Speed")
+st.caption("Оптимизировано через Optimum + BetterTransformer")
+
+# 2. Загрузка модели и токенайзера
 @st.cache_resource
-def load_model():
+def load_optimized_model():
     model_id = "Qwen/Qwen2.5-1.5B-Instruct"
+    
+    # Загружаем токенайзер (нужен sentencepiece)
     tokenizer = AutoTokenizer.from_pretrained(model_id)
+    
+    # Загружаем модель
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        torch_dtype=torch.float32, 
+        torch_dtype=torch.float32,
         low_cpu_mem_usage=True,
         device_map="cpu"
     )
+
+    # Включаем Optimum BetterTransformer для ускорения CPU
+    try:
+        from optimum.bettertransformer import BetterTransformer
+        model = BetterTransformer.transform(model)
+        st.sidebar.success("🚀 Optimum BetterTransformer активен!")
+    except Exception as e:
+        st.sidebar.warning(f"Optimum не запустился, работаем в обычном режиме.")
+        
     return tokenizer, model
 
-tokenizer, model = load_model()
+with st.spinner("Запуск двигателя ИИ..."):
+    tokenizer, model = load_optimized_model()
 
-# 2. Инициализация памяти
+# 3. Работа с историей сообщений
 if "messages" not in st.session_state:
     st.session_state.messages = []
 
-# Боковая панель
-with st.sidebar:
-    st.title("⚙️ Настройки")
-    if st.button("🗑️ Очистить чат"):
-        st.session_state.messages = []
-        st.rerun()
-    st.info("ИИ оптимизирован для работы на CPU")
-
-# Отображение истории (последние 10 сообщений)
+# Отображаем чат (последние 10 реплик для экономии памяти)
 for message in st.session_state.messages[-10:]:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# 3. Основная логика работы
-if prompt := st.chat_input("Напиши HiperAI..."):
+# 4. Поле ввода и логика ответа
+if prompt := st.chat_input("Спроси HiperAI..."):
+    # Показываем сообщение пользователя
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.markdown(prompt)
 
+    # Ответ ассистента со стримингом
     with st.chat_message("assistant"):
-        # Формируем компактный контекст для ускорения
-        context = [{"role": "system", "content": "Ты HiperAI, быстрый и умный помощник. Отвечай на русском."}]
-        context += st.session_state.messages[-5:] # Помним только последние 5 реплик
+        # Ограниченный контекст для скорости
+        context = [{"role": "system", "content": "Ты HiperAI, быстрый и точный ИИ."}]
+        context += st.session_state.messages[-5:]
         
         inputs = tokenizer.apply_chat_template(
             context, 
@@ -57,32 +73,38 @@ if prompt := st.chat_input("Напиши HiperAI..."):
             return_tensors="pt"
         ).to(model.device)
 
-        # Настройка стримера для эффекта печати
+        # Стример для вывода текста по буквам
         streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
         
-        # Параметры для самой быстрой генерации
+        # Параметры генерации
         generation_kwargs = dict(
             input_ids=inputs,
             streamer=streamer,
             max_new_tokens=512,
-            do_sample=True, # Оставляем немного креативности
+            do_sample=True,
             temperature=0.7,
-            top_k=50,
-            top_p=0.95,
-            use_cache=True # Ускоряет повторную генерацию
+            use_cache=True
         )
 
-        # Запуск в отдельном потоке (Thread)
+        # Запуск в отдельном потоке
         thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
 
-        # Вывод текста по мере появления
+        # Функция для st.write_stream
         def stream_output():
-            full_response = ""
+            full_text = ""
             for new_text in streamer:
-                full_response += new_text
+                full_text += new_text
                 yield new_text
-            st.session_state.messages.append({"role": "assistant", "content": full_response})
+            st.session_state.messages.append({"role": "assistant", "content": full_text})
 
+        # Запускаем эффект печати
         st.write_stream(stream_output)
+
+# Боковое меню
+with st.sidebar:
+    st.title("Управление")
+    if st.button("🗑️ Очистить память"):
+        st.session_state.messages = []
+        st.rerun()