File size: 3,790 Bytes
e23537f cb192d4 e23537f | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """Hugging Face Inference Endpoints handler: a description (or an image, when a vision
model is reachable) in, ranked Vienna codes out.
Copied to the root of the model repository by ``i2vienna hf-export``. The repository
also carries ``scheme/`` and ``index/`` (parquet), ``image2vienna.json`` (which index
to serve) and a copy of the ``image2vienna`` package, so no PyPI release is required.
Request, text (the description of the figurative elements, stage two alone)::
{"inputs": "Three stars above a crescent moon. Gold and blue.",
"parameters": {"level": "section", "top_k": 5}}
Request, image (needs a describer the endpoint can reach, named by the environment
variable ``IMAGE2VIENNA_DESCRIBER``, e.g. ``hf:florence-community/Florence-2-base-ft``
to run Florence-2 on the endpoint itself; without the variable an image is refused)::
{"inputs": {"image": "<base64 bytes>"}, "parameters": {"level": "division"}}
Response (one list per input; a single input yields a single list)::
[{"code": "1.1.4", "padded": "01.01.04", "level": "section", "auxiliary": true,
"score": 0.83, "similarity": 0.83, "title": "Three stars", "path": "..."}, ...]
"""
from __future__ import annotations
import base64
import json
import os
import sys
from pathlib import Path
from typing import Any
ALLOWED_PARAMS = {
"level",
"top_k",
"gap",
"auto_margin",
"normalize",
"chunking",
"principal_only",
"exclude_codes",
}
class EndpointHandler:
def __init__(self, path: str = ""):
root = Path(path or ".").resolve()
if str(root) not in sys.path:
sys.path.insert(0, str(root)) # vendored image2vienna package
from image2vienna.classifier import ViennaClassifier
cfg = json.loads((root / "image2vienna.json").read_text())
self.clf = ViennaClassifier(
cfg["edition"],
lang=cfg["lang"],
model=cfg["model"],
notes=bool(cfg.get("notes", False)),
describer=os.environ.get("IMAGE2VIENNA_DESCRIBER") or "fixed:/nonexistent",
root=root,
)
self.has_describer = bool(os.environ.get("IMAGE2VIENNA_DESCRIBER"))
_ = self.clf.index # load eagerly so the first request is fast
_ = self.clf.embedder
def __call__(self, data: dict[str, Any]) -> list[dict] | list[list[dict]]:
inputs = data.get("inputs", "")
params = {k: v for k, v in (data.get("parameters") or {}).items() if k in ALLOWED_PARAMS}
if "exclude_codes" in params:
params["exclude_codes"] = tuple(params["exclude_codes"] or ())
if isinstance(inputs, list):
return [self._one(item, params) for item in inputs]
return self._one(inputs, params)
def _one(self, item: Any, params: dict[str, Any]) -> list[dict]:
if isinstance(item, dict) and "image" in item:
if not self.has_describer:
raise ValueError(
"This endpoint has no vision model: send the description of the figurative "
"elements as text, or run the image2vienna package locally for images"
)
matches = self.clf.classify(base64.b64decode(item["image"]), **params)
else:
text = item["text"] if isinstance(item, dict) else str(item)
matches = self.clf.classify_text(text, **params)
return [
{
"code": m.code,
"padded": m.padded,
"level": m.level,
"auxiliary": m.auxiliary,
"score": round(m.score, 4),
"similarity": round(m.similarity, 4),
"title": m.title,
"path": m.text,
}
for m in matches
]
|