from fastapi import FastAPI
from fastapi.responses import StreamingResponse, HTMLResponse
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import torch
import threading
app = FastAPI()
MODEL_NAME = "cygnisai/Cygnis-Alpha-1.7B-v0.1-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map="cpu",
torch_dtype=torch.float32
)
class Request(BaseModel):
message: str
@app.get("/", response_class=HTMLResponse)
def home():
return """
Server Status
"""
@app.post("/chat-stream")
def chat_stream(req: Request):
inputs = tokenizer(req.message, return_tensors="pt")
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt=True,
skip_special_tokens=True
)
def generate():
model.generate(
**inputs,
max_new_tokens=150,
temperature=0.7,
do_sample=True,
streamer=streamer
)
thread = threading.Thread(target=generate)
thread.start()
def token_stream():
for token in streamer:
yield token
return StreamingResponse(token_stream(), media_type="text/plain")