gemma-coder-api / app.py
ryan33121's picture
Update app.py
188bb6e verified
Raw
History Blame Contribute Delete
1.83 kB
import json
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
from pydantic import BaseModel
from typing import List
# Initialize FastAPI
app = FastAPI()
# Enable CORS for frontend communication
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
# Load the model
# We use n_ctx=512 to keep memory footprint low
print("Loading model... please wait.")
MODEL_PATH = hf_hub_download(
repo_id="Qwen/Qwen2.5-Coder-1.5B-Instruct-GGUF",
filename="qwen2.5-coder-1.5b-instruct-q4_k_m.gguf"
)
# Force CPU mode (n_gpu_layers=0) for maximum stability on Free Tier
llm = Llama(
model_path=MODEL_PATH,
n_ctx=512,
n_gpu_layers=0,
n_threads=1
)
print("Model loaded successfully.")
# Pydantic models for request validation
class ChatMessage(BaseModel):
role: str
content: str
class ChatRequest(BaseModel):
messages: List[ChatMessage]
# Endpoints
@app.get("/")
def read_root():
return {"message": "Server is up and waiting for chat!"}
@app.post("/v1/chat/completions")
async def chat_completion(request: ChatRequest):
print("Received chat request...")
try:
# Convert request to history format
history = [{"role": msg.role, "content": msg.content} for msg in request.messages]
# Run inference with stream=False to ensure stable connection
response = llm.create_chat_completion(
messages=history,
temperature=0.2,
stream=False
)
print("Inference finished.")
return response
except Exception as e:
print(f"ERROR: {e}")
raise HTTPException(status_code=500, detail=str(e))