File size: 2,445 Bytes
5e61dc8 700f6af 5e61dc8 4e397a7 700f6af 88fca26 005d065 3cdaf44 4e397a7 5e61dc8 88fca26 700f6af 3cdaf44 700f6af 5e61dc8 4e397a7 5e61dc8 0be9221 700f6af 4e397a7 0be9221 4e397a7 5eda153 3cdaf44 88fca26 5eda153 5e61dc8 5eda153 4e397a7 005d065 5eda153 4e397a7 51bd9ab 4e397a7 | 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 | import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from starlette.routing import Route
import gradio as gr
MODEL_NAME = "teknium/OpenHermes-2.5-Mistral-7B"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.float16, device_map="cpu")
model.eval()
@spaces.GPU(duration=120)
def generate_reply(prompt: str, max_tokens: int = 512, temperature: float = 0.7) -> str:
model.to("cuda")
try:
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
with torch.inference_mode():
output = model.generate(
**inputs,
max_new_tokens=max_tokens,
do_sample=True,
temperature=temperature,
pad_token_id=tokenizer.eos_token_id,
)
return tokenizer.decode(output[0], skip_special_tokens=True)
finally:
model.to("cpu")
torch.cuda.empty_cache()
with gr.Blocks() as demo:
gr.Markdown("API is running. POST to `/v1/chat/completions`.")
# --- Register the custom route FIRST ---
async def chat_completions(request: Request):
payload = await request.json()
messages = payload.get("messages", [])
temperature = float(payload.get("temperature", 0.7))
max_tokens = int(payload.get("max_tokens", 512))
prompt = ""
for msg in messages:
role = msg.get("role", "user")
content = msg.get("content", "")
prompt += f"{role}: {content}\n"
prompt += "assistant:"
reply = generate_reply(prompt, max_tokens=max_tokens, temperature=temperature)
return JSONResponse({
"id": "openhermes-2.5-mistral-7b-response",
"object": "chat.completion",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": reply},
"finish_reason": "stop",
}],
})
# Insert at position 0 so it runs before Gradio's catch-all
demo.app.router.routes.insert(0, Route("/v1/chat/completions", chat_completions, methods=["POST"]))
# --- CORS middleware ---
demo.app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False) |