File size: 2,614 Bytes
dfdd6ea 71de84b dfdd6ea 71de84b dfdd6ea 71de84b dfdd6ea 7f717aa dfdd6ea 7f717aa dfdd6ea 71de84b dfdd6ea 71de84b dfdd6ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import os
import json
import asyncio
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import List
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
app = FastAPI(title="Qwen 0.5B Streaming API")
model_id = "Qwen/Qwen2.5-0.5B-Instruct"
# Initialize model directly onto CPU
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float32,
device_map="cpu"
)
class Message(BaseModel):
role: str
content: str
class ChatPayload(BaseModel):
messages: List[Message]
temperature: float = 0.7
max_tokens: int = 512
@app.post("/v1/chat/completions")
async def chat_completion(payload: ChatPayload):
# Convert incoming array to ChatML format
formatted_messages = [{"role": msg.role, "content": msg.content} for msg in payload.messages]
formatted_prompt = tokenizer.apply_chat_template(
formatted_messages,
tokenize=False,
add_generation_prompt=True
)
inputs = tokenizer([formatted_prompt], return_tensors="pt").to("cpu")
# Initialize the background streamer tool
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
generation_kwargs = dict(
inputs,
streamer=streamer,
max_new_tokens=payload.max_tokens,
do_sample=True,
temperature=payload.temperature,
pad_token_id=tokenizer.eos_token_id
)
# Run text generation in a separate thread to keep the main event loop unblocked
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
async def event_generator():
for new_text in streamer:
if new_text:
# Mirror standard OpenAI stream object format
chunk = {
"choices": [
{
"delta": {
"content": new_text
}
}
]
}
yield f"data: {json.dumps(chunk)}\n\n"
await asyncio.sleep(0.01) # Small yield breather for network loop
yield "data: [DONE]\n\n"
return StreamingResponse(event_generator(), media_type="text/event-stream")
@app.get("/")
async def health():
return {"status": "online"}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) |