Update app.py

app.py

@@ -1,7 +1,7 @@
 from fastapi import FastAPI
 from fastapi.responses import PlainTextResponse
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import torch
 import uvicorn
 import threading
@@ -18,60 +18,68 @@ model.eval()
 
 # 🔹 настройки
 MAX_HISTORY = 40
-NUM_WORKERS = 3  # 🔥 ВАЖНО: количество потоков
+NUM_WORKERS = 3
 
-db = OrderedDict()
+db = OrderedDict()   # message -> {status, reply}
 queue = []
 
 class Message(BaseModel):
     message: str
 
 
-def generate_ai(message: str):
+# 🔥 STREAMING GENERATION
+def generate_ai_stream(message: str):
     prompt = f"User: {message}\nAssistant: Answer clearly and fully:\n"
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=180,
-            min_new_tokens=20,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
-            eos_token_id=tokenizer.eos_token_id
-        )
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
 
-    input_length = inputs.input_ids.shape[1]
-    generated_tokens = outputs[0][input_length:]
+    gen_kwargs = dict(
+        **inputs,
+        max_new_tokens=300,
+        min_new_tokens=30,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        streamer=streamer,
+        eos_token_id=tokenizer.eos_token_id
+    )
 
-    reply = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
+    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
+    thread.start()
 
-    return reply
+    partial = ""
 
+    for text in streamer:
+        partial += text
 
-# 🔥 Поток-воркер
+        # 🔥 обновляем ответ в реальном времени
+        if message in db:
+            db[message]["reply"] = partial
+
+    return partial.strip()
+
+
+# 🔥 WORKER
 def worker():
     while True:
         if queue:
             message = queue.pop(0)
 
-            # ⚡ если уже есть ответ — пропускаем
             if message in db and db[message]["status"] == "done":
                 continue
 
-            reply = generate_ai(message)
+            reply = generate_ai_stream(message)
 
             if message in db:
                 db[message]["status"] = "done"
                 db[message]["reply"] = reply
-
         else:
-            time.sleep(0.01)  # 🔥 меньше лаг
+            time.sleep(0.01)
 
 
-# 🔥 запускаем несколько воркеров
+# 🔥 запускаем 3 воркера (ускорение x2-x3)
 for _ in range(NUM_WORKERS):
     threading.Thread(target=worker, daemon=True).start()
 
@@ -81,16 +89,19 @@ async def root():
     return PlainTextResponse("AI server работает")
 
 
-# 🔹 отправка запроса
+# 🔹 ASK
 @app.get("/ask")
 async def ask(message: str):
 
-    # ⚡ МГНОВЕННЫЙ КЭШ
+    # кеш
     if message in db and db[message]["status"] == "done":
         return PlainTextResponse("cached")
 
     if message not in db:
-        db[message] = {"status": "pending", "reply": ""}
+        db[message] = {
+            "status": "pending",
+            "reply": ""
+        }
         queue.append(message)
 
         if len(db) > MAX_HISTORY:
@@ -99,7 +110,7 @@ async def ask(message: str):
     return PlainTextResponse("accepted")
 
 
-# 🔹 получение ответа
+# 🔹 GET (визуальный стриминг)
 @app.get("/get")
 async def get(message: str):
 
@@ -109,7 +120,7 @@ async def get(message: str):
     data = db[message]
 
     if data["status"] == "pending":
-        return PlainTextResponse("processing")
+        return PlainTextResponse(data["reply"] or "thinking...")
 
     return PlainTextResponse(data["reply"])