Update app.py

app.py

@@ -1,99 +1,107 @@
 import os
 import torch
 import telebot
-from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
-from threading import Thread
-import gradio as gr
+import time
 import gc
+import gradio as gr
+from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
 
-# --- НАСТРОЙКИ ---
+# --- КОНФИГ ---
 MODEL_PATH = "."
-# Берем токены из секретов Space
 TG_TOKEN = os.environ.get("TG_BOT")
-HF_TOKEN = os.environ.get("HF_TOKEN")
+
+user_status = {}
 
 # --- ЗАГРУЗКА МОДЕЛИ ---
+print("Загружаю модель...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
-
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_PATH,
-    torch_dtype=torch.float32, # Для CPU float32 стабильнее
+    torch_dtype=torch.float32,
     low_cpu_mem_usage=True,
-    use_cache=True
+    device_map="cpu"
 )
-model.eval()
-
-bot = telebot.TeleBot(TG_TOKEN)
 
 # --- ФИЛЬТРЫ ---
 class StopOnUser(StoppingCriteria):
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        stop_words = ["Юзер", "User", "Пользователь", "Персонаж"]
-        text = tokenizer.decode(input_ids[0][-6:])
+    def __call__(self, input_ids, scores, **kwargs):
+        stop_words = ["User", "Юзер", "Пользователь", "Персонаж"]
+        text = tokenizer.decode(input_ids[0][-10:])
         return any(sw in text for sw in stop_words)
 
-def generate_answer(message_text):
-    gc.collect() # Чистим мусор перед генерацией
-    
-    # Формат Saiga (минималистичный)
-    prompt = f"<s>system\nТы — дерзкий ИИ с характером.</s>\n<s>user\n{message_text}</s>\n<s>bot\n"
-    
+def generate_answer(text):
+    gc.collect()
+    # Убрали системный промпт, оставили только структуру диалога
+    prompt = f"<s>user\n{text}</s>\n<s>bot\n"
     inputs = tokenizer(prompt, return_tensors="pt")
     
-    with torch.inference_mode():
+    with torch.no_grad():
         output_ids = model.generate(
             **inputs,
-            max_new_tokens=256,
+            max_new_tokens=150,
             do_sample=True,
-            temperature=0.5, # Чуть выше для "жизни"
-            top_p=0.9,
-            top_k=70,
+            temperature=0.4, # Установили 0.4 для большей точности
             repetition_penalty=1.2,
-            no_repeat_ngram_size=3,
-            stopping_criteria=StoppingCriteriaList([StopOnUser()]),
-            pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=tokenizer.eos_token_id
+            stopping_criteria=StoppingCriteriaList([StopOnUser()])
         )
     
-    # Декодируем только ответ бота
     full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     answer = full_text.split("<s>bot\n")[-1].strip()
     
-    # Финальная чистка хвостов
-    for stop in ["Юзер", "User", "Пользователь", "Персонаж"]:
-        answer = answer.split(stop)[0]
-    
+    for s in ["User", "Юзер", "Пользователь", "<s>"]:
+        answer = answer.split(s)[0]
     return answer.strip()
 
-# --- ЛОГИКА ТЕЛЕГРАМ ---
-@bot.message_handler(commands=['start'])
-def send_welcome(message):
-    bot.reply_to(message, "Я живая. Че хотел?")
+# --- ТЕЛЕГРАМ ---
+bot = telebot.TeleBot(TG_TOKEN, threaded=False)
+
+@bot.message_handler(func=lambda m: True)
+def chat(m):
+    user_id = m.from_user.id
+    current_time = time.time()
+    
+    if user_id not in user_status:
+        user_status[user_id] = {"last_time": 0, "is_thinking": False}
+
+    if user_status[user_id]["is_thinking"]:
+        bot.reply_to(m, "Модель думает. Не нагружайте сервер.")
+        return
+
+    time_passed = current_time - user_status[user_id]["last_time"]
+    if time_passed < 15:
+        wait_time = int(15 - time_passed)
+        bot.reply_to(m, f"Слишком быстро! Подожди еще {wait_time} сек.")
+        return
 
-@bot.message_handler(func=lambda message: True)
-def echo_all(message):
     try:
-        # Отправляем статус "печать", чтобы юзер не скучал
-        bot.send_chat_action(message.chat.id, 'typing')
-        response = generate_answer(message.text)
-        if not response:
-            response = "Мне нечего сказать на это..."
-        bot.reply_to(message, response)
+        user_status[user_id]["is_thinking"] = True
+        bot.send_chat_action(m.chat.id, 'typing')
+        
+        ans = generate_answer(m.text)
+        
+        bot.reply_to(m, ans if ans else "...")
+        user_status[user_id]["last_time"] = time.time()
     except Exception as e:
-        bot.reply_to(message, "У меня мозги заклинило, попробуй еще раз.")
+        print(f"Ошибка: {e}")
+        bot.reply_to(m, "Ошибка при обработке запроса.")
+    finally:
+        user_status[user_id]["is_thinking"] = False
 
-# --- ЗАПУСК ---
 def run_bot():
-    bot.infinity_polling()
+    print("Цикл опроса запущен...")
+    while True:
+        try:
+            bot.polling(none_stop=True, interval=2, timeout=25)
+        except Exception as e:
+            print(f"Сеть спит (ошибка: {e})")
+            time.sleep(10)
 
-# Запускаем бота в отдельном потоке
 Thread(target=run_bot, daemon=True).start()
 
-# Оставляем Gradio, чтобы Space жил
+# --- ИНТЕРФЕЙС ---
 with gr.Blocks() as demo:
-    gr.Markdown("## 🍌 BananaBot в Telegram запущен!")
-    gr.Markdown("Можешь писать в ТГ, а здесь просто висит панель управления.")
+    gr.Markdown("# 🍌 Banana Bot (Temp 0.4)")
+    gr.Markdown("Бот работает без системного промпта.")
 
 demo.launch(server_name="0.0.0.0", server_port=7860)