Update app.py

app.py

@@ -1,24 +1,34 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
 from threading import Thread
 import gradio as gr
 
 model_path = "."
 
-# 1. Загрузка модели и токенизатора
-print("🍌 Загрузка BananaGPT...")
+# 1. Загрузка модели (без квантования для стабильной скорости на vCPU)
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-model = AutoModelForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True)
-model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
-print("✅ Готово!")
+model = AutoModelForCausalLM.from_pretrained(
+    model_path, 
+    low_cpu_mem_usage=True, 
+    use_cache=True,
+    device_map="cpu"
+)
+model.eval()
 
-def predict(message, history, temperature, top_p, rep_penalty, max_tokens):
-    # Собираем промпт
+# 2. Фильтр-стоппер на "Юзер" (чтобы не роллила за тебя)
+class StopOnUser(StoppingCriteria):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        stop_words = ["Юзер", "User"]
+        decoded_tail = tokenizer.decode(input_ids[0][-5:]) # Проверяем последние 5 токенов
+        return any(sw in decoded_tail for sw in stop_words)
+
+def predict(message, history, temperature, top_p, top_k, rep_penalty, no_repeat_ngram):
+    # Формируем контекст (последние 4 пары сообщений)
     prompt = ""
-    for user_msg, bot_msg in history:
+    for user_msg, bot_msg in history[-4:]:
         prompt += f"Юзер: {user_msg}\nБот: {bot_msg}\n"
     prompt += f"Юзер: {message}\nБот:"
 
@@ -26,50 +36,52 @@ def predict(message, history, temperature, top_p, rep_penalty, max_tokens):
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
     generation_kwargs = dict(
-        **inputs,
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
         streamer=streamer,
-        max_new_tokens=max_tokens,
+        max_new_tokens=512,
         do_sample=True,
         top_p=top_p,
+        top_k=int(top_k), # Твой фильтр на 70
         temperature=temperature,
         repetition_penalty=rep_penalty,
-        eos_token_id=tokenizer.eos_token_id,
+        no_repeat_ngram_size=int(no_repeat_ngram), # Против шизы и циклов
+        stopping_criteria=StoppingCriteriaList([StopOnUser()]),
         pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
     )
 
-    # Запуск в отдельном потоке
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
+    # Запуск генерации в потоке
+    with torch.inference_mode():
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
 
-    generated_text = ""
-    for new_text in streamer:
-        generated_text += new_text
-        # Наш стоп-кран для ролевика внутри
-        if "Юзер:" in generated_text:
-            generated_text = generated_text.split("Юзер:")[0].strip()
+        generated_text = ""
+        for new_text in streamer:
+            generated_text += new_text
+            
+            # Если проскочил "Юзер", обрезаем и выходим
+            if "Юзер:" in generated_text or "User:" in generated_text:
+                for stop in ["Юзер:", "User:"]:
+                    generated_text = generated_text.split(stop)[0]
+                yield generated_text.strip()
+                break
+                
             yield generated_text
-            break
-        yield generated_text
 
-# Настройка интерфейса с ползунками
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="slate")) as demo:
-    gr.Markdown("# 🍌 BananaGPT Lab")
-    
-    with gr.Row():
-        with gr.Column(scale=4):
-            # Чат интерфейс
-            chat = gr.ChatInterface(
-                fn=predict,
-                additional_inputs=[
-                    gr.Slider(0.1, 1.5, value=0.34, label="Температура (Хаос)"),
-                    gr.Slider(0.1, 1.0, value=0.9, label="Top-p (Плотность)"),
-                    gr.Slider(1.0, 2.0, value=1.3, label="Repetition Penalty"),
-                    gr.Slider(16, 1024, value=512, step=16, label="Макс. токенов"),
-                ]
-            )
-            
-    gr.Markdown("---")
-    gr.Markdown("ℹ️ *Если бот начал ролить за тебя, просто уменьши температуру или нажми Очистить.*")
+# 3. Интерфейс
+with gr.Blocks() as demo:
+    gr.Markdown("## 🍌 BananaGPT: Режим Анти-Шиза")
+    chat = gr.ChatInterface(
+        fn=predict,
+        additional_inputs=[
+            gr.Slider(0.1, 1.0, value=0.34, label="Temperature"),
+            gr.Slider(0.1, 1.0, value=0.9, label="Top-P"),
+            gr.Slider(1, 100, value=70, step=1, label="Top-K (Отсечка мусора)"),
+            gr.Slider(1.0, 2.0, value=1.2, label="Repetition Penalty"),
+            gr.Slider(0, 10, value=3, step=1, label="No Repeat N-Gram (Запрет циклов)"),
+        ]
+    )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)