Update app.py

app.py

@@ -1,21 +1,36 @@
 from fastapi import FastAPI
 from fastapi.responses import PlainTextResponse
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from llama_cpp import Llama
 from deep_translator import GoogleTranslator
-import torch
 import uvicorn
 import threading
 import time
 from collections import OrderedDict
+import os
+import requests
 
 app = FastAPI()
 
-model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-model.eval()
+# 🔥 MODEL DOWNLOAD
+MODEL_URL = "https://huggingface.co/TheBloke/phi-2-GGUF/resolve/main/phi-2.Q4_K_M.gguf"
+MODEL_PATH = "phi-2.gguf"
+
+if not os.path.exists(MODEL_PATH):
+    print("Downloading model... (это может занять время)")
+    with requests.get(MODEL_URL, stream=True) as r:
+        with open(MODEL_PATH, "wb") as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                if chunk:
+                    f.write(chunk)
+    print("Model downloaded!")
+
+# 🔥 LOAD MODEL
+llm = Llama(
+    model_path=MODEL_PATH,
+    n_ctx=2048,
+    n_threads=4
+)
 
 MAX_HISTORY = 40
 NUM_WORKERS = 3
@@ -27,47 +42,36 @@ class Message(BaseModel):
     message: str
 
 
-# 🔥 РАЗБИВКА ТЕКСТА
-def split_text(text, max_len=80):
+# 🔥 split text
+def split_text(text, max_len=25):
     return "\n".join([text[i:i+max_len] for i in range(0, len(text), max_len)])
 
 
-# 🔥 СТРИМИНГ ГЕНЕРАЦИИ
+# 🔥 GENERATION
 def generate_ai_stream(message: str):
     prompt = f"User: {message}\nAssistant: Answer clearly and fully:\n"
 
-    inputs = tokenizer(prompt, return_tensors="pt")
-
-    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
-
-    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
+    output = llm(
+        prompt,
+        max_tokens=300,
+        stop=["User:", "\n\n"],
+        stream=True
     )
 
-    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
-    thread.start()
-
     partial = ""
 
-    # 🔥 ПОТОКОВАЯ ГЕНЕРАЦИЯ
-    for text in streamer:
-        partial += text
+    for chunk in output:
+        token = chunk["choices"][0]["text"]
+        partial += token
 
         if message in db:
             db[message]["reply"] = split_text(partial)
 
-    # 🔥 ПЕРЕВОД В КОНЦЕ
+    # 🔥 перевод
     try:
         translated = GoogleTranslator(source='en', target='ru').translate(partial.strip())
     except:
-        translated = partial.strip()  # fallback если перевод упал
+        translated = partial.strip()
 
     final_text = split_text(translated) + " full generated"
 
@@ -77,7 +81,7 @@ def generate_ai_stream(message: str):
     return final_text
 
 
-# 🔥 ВОРКЕР
+# 🔥 WORKER
 def worker():
     while True:
         if queue:
@@ -95,17 +99,17 @@ def worker():
             time.sleep(0.01)
 
 
-# 🔥 ЗАПУСК НЕСКОЛЬКИХ ПОТОКОВ
+# 🔥 START WORKERS
 for _ in range(NUM_WORKERS):
     threading.Thread(target=worker, daemon=True).start()
 
 
 @app.get("/")
 async def root():
-    return PlainTextResponse("AI server работает")
+    return PlainTextResponse("AI server running (Phi-2 GGUF + llama.cpp)")
 
 
-# 🔹 ASK
+# 🔥 ASK
 @app.get("/ask")
 async def ask(message: str):
 
@@ -122,7 +126,7 @@ async def ask(message: str):
     return PlainTextResponse("accepted")
 
 
-# 🔹 GET (стриминг)
+# 🔥 GET
 @app.get("/get")
 async def get(message: str):