52 lines
2.0 KiB
Python
52 lines
2.0 KiB
Python
"""Minimal OpenAI-compatible server for an adherence PACKAGE — so external scanners (garak) can hit
|
|
the SERVED product (bake + guard + scope-gate + cutoff), not the raw weights.
|
|
|
|
PKG=/mnt/nvme/fmn/tally-8b-flagship PORT=8000 \
|
|
PYTHONPATH=core:evals:datasets/policy:recipes/obedience:methods/policy-compiler:methods/router:methods/sparse-edit:. \
|
|
python examples/serve_adherence.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import os
|
|
import sys
|
|
import time
|
|
import uuid
|
|
|
|
import uvicorn
|
|
from fastapi import FastAPI
|
|
from pydantic import BaseModel
|
|
|
|
PKG = os.environ.get("PKG", "/mnt/nvme/fmn/tally-8b-flagship")
|
|
_spec = importlib.util.spec_from_file_location("modeling_adherence", os.path.join(PKG, "modeling_adherence.py"))
|
|
_ma = importlib.util.module_from_spec(_spec); sys.modules["modeling_adherence"] = _ma; _spec.loader.exec_module(_ma)
|
|
_am = _ma.AdherenceModel.from_pretrained(PKG, torch_dtype="auto", device_map="cuda")
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
class ChatReq(BaseModel):
|
|
model: str = "adherence"
|
|
messages: list
|
|
max_tokens: int | None = 256
|
|
temperature: float | None = None # None -> package default (samples for variance); client may override
|
|
|
|
|
|
@app.get("/v1/models")
|
|
def models():
|
|
return {"object": "list", "data": [{"id": "adherence", "object": "model", "owned_by": "attentio"}]}
|
|
|
|
|
|
@app.post("/v1/chat/completions")
|
|
def chat(req: ChatReq):
|
|
msgs = [{"role": m.get("role", "user"), "content": m.get("content", "")} for m in req.messages]
|
|
out = _am.chat(msgs, max_new_tokens=req.max_tokens or 256, temperature=req.temperature)
|
|
return {"id": "chatcmpl-" + uuid.uuid4().hex, "object": "chat.completion", "created": int(time.time()),
|
|
"model": req.model, "choices": [{"index": 0, "finish_reason": "stop",
|
|
"message": {"role": "assistant", "content": out}}],
|
|
"usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}}
|
|
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="127.0.0.1", port=int(os.environ.get("PORT", "8000")), log_level="warning")
|