davanstrien HF Staff
Lead with the Jev / System One decision API; add jev_api.py server (tested)
c1ac1fc verified | # /// script | |
| # requires-python = ">=3.11" | |
| # dependencies = ["gliclass>=0.1.21", "torch", "transformers>=5.20", "fastapi", "uvicorn"] | |
| # /// | |
| """Jev / System One decision API for davanstrien/ultrabert-380m-curation. | |
| A request is a state (text, list or JSON record) plus named, typed questions; the response has one answer per | |
| question and no generated text. Question types: `choice`, `noul` (yes / no), `score` (ordered levels), and `choice` | |
| with `"multiple": True` (several options can be true). Each question is one GLiClass pass: the state is the text, | |
| `instructions` the prompt, each option a label. Inputs longer than MAX_LENGTH tokens are refused, never truncated. | |
| Serve it: | |
| uv run jev_api.py # or: uvicorn jev_api:app --host 0.0.0.0 --port 8765 | |
| curl -s localhost:8765/v1/systemone -H 'Content-Type: application/json' -d '{"state": "...", "questions": {...}}' | |
| or call it in Python: | |
| from jev_api import system_one | |
| system_one("page text", {"spam": {"type": "noul", "instructions": "Is this page spam?"}}) | |
| For large batches on a GPU, fast_inference.py (same mapping, packed and compiled) is about 3x faster. | |
| """ | |
| import json | |
| import os | |
| import torch | |
| from fastapi import FastAPI, HTTPException | |
| from gliclass import GLiClassModel, ZeroShotClassificationPipeline | |
| from pydantic import BaseModel | |
| from transformers import AutoTokenizer | |
| MODEL_ID = os.environ.get("MODEL", "davanstrien/ultrabert-380m-curation") | |
| MAX_LENGTH = int(os.environ.get("MAX_LENGTH", "4096")) # the training length; longer inputs are refused | |
| DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu" | |
| DTYPE = torch.bfloat16 if DEVICE != "cpu" else torch.float32 | |
| model = GLiClassModel.from_pretrained(MODEL_ID, dtype=DTYPE).to(DEVICE).eval() | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) | |
| _pipe = ZeroShotClassificationPipeline(model, tokenizer, max_length=10**9, classification_type="single-label", | |
| device=DEVICE, progress_bar=False) | |
| pipe = getattr(_pipe, "pipe", _pipe) | |
| app = FastAPI(title="System One API") | |
| class Request(BaseModel): | |
| state: str | dict | list | None = None | |
| questions: dict[str, dict] | |
| class RequestError(ValueError): | |
| """A request the model cannot answer (reported as HTTP 422).""" | |
| def as_text(value): | |
| return value if isinstance(value, str) else json.dumps(value, ensure_ascii=False) | |
| def render_state(state): | |
| """A state as the text the model reads: records become `key: value` lines, string lists `- item` lines.""" | |
| if state in (None, "", {}, []): | |
| return "" | |
| if isinstance(state, str): | |
| return state | |
| if isinstance(state, list): | |
| return "\n".join(f"- {as_text(item)}" for item in state) | |
| lines = [] | |
| for key, value in state.items(): | |
| if isinstance(value, list) and all(isinstance(item, str) for item in value): | |
| lines += [f"{key}:"] + [f"- {item}" for item in value] | |
| else: | |
| lines.append(f"{key}: {as_text(value)}") | |
| return "\n".join(lines) | |
| def question_options(question): | |
| """A question -> (answer keys, labels), as in training.""" | |
| kind, criteria = question.get("type"), question.get("criteria") | |
| if kind == "choice": | |
| if not criteria or len(criteria) < 2: | |
| raise RequestError("a choice question needs at least two criteria") | |
| if isinstance(criteria, dict): | |
| return list(criteria), [as_text(v) if v is not None else k for k, v in criteria.items()] | |
| return [str(c) for c in criteria], [str(c) for c in criteria] | |
| if kind == "noul": | |
| c = criteria if isinstance(criteria, dict) else {} | |
| return ["true", "false"], [as_text(c.get("true") or "yes"), as_text(c.get("false") or "no")] | |
| if kind == "score": | |
| if not criteria or len(criteria) < 2: | |
| raise RequestError("a score question needs at least two levels") | |
| return [str(i) for i in range(len(criteria))], [as_text(c) for c in criteria] | |
| raise RequestError(f"unsupported question type {kind!r}: this model answers choice, noul and score questions") | |
| def option_logits(text, prompt, labels): | |
| if len(set(labels)) != len(labels): | |
| raise RequestError("option descriptions must be unique") | |
| inputs = pipe.prepare_inputs([text], labels, same_labels=True, prompt=prompt) | |
| length = inputs["input_ids"].shape[-1] | |
| if length > MAX_LENGTH: | |
| raise RequestError(f"input of {length} tokens exceeds the {MAX_LENGTH}-token limit") | |
| inputs = {k: v.to(DEVICE) for k, v in inputs.items()} | |
| logits = model(**inputs, max_num_classes=len(labels)).logits[0, : len(labels)].float() | |
| return logits, length | |
| def answer(question, keys, labels, logits): | |
| kind = question["type"] | |
| if question.get("multiple"): # trained as an independent yes / no per option | |
| p = dict(zip(keys, torch.sigmoid(logits).tolist(), strict=True)) | |
| threshold = question.get("threshold", 0.5) | |
| return {"type": kind, "choices": [k for k in keys if p[k] >= threshold], "probabilities": p} | |
| p = dict(zip(keys, torch.softmax(logits, dim=-1).tolist(), strict=True)) | |
| best = max(keys, key=p.get) | |
| if kind == "noul": | |
| return {"type": kind, "noul": p["true"], "probabilities": p} | |
| if kind == "score": | |
| return {"type": kind, "score": sum(int(k) * v for k, v in p.items()), "confidence": p[best], | |
| "probabilities": p} | |
| return {"type": kind, "choice": best, "confidence": p[best], "probabilities": p} | |
| def system_one(state, questions): | |
| """Answer every question about one state. Raises RequestError for a request the model cannot answer.""" | |
| text, answers, tokens = render_state(state), {}, 0 | |
| for name, question in questions.items(): | |
| keys, labels = question_options(question) | |
| prompt = as_text(question.get("instructions") or "") or None | |
| logits, length = option_logits(text, prompt, labels) | |
| tokens += length | |
| answers[name] = answer(question, keys, labels, logits) | |
| return {"model": MODEL_ID, "answers": answers, "usage": {"input_tokens": tokens, "output_tokens": 0}} | |
| def models(): | |
| return {"models": [{"name": MODEL_ID, "description": "GLiClass curation decision model (UltraBERT-380M)"}]} | |
| def systemone(request: Request): | |
| try: | |
| return system_one(request.state, request.questions) | |
| except RequestError as e: | |
| raise HTTPException(422, str(e)) from e | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8765"))) | |