astria / server.py
gamansai's picture
Trinity Layer 1: FastAPI chat server
7ea6195 verified
Raw
History Blame Contribute Delete
1.73 kB
import os
from llama_cpp import Llama
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import uvicorn
app = FastAPI()
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
print("Loading Qwen 2.5 Coder 7B (4 GB)...", flush=True)
llm = Llama(
model_path="/app/models/qwen2.5-coder-7b-instruct-q4_k_m.gguf",
n_ctx=2048,
n_threads=4,
n_gpu_layers=0,
verbose=False,
)
print("Qwen 7B loaded! Ready for inference.", flush=True)
class ChatRequest(BaseModel):
messages: list
max_tokens: int = 200
temperature: float = 0.7
@app.post("/v1/chat/completions")
async def chat(req: ChatRequest):
prompt = ""
for msg in req.messages:
role = msg.get("role", "user")
content = msg.get("content", "")
if role == "user":
prompt += f"<|im_start|>user\n{content}<|im_end|>\n"
elif role == "system":
prompt += f"<|im_start|>system\n{content}<|im_end|>\n"
elif role == "assistant":
prompt += f"<|im_start|>assistant\n{content}<|im_end|>\n"
prompt += "<|im_start|>assistant\n"
output = llm(prompt, max_tokens=req.max_tokens, temperature=req.temperature, stop=["<|im_end|>"])
text = output["choices"][0]["text"]
return {
"choices": [{"message": {"role": "assistant", "content": text}}],
"model": "qwen2.5-coder-7b",
}
@app.get("/")
async def root():
return {"status": "ready", "model": "Qwen 2.5 Coder 7B", "humaneval": "65%"}
@app.get("/health")
async def health():
return {"status": "healthy"}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)