94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
"""Closed-vocabulary classification on the local model pool.
|
|
|
|
Given a text and a list of allowed slugs, the model must answer with
|
|
exactly one slug or ``none``. Anything else is treated as ``None`` so a
|
|
caller can queue the text for human review instead of inventing a value.
|
|
Self-hosted inference only (Ollama ``/api/chat`` on the largest live
|
|
host, the same route ``llm.rag.stream_answer`` takes).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any, Callable, Sequence
|
|
|
|
import httpx
|
|
|
|
from llm.pool import pick_model
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
Classifier = Callable[[str, Sequence[str]], str | None]
|
|
|
|
_SYSTEM = (
|
|
"You label short healthcare-policy text with exactly one slug from a "
|
|
"closed list. Reply with the slug only — no words, no punctuation. If no "
|
|
"slug fits, reply: none"
|
|
)
|
|
_TIMEOUT = httpx.Timeout(120.0, connect=10.0)
|
|
|
|
|
|
def build_prompt(text: str, choices: Sequence[str]) -> list[dict]:
|
|
listing = "\n".join(f"- {c}" for c in choices)
|
|
return [
|
|
{"role": "system", "content": _SYSTEM},
|
|
{
|
|
"role": "user",
|
|
"content": f"Allowed slugs:\n{listing}\n- none\n\nText:\n{text.strip()}\n\nSlug:",
|
|
},
|
|
]
|
|
|
|
|
|
def parse_choice(reply: str, choices: Sequence[str]) -> str | None:
|
|
token = (reply or "").strip().strip("\"'`.").strip().lower()
|
|
if token == "none":
|
|
return None
|
|
return token if token in {c.lower() for c in choices} else None
|
|
|
|
|
|
def _default_post(url: str, json: dict) -> dict:
|
|
with httpx.Client(timeout=_TIMEOUT) as client:
|
|
resp = client.post(url, json=json)
|
|
resp.raise_for_status()
|
|
return resp.json()
|
|
|
|
|
|
def closed_vocab_classifier(
|
|
cfg: Any,
|
|
pool: Any,
|
|
*,
|
|
post: Callable[..., dict] | None = None,
|
|
model: str | None = None,
|
|
) -> Classifier:
|
|
"""Return ``classify(text, choices) -> slug | None`` bound to *pool*.
|
|
|
|
``pool.check`` runs once here, not per call (I5) — it's a host-liveness
|
|
probe, not per-request work, and a code with hundreds of element lines
|
|
used to re-probe the pool hundreds of times.
|
|
|
|
*model* pins every call to that model on every host (a large-VRAM
|
|
host no longer upgrades it), for callers whose recorded outputs name
|
|
the model; the default picks per host (``pick_model``)."""
|
|
send = post or _default_post
|
|
pool.check(model or cfg.instruct_model)
|
|
|
|
def classify(text: str, choices: Sequence[str]) -> str | None:
|
|
with pool.acquire_generation() as host:
|
|
chosen = model or pick_model(cfg, pool, host)
|
|
data = send(
|
|
f"{host}/api/chat",
|
|
json={
|
|
"model": chosen,
|
|
"messages": build_prompt(text, choices),
|
|
"stream": False,
|
|
"options": {"num_ctx": cfg.chat_num_ctx, "temperature": 0},
|
|
},
|
|
)
|
|
reply = data.get("message", {}).get("content", "")
|
|
choice = parse_choice(reply, choices)
|
|
if choice is None and reply.strip().lower() != "none":
|
|
log.info("classify: unparseable reply %r for %r", reply[:60], text[:60])
|
|
return choice
|
|
|
|
return classify
|