Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,55 +1,88 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
 import torch
 
-# Настройка страницы
-st.set_page_config(page_title="HiperAI Pro", page_icon="⚡", layout="centered")
+# Настройка интерфейса
+st.set_page_config(page_title="HiperAI Ultra", page_icon="🚀")
+st.title("🚀 HiperAI Ultra")
+st.caption("Режим: Максимальная скорость + Стриминг")
 
-# Загрузка модели
+# 1. Загрузка модели с оптимизацией под CPU
 @st.cache_resource
 def load_model():
-    # Qwen2.5 — топ для русского языка
     model_id = "Qwen/Qwen2.5-1.5B-Instruct"
-    return pipeline("text-generation", model=model_id, device_map="auto", torch_dtype="auto")
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.float32, 
+        low_cpu_mem_usage=True,
+        device_map="cpu"
+    )
+    return tokenizer, model
 
-generator = load_model()
+tokenizer, model = load_model()
 
-# --- БОКОВАЯ ПАНЕЛЬ (SIDEBAR) ---
+# 2. Инициализация памяти
+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()
-    
-    temp = st.slider("Креативность (Temperature)", 0.1, 1.5, 0.7)
-    st.info("Выше = интереснее, ниже = точнее.")
-
-# --- ЛОГИКА ЧАТА ---
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+    st.info("ИИ оптимизирован для работы на CPU")
 
-# Отображение чата
-for message in st.session_state.messages:
+# Отображение истории (последние 10 сообщений)
+for message in st.session_state.messages[-10:]:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
 
-# Ввод сообщения
-if prompt := st.chat_input("Напиши сообщение..."):
+# 3. Основная логика работы
+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"):
-        with st.spinner("HiperAI думает..."):
-            # Формируем контекст с личностью
-            chat = [
-                {"role": "system", "content": "Ты — HiperAI, крутой и полезный ИИ-помощник. Отвечай кратко и по делу на русском языке."},
-                {"role": "user", "content": prompt}
-            ]
-            
-            output = generator(chat, max_new_tokens=512, do_sample=True, temperature=temp)
-            response = output[0]['generated_text'][-1]['content']
-            
-            st.markdown(response)
-            st.session_state.messages.append({"role": "assistant", "content": response})
-            
+        # Формируем компактный контекст для ускорения
+        context = [{"role": "system", "content": "Ты HiperAI, быстрый и умный помощник. Отвечай на русском."}]
+        context += st.session_state.messages[-5:] # Помним только последние 5 реплик
+        
+        inputs = tokenizer.apply_chat_template(
+            context, 
+            add_generation_prompt=True, 
+            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, # Оставляем немного креативности
+            temperature=0.7,
+            top_k=50,
+            top_p=0.95,
+            use_cache=True # Ускоряет повторную генерацию
+        )
+
+        # Запуск в отдельном потоке (Thread)
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        # Вывод текста по мере появления
+        def stream_output():
+            full_response = ""
+            for new_text in streamer:
+                full_response += new_text
+                yield new_text
+            st.session_state.messages.append({"role": "assistant", "content": full_response})
+
+        st.write_stream(stream_output)
+