api / app.py
wowimdaniel's picture
Update app.py
510bb60 verified
Raw History Blame Contribute Delete
2.47 kB
from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
from llama_cpp import Llama
import uuid
app = FastAPI()
# Загружаем Llama-3-8B (сжатую до Q4_K_M — это ~5ГБ)
# Она влезет в 16ГБ ОЗУ и будет на порядок умнее GPT-2 и TinyLlama
llm = Llama.from_pretrained(
repo_id="MaziyarPanahi/Llama-3-8B-Instruct-v0.1-GGUF",
filename="*Q4_K_M.gguf",
n_ctx=2048, # Контекст (память)
n_threads=2 # Используем оба ядра твоего бесплатного CPU
)
tasks = {}
class GenerateRequest(BaseModel):
prompt: str
system_prompt: str = "Ты — опытный ИИ-помощник. Отвечай четко и только на русском языке."
custom_id: str = None
max_tokens: int = 256
temperature: float = 0.7
def run_ai(task_id, prompt, system_prompt, max_tokens, temp):
try:
# Форматируем запрос под стандарт Llama 3
full_prompt = f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
output = llm(
full_prompt,
max_tokens=max_tokens,
temperature=temp,
stop=["<|eot_id|>", "<|end_of_text|>"],
echo=False
)
tasks[task_id]["result"] = output["choices"][0]["text"].strip()
tasks[task_id]["status"] = "ready"
except Exception as e:
tasks[task_id] = {"status": "error", "result": str(e)}
@app.post("/generate")
async def generate(req: GenerateRequest, background_tasks: BackgroundTasks):
task_id = req.custom_id if req.custom_id else str(uuid.uuid4())
tasks[task_id] = {"status": "pending", "result": ""}
background_tasks.add_task(
run_ai,
task_id,
req.prompt,
req.system_prompt,
req.max_tokens,
req.temperature
)
return {"task_id": task_id}
@app.get("/check/{task_id}")
async def check(task_id: str):
data = tasks.get(task_id)
if not data:
return {"task_id": task_id, "status": "not_found", "result": None}
return {"task_id": task_id, "status": data["status"], "result": data["result"]}
@app.get("/")
async def root():
return {"status": "running", "model": "Llama-3-8B-Instruct-GGUF"}