File size: 2,129 Bytes
7d03f2d | 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 | """
OpenAI-compatible shim for mailguard-jev-style-1.5b.
Exposes /v1/chat/completions so any OpenAI client can use the model locally.
Usage:
pip install fastapi uvicorn transformers torch
python serve_openai_shim.py
"""
import json
import re
import time
import uuid
import torch
import uvicorn
from fastapi import FastAPI
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "JackKozmo29/mailguard-jev-style-1.5b"
print(f"Loading {MODEL_ID} ...")
tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32).eval()
print("Model ready.")
app = FastAPI()
class Message(BaseModel):
role: str
content: str
class ChatRequest(BaseModel):
model: str = MODEL_ID
messages: list[Message]
max_tokens: int = 120
temperature: float = 0.0
@app.post("/v1/chat/completions")
def chat(req: ChatRequest):
messages = [m.model_dump() for m in req.messages]
prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
ids = tok(prompt, return_tensors="pt")
with torch.no_grad():
out = model.generate(
**ids,
max_new_tokens=req.max_tokens,
do_sample=False,
pad_token_id=tok.eos_token_id,
)
text = tok.decode(out[0][ids["input_ids"].shape[1]:], skip_special_tokens=True)
m = re.search(r"\{.*\}", text, re.DOTALL)
content = m.group(0) if m else text.strip()
return JSONResponse({
"id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
"object": "chat.completion",
"created": int(time.time()),
"model": req.model,
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": content},
"finish_reason": "stop",
}],
"usage": {"prompt_tokens": ids["input_ids"].shape[1], "completion_tokens": len(out[0]) - ids["input_ids"].shape[1], "total_tokens": len(out[0])},
})
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)
|