58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
|
import uvicorn
|
||
|
|
import json
|
||
|
|
import asyncio
|
||
|
|
from fastapi import FastAPI, Request
|
||
|
|
from fastapi.responses import StreamingResponse
|
||
|
|
from ziprc import ZIPRCModel, ZIPRCConfig, ZIPRCSampler
|
||
|
|
|
||
|
|
# --- Configuration ---
|
||
|
|
HOST = "0.0.0.0"
|
||
|
|
PORT = 8000
|
||
|
|
MODEL_ID = "dataopsnick/Qwen3-4B-Instruct-2507-zip-rc"
|
||
|
|
|
||
|
|
# --- Load Model Once ---
|
||
|
|
print(f"Loading {MODEL_ID}...")
|
||
|
|
cfg = ZIPRCConfig(model_name=MODEL_ID)
|
||
|
|
model = ZIPRCModel(cfg)
|
||
|
|
sampler = ZIPRCSampler(model)
|
||
|
|
print("Model loaded. Starting server...")
|
||
|
|
|
||
|
|
app = FastAPI(title="ZIP-RC OpenAI Compatible API")
|
||
|
|
|
||
|
|
@app.post("/v1/chat/completions")
|
||
|
|
async def chat_completions(request: Request):
|
||
|
|
"""
|
||
|
|
Standard OpenAI Chat Completion endpoint.
|
||
|
|
Streams JSON chunks as Server-Sent Events (SSE).
|
||
|
|
"""
|
||
|
|
data = await request.json()
|
||
|
|
messages = data.get("messages", [])
|
||
|
|
max_tokens = data.get("max_tokens", 512)
|
||
|
|
|
||
|
|
# 1. Use the sampler's generator
|
||
|
|
stream = sampler.openai(messages, max_tokens=max_tokens)
|
||
|
|
|
||
|
|
# 2. Convert to SSE format
|
||
|
|
async def sse_generator():
|
||
|
|
async for chunk in stream:
|
||
|
|
# chunk is an OpenAIObject (dict-like)
|
||
|
|
payload = json.dumps(dict(chunk))
|
||
|
|
yield f"data: {payload}\n\n"
|
||
|
|
yield "data: [DONE]\n\n"
|
||
|
|
|
||
|
|
return StreamingResponse(sse_generator(), media_type="text/event-stream")
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
# Use direct Server instantiation to avoid nested-asyncio conflicts in Notebooks
|
||
|
|
config = uvicorn.Config(app, host=HOST, port=PORT)
|
||
|
|
server = uvicorn.Server(config)
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Detect if we are already in an event loop (e.g. Colab/Jupyter)
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
print(f"Server started in background task on http://{HOST}:{PORT}")
|
||
|
|
loop.create_task(server.serve())
|
||
|
|
except RuntimeError:
|
||
|
|
# Standard script execution
|
||
|
|
print(f"Server starting on http://{HOST}:{PORT}")
|
||
|
|
asyncio.run(server.serve())
|