Spaces:
Running
Running
| 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 | |
| def read_root(): | |
| return {"message": "Server is up and waiting for chat!"} | |
| 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)) |