paperus-ml 0.0.8__tar.gz

paperus_ml-0.0.8/PKG-INFO

@@ -0,0 +1,53 @@
+Metadata-Version: 2.4
+Name: paperus-ml
+Version: 0.0.8
+Summary: Machine learning service of paperus: text embeddings (later reranking, OCR)
+Keywords: embeddings,onnx,semantic search,self-hosted
+Author: paperus contributors
+License-Expression: Apache-2.0
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Requires-Dist: fastapi>=0.162.0
+Requires-Dist: huggingface-hub>=2.4.0
+Requires-Dist: numpy>=2.5.4
+Requires-Dist: onnxruntime>=1.31.1
+Requires-Dist: prometheus-client>=0.26.0
+Requires-Dist: pydantic-settings>=2.16.2
+Requires-Dist: tokenizers>=0.23.3
+Requires-Dist: uvicorn[standard]>=0.54.0
+Requires-Python: >=3.13
+Project-URL: Homepage, https://github.com/qcode-uni-a/paperus
+Project-URL: Source, https://github.com/qcode-uni-a/paperus
+Project-URL: Issues, https://github.com/qcode-uni-a/paperus/issues
+Description-Content-Type: text/markdown
+
+# paperus-ml
+
+Machine learning service of [paperus](https://github.com/qcode-uni-a/paperus):
+text embeddings for search by meaning, with ONNX Runtime on the CPU. Models
+are downloaded from Hugging Face on first use and unloaded when idle.
+
+```sh
+pip install paperus-ml
+PAPERUS_ML_TOKEN=… paperus-ml --host 0.0.0.0   # port 3003
+```
+
+Settings are environment variables:
+
+| Variable | Default | Meaning |
+| --- | --- | --- |
+| `PAPERUS_ML_TOKEN` | none | Shared secret; requests need `Authorization: Bearer <token>` |
+| `PAPERUS_ML_CACHE_PATH` | `~/.cache/paperus-ml` | Downloaded models |
+| `PAPERUS_ML_PRELOAD` | `[]` | Models to load at start and keep, e.g. `["arctic-embed-m-v2"]` (paperus then sends the one in use) |
+| `PAPERUS_ML_MODEL_TTL` | `300` | Unload other models after this many idle seconds (0: never) |
+| `PAPERUS_ML_INTRA_OP_THREADS` | `0` | CPU cores per model at start (0: all it sees); paperus' administration sets it afterwards |
+| `PAPERUS_ML_MAX_BATCH_TEXTS` / `_TOKENS` | `32` / `16384` | Batch size (memory grows with both) |
+
+paperus sends its settings with `PUT /config` (cores, models to keep
+loaded). One batch runs at a time per model, using all its cores; simultaneous
+requests take turns batch by batch.
+
+Most installations use the `paperus-ml` container image instead. Status:
+early version (0.0.x); expect breaking changes.

paperus_ml-0.0.8/README.md

@@ -0,0 +1,28 @@
+# paperus-ml
+
+Machine learning service of [paperus](https://github.com/qcode-uni-a/paperus):
+text embeddings for search by meaning, with ONNX Runtime on the CPU. Models
+are downloaded from Hugging Face on first use and unloaded when idle.
+
+```sh
+pip install paperus-ml
+PAPERUS_ML_TOKEN=… paperus-ml --host 0.0.0.0   # port 3003
+```
+
+Settings are environment variables:
+
+| Variable | Default | Meaning |
+| --- | --- | --- |
+| `PAPERUS_ML_TOKEN` | none | Shared secret; requests need `Authorization: Bearer <token>` |
+| `PAPERUS_ML_CACHE_PATH` | `~/.cache/paperus-ml` | Downloaded models |
+| `PAPERUS_ML_PRELOAD` | `[]` | Models to load at start and keep, e.g. `["arctic-embed-m-v2"]` (paperus then sends the one in use) |
+| `PAPERUS_ML_MODEL_TTL` | `300` | Unload other models after this many idle seconds (0: never) |
+| `PAPERUS_ML_INTRA_OP_THREADS` | `0` | CPU cores per model at start (0: all it sees); paperus' administration sets it afterwards |
+| `PAPERUS_ML_MAX_BATCH_TEXTS` / `_TOKENS` | `32` / `16384` | Batch size (memory grows with both) |
+
+paperus sends its settings with `PUT /config` (cores, models to keep
+loaded). One batch runs at a time per model, using all its cores; simultaneous
+requests take turns batch by batch.
+
+Most installations use the `paperus-ml` container image instead. Status:
+early version (0.0.x); expect breaking changes.

paperus_ml-0.0.8/pyproject.toml

@@ -0,0 +1,74 @@
+[project]
+name = "paperus-ml"
+version = "0.0.8"
+description = "Machine learning service of paperus: text embeddings (later reranking, OCR)"
+readme = "README.md"
+license = "Apache-2.0"
+keywords = [
+    "embeddings",
+    "onnx",
+    "semantic search",
+    "self-hosted",
+]
+classifiers = [
+    "Development Status :: 2 - Pre-Alpha",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering",
+]
+requires-python = ">=3.13"
+dependencies = [
+    "fastapi>=0.162.0",
+    "huggingface-hub>=2.4.0",
+    "numpy>=2.5.4",
+    "onnxruntime>=1.31.1",
+    "prometheus-client>=0.26.0",
+    "pydantic-settings>=2.16.2",
+    "tokenizers>=0.23.3",
+    "uvicorn[standard]>=0.54.0",
+]
+
+[[project.authors]]
+name = "paperus contributors"
+
+[project.urls]
+Homepage = "https://github.com/qcode-uni-a/paperus"
+Source = "https://github.com/qcode-uni-a/paperus"
+Issues = "https://github.com/qcode-uni-a/paperus/issues"
+
+[project.scripts]
+paperus-ml = "paperus_ml.cli:main"
+
+[build-system]
+requires = ["uv_build>=0.13,<0.14"]
+build-backend = "uv_build"
+
+[tool.ruff]
+line-length = 88
+
+[tool.ruff.lint]
+select = [
+    "E",
+    "F",
+    "W",
+    "I",
+    "B",
+    "UP",
+    "SIM",
+    "RUF",
+    "ANN",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+filterwarnings = ["error"]
+markers = ["download: downloads a real model (set PAPERUS_ML_TEST_DOWNLOAD=1)"]
+
+[dependency-groups]
+dev = [
+    "httpx2>=2.16.0",
+    "onnx>=1.23.2",
+    "pytest>=9.2.0",
+    "ruff>=0.17.3",
+    "ty>=0.0.97",
+]

paperus_ml-0.0.8/pyproject.toml.orig

@@ -0,0 +1,57 @@
+[project]
+name = "paperus-ml"
+version = "0.0.8"
+description = "Machine learning service of paperus: text embeddings (later reranking, OCR)"
+readme = "README.md"
+license = "Apache-2.0"
+authors = [{ name = "paperus contributors" }]
+keywords = ["embeddings", "onnx", "semantic search", "self-hosted"]
+classifiers = [
+    "Development Status :: 2 - Pre-Alpha",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering",
+]
+requires-python = ">=3.13"
+dependencies = [
+    "fastapi>=0.162.0",
+    "huggingface-hub>=2.4.0",
+    "numpy>=2.5.4",
+    "onnxruntime>=1.31.1",
+    "prometheus-client>=0.26.0",
+    "pydantic-settings>=2.16.2",
+    "tokenizers>=0.23.3",
+    "uvicorn[standard]>=0.54.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/qcode-uni-a/paperus"
+Source = "https://github.com/qcode-uni-a/paperus"
+Issues = "https://github.com/qcode-uni-a/paperus/issues"
+
+[project.scripts]
+paperus-ml = "paperus_ml.cli:main"
+
+[build-system]
+requires = ["uv_build>=0.13,<0.14"]
+build-backend = "uv_build"
+
+[tool.ruff]
+line-length = 88
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "B", "UP", "SIM", "RUF", "ANN"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+filterwarnings = ["error"]
+markers = ["download: downloads a real model (set PAPERUS_ML_TEST_DOWNLOAD=1)"]
+
+[dependency-groups]
+dev = [
+    "httpx2>=2.16.0",
+    "onnx>=1.23.2",
+    "pytest>=9.2.0",
+    "ruff>=0.17.3",
+    "ty>=0.0.97",
+]

paperus_ml-0.0.8/src/paperus_ml/__init__.py

@@ -0,0 +1,13 @@
+"""paperus-ml: the machine learning service of paperus.
+
+No telemetry: Hugging Face Hub reads its switch when imported, so it is set
+here, before any submodule imports it; ONNX Runtime's is switched off too.
+"""
+
+import os
+
+os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
+
+import onnxruntime
+
+onnxruntime.disable_telemetry_events()

paperus_ml-0.0.8/src/paperus_ml/cli.py

@@ -0,0 +1,17 @@
+"""``paperus-ml``: run the service (``--host 0.0.0.0`` in containers)."""
+
+import argparse
+import logging
+
+import uvicorn
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(prog="paperus-ml")
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=3003)
+    args = parser.parse_args(argv)
+    logging.basicConfig(level=logging.INFO)
+    uvicorn.run(
+        "paperus_ml.main:create_app", factory=True, host=args.host, port=args.port
+    )

paperus_ml-0.0.8/src/paperus_ml/config.py

@@ -0,0 +1,38 @@
+"""Settings, read from environment variables prefixed with ``PAPERUS_ML_``."""
+
+import os
+from pathlib import Path
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Settings(BaseSettings):
+    model_config = SettingsConfigDict(env_prefix="PAPERUS_ML_", env_file=".env")
+
+    # Downloaded models, one folder per task and model:
+    # <cache>/embedding/bge-m3/. "/cache" (a volume) in the container.
+    cache_path: Path = Field(default_factory=lambda: default_cache_path())
+    # Shared secret: requests need "Authorization: Bearer <token>" when set.
+    token: str | None = None
+    # Unload a model after this many idle seconds (0: never).
+    model_ttl: int = 300
+    # Models to load at startup, never unloaded, e.g. ["bge-m3"].
+    preload: list[str] = []
+    # ONNX Runtime threads per model; 0 lets ONNX Runtime decide.
+    intra_op_threads: int = 0
+    # Most texts and tokens per inference batch (memory grows with both).
+    max_batch_texts: int = 32
+    max_batch_tokens: int = 16384
+    # How long a request waits for a model that is still loading before the
+    # service answers 503 "loading" (the model keeps loading meanwhile).
+    load_wait: float = 2.0
+    # Prometheus metrics on this port (inside the container network); off
+    # unless set, e.g. 9465.
+    metrics_port: int | None = None
+
+
+def default_cache_path() -> Path:
+    """~/.cache/paperus-ml, or under $XDG_CACHE_HOME."""
+    base = os.environ.get("XDG_CACHE_HOME") or Path.home() / ".cache"
+    return Path(base) / "paperus-ml"

paperus_ml-0.0.8/src/paperus_ml/encoder.py

@@ -0,0 +1,116 @@
+"""Running an embedding model with ONNX Runtime."""
+
+import threading
+import time
+from collections.abc import Sequence
+from pathlib import Path
+
+import numpy as np
+import numpy.typing as npt
+import onnxruntime
+from tokenizers import Encoding, Tokenizer
+
+from paperus_ml import metrics
+from paperus_ml.models import Kind, ModelSpec
+
+Vectors = npt.NDArray[np.float32]
+
+
+def session_options(intra_op_threads: int) -> onnxruntime.SessionOptions:
+    options = onnxruntime.SessionOptions()
+    options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
+    # One graph at a time, parallel inside operators.
+    options.inter_op_num_threads = 1
+    options.intra_op_num_threads = intra_op_threads
+    # Denormal floats are slow on CPUs and irrelevant for embeddings.
+    options.add_session_config_entry("session.set_denormal_as_zero", "1")
+    return options
+
+
+def pool(
+    hidden: npt.NDArray[np.float32],
+    mask: npt.NDArray[np.int64],
+    how: str,
+) -> Vectors:
+    """Sentence vectors from token vectors, L2-normalized."""
+    if how == "cls":
+        vectors = hidden[:, 0]
+    else:
+        weights = mask[..., None].astype(np.float32)
+        vectors = (hidden * weights).sum(axis=1) / np.maximum(weights.sum(axis=1), 1e-9)
+    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
+    return (vectors / np.maximum(norms, 1e-12)).astype(np.float32)
+
+
+def batches(lengths: Sequence[int], max_texts: int, max_tokens: int) -> list[list[int]]:
+    """Group text indices by length, so little padding is computed."""
+    order = sorted(range(len(lengths)), key=lambda i: lengths[i])
+    groups: list[list[int]] = []
+    current: list[int] = []
+    for index in order:
+        padded = (len(current) + 1) * lengths[index]
+        if current and (len(current) >= max_texts or padded > max_tokens):
+            groups.append(current)
+            current = []
+        current.append(index)
+    if current:
+        groups.append(current)
+    return groups
+
+
+class Encoder:
+    def __init__(
+        self,
+        spec: ModelSpec,
+        model_dir: Path,
+        *,
+        intra_op_threads: int = 0,
+        max_batch_texts: int = 32,
+        max_batch_tokens: int = 16384,
+    ) -> None:
+        self.spec = spec
+        self.max_batch_texts = max_batch_texts
+        self.max_batch_tokens = max_batch_tokens
+        self.tokenizer = Tokenizer.from_file(str(model_dir / spec.tokenizer_file))
+        self.tokenizer.enable_truncation(spec.max_tokens)
+        self.tokenizer.no_padding()
+        self.session = onnxruntime.InferenceSession(
+            str(model_dir / spec.onnx_file),
+            sess_options=session_options(intra_op_threads),
+            providers=["CPUExecutionProvider"],
+        )
+        self.input_names = {i.name for i in self.session.get_inputs()}
+        # One batch at a time: each already uses all cores, so batches of
+        # simultaneous requests would only compete for them. Per batch, not
+        # per request, so a search query waits for one batch at most.
+        # Reentrant: the speed test holds it across all its batches.
+        self.running = threading.RLock()
+
+    def embed(self, texts: Sequence[str], kind: Kind) -> Vectors:
+        prefix = self.spec.prefix(kind)
+        encodings = self.tokenizer.encode_batch([prefix + t for t in texts])
+        result = np.zeros((len(texts), self.spec.dimensions), dtype=np.float32)
+        groups = batches(
+            [len(e.ids) for e in encodings], self.max_batch_texts, self.max_batch_tokens
+        )
+        for group in groups:
+            result[group] = self._run([encodings[i] for i in group])
+        return result
+
+    def _run(self, encodings: list[Encoding]) -> Vectors:
+        width = max(len(e.ids) for e in encodings)
+        ids = np.zeros((len(encodings), width), dtype=np.int64)
+        mask = np.zeros((len(encodings), width), dtype=np.int64)
+        for row, encoding in enumerate(encodings):
+            ids[row, : len(encoding.ids)] = encoding.ids
+            mask[row, : len(encoding.ids)] = 1
+        feed = {"input_ids": ids, "attention_mask": mask}
+        if "token_type_ids" in self.input_names:
+            feed["token_type_ids"] = np.zeros_like(ids)
+        with self.running:
+            started = time.perf_counter()
+            hidden = self.session.run(None, feed)[0]
+            metrics.BATCH_SECONDS.labels(self.spec.name).observe(
+                time.perf_counter() - started
+            )
+        return pool(np.asarray(hidden, dtype=np.float32), mask, self.spec.pooling)

paperus_ml-0.0.8/src/paperus_ml/main.py

@@ -0,0 +1,250 @@
+"""HTTP API of paperus-ml.
+
+``POST /embed`` with ``{"model": "bge-m3", "kind": "passage", "texts": [...]}``
+returns one L2-normalized vector per text. The caller (paperus) names the
+model in every request, so this service needs no model configuration.
+"""
+
+import asyncio
+import contextlib
+import os
+import secrets
+import time
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from datetime import UTC, datetime
+from typing import Annotated
+
+from fastapi import BackgroundTasks, Depends, FastAPI, Header, HTTPException
+from pydantic import BaseModel, Field
+
+from paperus_ml import metrics
+from paperus_ml.config import Settings
+from paperus_ml.encoder import Encoder
+from paperus_ml.manager import ModelFailed, ModelLoading, ModelManager, UnknownModel
+from paperus_ml.models import MODELS, SCORES, Kind
+from paperus_ml.samples import sample_passages
+
+MAX_TEXTS = 256
+UNLOAD_CHECK_SECONDS = 30
+
+
+class EmbedRequest(BaseModel):
+    model: str
+    kind: Kind = "passage"
+    """Queries and passages get different prefixes in some models."""
+    texts: list[str] = Field(min_length=1, max_length=MAX_TEXTS)
+
+
+class EmbedResponse(BaseModel):
+    model: str
+    dimensions: int
+    embeddings: list[list[float]]
+
+
+class ModelInfo(BaseModel):
+    name: str
+    task: str
+    repo: str
+    revision: str
+    dimensions: int
+    max_tokens: int
+    loaded: bool
+    cached_bytes: int | None
+    """Size of the downloaded files; None if not downloaded yet."""
+    description: str
+    licence: str
+    download_mb: int
+    scores: dict[str, float]
+    """Published retrieval scores (nDCG@10, out of 100) by benchmark."""
+    benchmarks: dict[str, str] = Field(default_factory=lambda: dict(SCORES))
+    """What the score keys mean."""
+
+
+class ServiceConfig(BaseModel):
+    threads: int | None = Field(default=None, ge=0, le=1024)
+    """CPU threads per model; 0: all the service sees; None: its own
+    setting (``PAPERUS_ML_INTRA_OP_THREADS``)."""
+    keep: list[str]
+    """Models to keep loaded (paperus sends the one in use)."""
+
+
+class ServiceState(BaseModel):
+    threads: int
+    keep: list[str]
+    cores: int
+    """CPU cores this service may use."""
+    loaded: list[str]
+    loading: list[str]
+    started: str
+    """When this service started (an ID: a new value means a restart)."""
+
+
+class Speed(BaseModel):
+    model: str
+    passages: int
+    seconds: float
+    passages_per_second: float
+
+
+def usable_cores() -> int:
+    """Cores this process may run on (a container's CPU set, if limited)."""
+    try:
+        return len(os.sched_getaffinity(0))
+    except AttributeError:  # not on Linux
+        return os.cpu_count() or 1
+
+
+def create_app(
+    settings: Settings | None = None, manager: ModelManager | None = None
+) -> FastAPI:
+    settings = settings or Settings()
+    models = manager or ModelManager(settings)
+    started = datetime.now(UTC).isoformat()
+
+    @asynccontextmanager
+    async def lifespan(app: FastAPI) -> AsyncGenerator[None]:
+        if settings.metrics_port:
+            metrics.serve(
+                settings.metrics_port,
+                metrics.ModelsCollector(models.loaded, models.loading),
+            )
+        await asyncio.to_thread(models.preload)
+
+        async def unload_loop() -> None:
+            while True:
+                await asyncio.sleep(UNLOAD_CHECK_SECONDS)
+                await asyncio.to_thread(models.unload_idle)
+
+        task = asyncio.create_task(unload_loop())
+        yield
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await task
+
+    def check_token(authorization: Annotated[str | None, Header()] = None) -> None:
+        if settings.token is None:
+            return
+        scheme, _, token = (authorization or "").partition(" ")
+        valid = scheme.lower() == "bearer" and secrets.compare_digest(
+            token.encode(), settings.token.encode()
+        )
+        if not valid:
+            raise HTTPException(401, "token required")
+
+    app = FastAPI(
+        title="paperus-ml",
+        lifespan=lifespan,
+        generate_unique_id_function=lambda route: route.name,
+    )
+    app.state.models = models
+
+    @app.get("/ping")
+    async def ping() -> str:
+        # Async: answers even when every worker thread is busy.
+        return "pong"
+
+    @app.get("/models", dependencies=[Depends(check_token)])
+    def list_models() -> list[ModelInfo]:
+        loaded = set(models.loaded())
+        return [
+            ModelInfo(
+                name=s.name,
+                task=s.task,
+                repo=s.repo,
+                revision=s.revision,
+                dimensions=s.dimensions,
+                max_tokens=s.max_tokens,
+                loaded=s.name in loaded,
+                cached_bytes=models.cached_bytes(s),
+                description=s.description,
+                licence=s.licence,
+                download_mb=s.download_mb,
+                scores=dict(s.scores),
+            )
+            for s in MODELS.values()
+        ]
+
+    @app.delete("/models/{name}", status_code=204, dependencies=[Depends(check_token)])
+    def delete_model(name: str) -> None:
+        """Remove a downloaded model to free disk space."""
+        try:
+            models.delete_cached(name)
+        except UnknownModel as error:
+            raise HTTPException(404, f"unknown model {name!r}") from error
+
+    @app.get("/config", dependencies=[Depends(check_token)])
+    def get_config() -> ServiceState:
+        return ServiceState(
+            threads=models.threads,
+            keep=models.kept,
+            cores=usable_cores(),
+            loaded=models.loaded(),
+            loading=models.loading(),
+            started=started,
+        )
+
+    @app.put("/config", dependencies=[Depends(check_token)])
+    def set_config(body: ServiceConfig, tasks: BackgroundTasks) -> ServiceState:
+        """Apply paperus' settings; kept models load in the background."""
+        try:
+            missing = models.configure(threads=body.threads, keep=body.keep)
+        except UnknownModel as error:
+            raise HTTPException(404, f"unknown model {error.args[0]!r}") from error
+        if missing:
+            tasks.add_task(models.load_kept)
+        return get_config()
+
+    @app.post("/models/{name}/speed", dependencies=[Depends(check_token)])
+    def measure_speed(name: str) -> Speed:
+        """Embed sample passages (about 250 words each, like paperus'
+        passages) and time it: speed differs much between machines. A model
+        not loaded yet starts loading and the answer is 503 "loading": ask
+        again. Other requests for the model wait while it measures, so they
+        do not slow the measurement (other models still can)."""
+        encoder = _model(name)
+        passages = sample_passages()
+        with encoder.running:
+            encoder.embed(passages[:2], "passage")  # warm up
+            started = time.monotonic()
+            encoder.embed(passages, "passage")
+            seconds = time.monotonic() - started
+        return Speed(
+            model=name,
+            passages=len(passages),
+            seconds=round(seconds, 2),
+            passages_per_second=round(len(passages) / max(seconds, 1e-6), 1),
+        )
+
+    def _model(name: str) -> Encoder:
+        try:
+            return models.get(name, wait=settings.load_wait)
+        except UnknownModel as error:
+            raise HTTPException(404, f"unknown model {name!r}") from error
+        except ModelLoading as error:
+            raise HTTPException(
+                503, f"model {name!r} is loading", headers={"Retry-After": "10"}
+            ) from error
+        except ModelFailed as error:
+            raise HTTPException(
+                502, f"loading model {name!r} failed: {error.args[1]}"
+            ) from error
+
+    @app.post("/embed", dependencies=[Depends(check_token)])
+    def embed(body: EmbedRequest) -> EmbedResponse:
+        """Runs in FastAPI's thread pool; ONNX Runtime releases the GIL."""
+        try:
+            vectors = _model(body.model).embed(body.texts, body.kind)
+        except HTTPException as error:
+            outcome = "loading" if error.status_code == 503 else "error"
+            metrics.REQUESTS.labels(body.model, body.kind, outcome).inc()
+            raise
+        metrics.REQUESTS.labels(body.model, body.kind, "ok").inc()
+        metrics.TEXTS.labels(body.model, body.kind).inc(len(body.texts))
+        return EmbedResponse(
+            model=body.model,
+            dimensions=vectors.shape[1],
+            embeddings=vectors.tolist(),
+        )
+
+    return app

paperus_ml-0.0.8/src/paperus_ml/manager.py

@@ -0,0 +1,217 @@
+"""Loading models on first use, sharing them, and unloading idle ones."""
+
+import logging
+import shutil
+import threading
+import time
+from collections.abc import Callable
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+from paperus_ml.config import Settings
+from paperus_ml.encoder import Encoder
+from paperus_ml.models import MODELS, ModelSpec
+
+log = logging.getLogger(__name__)
+
+Loader = Callable[[ModelSpec], Encoder]
+
+
+class UnknownModel(KeyError):
+    pass
+
+
+class ModelLoading(Exception):
+    """The model is still downloading or loading; ask again shortly."""
+
+
+class ModelFailed(Exception):
+    """Loading the model failed (e.g. the download); the message says why."""
+
+
+class ModelManager:
+    def __init__(self, settings: Settings, loader: Loader | None = None) -> None:
+        self.settings = settings
+        self._load = loader or self._download_and_load
+        # name → (encoder, last used, threads it was loaded with)
+        self._models: dict[str, tuple[Encoder, float, int]] = {}
+        self._loading: dict[str, threading.Event] = {}
+        self._errors: dict[str, str] = {}
+        self._guard = threading.Lock()
+        self._locks: dict[str, threading.Lock] = {
+            name: threading.Lock() for name in MODELS
+        }
+        self._keep = set(settings.preload)
+        self.threads = settings.intra_op_threads
+        """CPU threads per model (0: all ONNX Runtime sees); set by paperus."""
+
+    def model_dir(self, spec: ModelSpec) -> Path:
+        """<cache>/<task>/<model>, with plain files: the shared Hugging Face
+        cache uses symlinks, and ONNX Runtime refuses weight files that
+        resolve outside the model's folder."""
+        return self.settings.cache_path / spec.task / spec.name
+
+    def cached_bytes(self, spec: ModelSpec) -> int | None:
+        """Size on disk if the model is downloaded, else None."""
+        directory = self.model_dir(spec)
+        if not (directory / spec.onnx_file).is_file():
+            return None
+        return sum(p.stat().st_size for p in directory.rglob("*") if p.is_file())
+
+    def delete_cached(self, name: str) -> None:
+        """Unload a model and remove its files; it downloads again when used."""
+        spec = MODELS.get(name)
+        if spec is None:
+            raise UnknownModel(name)
+        with self._locks[name]:
+            self._models.pop(name, None)
+            self._errors.pop(name, None)
+            shutil.rmtree(self.model_dir(spec), ignore_errors=True)
+
+    def _download_and_load(self, spec: ModelSpec) -> Encoder:
+        model_dir = self.model_dir(spec)
+        # Revision changes download the changed files into the same folder.
+        snapshot_download(
+            spec.repo,
+            revision=spec.revision,
+            local_dir=model_dir,
+            allow_patterns=[f"{spec.onnx_file}*", spec.tokenizer_file],
+        )
+        return Encoder(
+            spec,
+            model_dir,
+            intra_op_threads=self.threads,
+            max_batch_texts=self.settings.max_batch_texts,
+            max_batch_tokens=self.settings.max_batch_tokens,
+        )
+
+    def get(self, name: str, wait: float | None = None) -> Encoder:
+        """The model, downloading and loading it on first use (and again when
+        the threads changed). ``wait=None`` blocks until it is ready; with
+        ``wait`` in seconds, loading continues in the background and
+        ``ModelLoading`` is raised if it takes longer, so requests do not pile
+        up behind a download."""
+        spec = MODELS.get(name)
+        if spec is None:
+            raise UnknownModel(name)
+        ready = self._ready(name)
+        if ready is not None:
+            return ready
+        if wait is None:
+            return self._load_now(name, spec)
+        done = self._load_in_background(name)
+        done.wait(wait)
+        ready = self._ready(name)
+        if ready is not None:
+            return ready
+        if name in self._errors and not self._loading_now(name):
+            raise ModelFailed(name, self._errors[name])
+        raise ModelLoading(name)
+
+    def _ready(self, name: str) -> Encoder | None:
+        """The loaded model, if loaded with the current threads."""
+        entry = self._models.get(name)
+        if entry is None or entry[2] != self.threads:
+            return None
+        self._models[name] = (entry[0], time.monotonic(), entry[2])
+        return entry[0]
+
+    def _load_now(self, name: str, spec: ModelSpec) -> Encoder:
+        with self._locks[name]:
+            ready = self._ready(name)
+            if ready is not None:
+                return ready
+            # Drop an old copy (other threads) before loading the new one.
+            self._models.pop(name, None)
+            threads = self.threads
+            log.info("loading model %s", name)
+            started = time.monotonic()
+            try:
+                encoder = self._load(spec)
+            except Exception as error:
+                self._errors[name] = str(error) or type(error).__name__
+                raise
+            self._errors.pop(name, None)
+            log.info("loaded %s in %.1f s", name, time.monotonic() - started)
+            self._models[name] = (encoder, time.monotonic(), threads)
+            return encoder
+
+    def _load_in_background(self, name: str) -> threading.Event:
+        with self._guard:
+            if name in self._loading:
+                return self._loading[name]
+            done = threading.Event()
+            self._loading[name] = done
+
+        def load() -> None:
+            try:
+                self.get(name)
+            except Exception:
+                log.warning("loading %s failed", name, exc_info=True)
+            finally:
+                with self._guard:
+                    self._loading.pop(name, None)
+                done.set()
+
+        threading.Thread(target=load, name=f"load-{name}", daemon=True).start()
+        return done
+
+    def _loading_now(self, name: str) -> bool:
+        with self._guard:
+            return name in self._loading
+
+    def loading(self) -> list[str]:
+        with self._guard:
+            return sorted(self._loading)
+
+    def loaded(self) -> list[str]:
+        return sorted(self._models)
+
+    def unload_idle(self, now: float | None = None) -> list[str]:
+        """Unload models idle longer than the TTL, except kept ones. Skips a
+        model that is busy loading; checks again under its lock."""
+        ttl = self.settings.model_ttl
+        if ttl <= 0:
+            return []
+        now = time.monotonic() if now is None else now
+        unloaded = []
+        for name in list(self._models):
+            lock = self._locks[name]
+            if not lock.acquire(blocking=False):
+                continue
+            try:
+                entry = self._models.get(name)
+                if entry is None or name in self._keep or now - entry[1] < ttl:
+                    continue
+                self._models.pop(name, None)
+            finally:
+                lock.release()
+            unloaded.append(name)
+            log.info("unloaded idle model %s", name)
+        return unloaded
+
+    @property
+    def kept(self) -> list[str]:
+        return sorted(self._keep)
+
+    def configure(self, *, threads: int | None, keep: list[str]) -> list[str]:
+        """Settings from paperus' administration: threads per model (None:
+        this service's own setting), and the models to keep loaded (the one
+        in use). Models loaded with other threads load again when next used
+        (requests running keep theirs). Returns the kept models still to load
+        (``load_kept`` loads them, which can take a while)."""
+        unknown = [name for name in keep if name not in MODELS]
+        if unknown:
+            raise UnknownModel(unknown[0])
+        self.threads = self.settings.intra_op_threads if threads is None else threads
+        self._keep = set(keep)
+        return [name for name in keep if self._ready(name) is None]
+
+    def load_kept(self) -> None:
+        for name in self.kept:
+            self.get(name)
+
+    def preload(self) -> None:
+        for name in self.settings.preload:
+            self.get(name)

paperus_ml-0.0.8/src/paperus_ml/metrics.py

@@ -0,0 +1,55 @@
+"""Prometheus metrics of the ML service, on a port of their own
+(``PAPERUS_ML_METRICS_PORT``): requests, batch times, loaded models."""
+
+from collections.abc import Callable, Iterator
+
+from prometheus_client import REGISTRY, Counter, Histogram, start_http_server
+from prometheus_client.core import GaugeMetricFamily, Metric
+from prometheus_client.registry import Collector
+
+REQUESTS = Counter(
+    "paperus_ml_requests_total",
+    "Embedding requests, by model, kind and outcome",
+    ["model", "kind", "outcome"],
+)
+TEXTS = Counter("paperus_ml_texts_total", "Texts embedded", ["model", "kind"])
+BATCH_SECONDS = Histogram(
+    "paperus_ml_batch_duration_seconds",
+    "Time to run one batch through a model",
+    ["model"],
+    buckets=(0.01, 0.05, 0.1, 0.25, 0.5, 1, 2, 5, 10, 30, 60),
+)
+
+
+class ModelsCollector(Collector):
+    """Loaded and loading models, read at each scrape."""
+
+    def __init__(
+        self, loaded: Callable[[], list[str]], loading: Callable[[], list[str]]
+    ) -> None:
+        self.loaded = loaded
+        self.loading = loading
+
+    def collect(self) -> Iterator[Metric]:
+        gauge = GaugeMetricFamily(
+            "paperus_ml_model_loaded",
+            "1 for each model in memory, 0.5 while loading",
+            labels=["model"],
+        )
+        for name in self.loaded():
+            gauge.add_metric([name], 1.0)
+        for name in self.loading():
+            gauge.add_metric([name], 0.5)
+        yield gauge
+
+
+_started = False
+
+
+def serve(port: int, collector: Collector) -> None:
+    global _started
+    if _started:
+        return
+    REGISTRY.register(collector)
+    start_http_server(port)
+    _started = True

paperus_ml-0.0.8/src/paperus_ml/models.py

@@ -0,0 +1,214 @@
+"""The models paperus-ml can serve.
+
+Each is an ONNX export from Hugging Face, pinned to a revision so a changed
+file upstream cannot silently change the vectors stored in paperus.
+
+Scores are the published retrieval averages of the full-precision models
+(int8 exports typically lose about a point), checked 2026-11-05: MTEB
+(eng, v2) and MTEB (deu, v1) Retrieval from the MTEB leaderboard, MMTEB
+(Multilingual, v2) Retrieval likewise, or from the MMTEB paper
+(arXiv 2502.13595) where the leaderboard shows no average.
+"""
+
+from dataclasses import dataclass
+from typing import Literal
+
+Task = Literal["embedding"]
+
+SCORES = {
+    "en": "Retrieval, English (MTEB eng v2)",
+    "de": "Retrieval, German (MTEB deu v1)",
+    "multilingual": "Retrieval, multilingual (MMTEB v2)",
+}
+Pooling = Literal["cls", "mean"]
+Kind = Literal["query", "passage"]
+
+
+@dataclass(frozen=True)
+class ModelSpec:
+    name: str
+    repo: str
+    revision: str
+    onnx_file: str
+    tokenizer_file: str
+    pooling: Pooling
+    dimensions: int
+    max_tokens: int
+    query_prefix: str = ""
+    passage_prefix: str = ""
+    task: Task = "embedding"
+    # For choosing a model (shown in paperus' administration).
+    description: str = ""
+    licence: str = ""
+    download_mb: int = 0
+    scores: tuple[tuple[str, float], ...] = ()
+    """Published retrieval scores (nDCG@10, out of 100) by benchmark, see SCORES."""
+
+    def prefix(self, kind: Kind) -> str:
+        return self.query_prefix if kind == "query" else self.passage_prefix
+
+
+MODELS = {
+    spec.name: spec
+    for spec in [
+        # Multilingual, 1024 dimensions, long inputs; ONNX from the authors.
+        ModelSpec(
+            name="bge-m3",
+            repo="BAAI/bge-m3",
+            revision="5617a9f61b028005a4858fdac845db406aefb181",
+            onnx_file="onnx/model.onnx",
+            tokenizer_file="onnx/tokenizer.json",
+            pooling="cls",
+            dimensions=1024,
+            max_tokens=8192,
+            description=(
+                "Multilingual, long inputs; strong but slow on a CPU (full precision)."
+            ),
+            licence="MIT",
+            download_mb=2270,
+            scores=(("multilingual", 54.6),),
+        ),
+        # The same model in int8 (570 MB): several times faster on a CPU,
+        # nearly the same quality; vectors differ, so it counts as its own.
+        ModelSpec(
+            name="bge-m3-int8",
+            repo="Xenova/bge-m3",
+            revision="4de13258303883538bd53b696b452bf8099f0858",
+            onnx_file="onnx/model_int8.onnx",
+            tokenizer_file="tokenizer.json",
+            pooling="cls",
+            dimensions=1024,
+            max_tokens=8192,
+            description=(
+                "bge-m3 in int8: faster and a quarter of the size, nearly the same "
+                "quality."
+            ),
+            licence="MIT",
+            download_mb=570,
+            scores=(("multilingual", 54.6),),
+        ),
+        # Snowflake Arctic Embed v2 (Apache-2.0): multilingual, made for
+        # retrieval, long inputs; int8.
+        ModelSpec(
+            name="arctic-embed-m-v2",
+            repo="Snowflake/snowflake-arctic-embed-m-v2.0",
+            revision="95c2741480856aa9666782eb4afe11959938017f",
+            onnx_file="onnx/model_quantized.onnx",
+            tokenizer_file="tokenizer.json",
+            pooling="cls",
+            dimensions=768,
+            max_tokens=8192,
+            query_prefix="query: ",
+            description=(
+                "Recommended: multilingual and made for search; good quality, fast, "
+                "small."
+            ),
+            licence="Apache-2.0",
+            download_mb=311,
+            scores=(
+                ("en", 58.4),
+                ("de", 55.2),
+                ("multilingual", 54.8),
+            ),
+        ),
+        ModelSpec(
+            name="arctic-embed-l-v2",
+            repo="Snowflake/snowflake-arctic-embed-l-v2.0",
+            revision="ac6544c8a46e00af67e330e85a9028c66b8cfd9a",
+            onnx_file="onnx/model_quantized.onnx",
+            tokenizer_file="tokenizer.json",
+            pooling="cls",
+            dimensions=1024,
+            max_tokens=8192,
+            query_prefix="query: ",
+            description=(
+                "The best quality here, multilingual; about half as fast as the medium "
+                "one."
+            ),
+            licence="Apache-2.0",
+            download_mb=570,
+            scores=(
+                ("en", 58.6),
+                ("de", 57.1),
+                ("multilingual", 58.4),
+            ),
+        ),
+        # Alibaba GTE multilingual (Apache-2.0): long inputs, int8.
+        ModelSpec(
+            name="gte-multilingual-base",
+            repo="onnx-community/gte-multilingual-base",
+            revision="2edbf5e672aab465f9ed4c154a8b61791c082c69",
+            onnx_file="onnx/model_quantized.onnx",
+            tokenizer_file="tokenizer.json",
+            pooling="cls",
+            dimensions=768,
+            max_tokens=8192,
+            description="Multilingual, long inputs, fast.",
+            licence="Apache-2.0",
+            download_mb=340,
+            scores=(("en", 50.9),),
+        ),
+        # E5 (MIT), int8: base and large; inputs up to 512 tokens.
+        ModelSpec(
+            name="multilingual-e5-base",
+            repo="Xenova/multilingual-e5-base",
+            revision="1ec9243030a27d1a115d5c340572074c125b58b2",
+            onnx_file="onnx/model_quantized.onnx",
+            tokenizer_file="tokenizer.json",
+            pooling="mean",
+            dimensions=768,
+            max_tokens=512,
+            query_prefix="query: ",
+            passage_prefix="passage: ",
+            description="Multilingual, fast; inputs up to 512 tokens.",
+            licence="MIT",
+            download_mb=279,
+            scores=(
+                ("en", 49.0),
+                ("de", 48.8),
+                ("multilingual", 52.7),
+            ),
+        ),
+        ModelSpec(
+            name="multilingual-e5-large",
+            repo="Xenova/multilingual-e5-large",
+            revision="00fc3aeb3dbb95842de2ac1961d33c6319acf57b",
+            onnx_file="onnx/model_quantized.onnx",
+            tokenizer_file="tokenizer.json",
+            pooling="mean",
+            dimensions=1024,
+            max_tokens=512,
+            query_prefix="query: ",
+            passage_prefix="passage: ",
+            description="Multilingual; inputs up to 512 tokens; slow.",
+            licence="MIT",
+            download_mb=562,
+            scores=(
+                ("en", 51.5),
+                ("de", 50.0),
+                ("multilingual", 54.1),
+            ),
+        ),
+        # Small and fast (int8, 118 MB); for development and small machines.
+        ModelSpec(
+            name="multilingual-e5-small",
+            repo="Xenova/multilingual-e5-small",
+            revision="761b726dd34fb83930e26aab4e9ac3899aa1fa78",
+            onnx_file="onnx/model_quantized.onnx",
+            tokenizer_file="tokenizer.json",
+            pooling="mean",
+            dimensions=384,
+            max_tokens=512,
+            query_prefix="query: ",
+            passage_prefix="passage: ",
+            description="Small and very fast, for weak machines; lower quality.",
+            licence="MIT",
+            download_mb=118,
+            scores=(
+                ("en", 46.4),
+                ("de", 48.2),
+                ("multilingual", 49.3),
+            ),
+        ),
+    ]
+}

paperus_ml-0.0.8/src/paperus_ml/samples.py

@@ -0,0 +1,26 @@
+# ruff: noqa: E501 (one sample sentence per line)
+"""Sample passages for measuring speed: research prose of the length
+paperus embeds (about 250 words), varied so no two are the same."""
+
+_SENTENCES = [
+    "We study how blind zones below the radar beam hide light rainfall in mountain valleys.",
+    "The model fills the missing reflectivity from neighbouring heights and nearby stations.",
+    "Results are compared against rain gauges over three winters and two summer seasons.",
+    "Errors grow with distance from the radar and with the height of the lowest usable beam.",
+    "Die Messungen zeigen, dass flache Niederschläge in Tälern häufig unterschätzt werden.",
+    "A diffusion model trained on complete scans learns the vertical structure of storms.",
+    "We release the code and the data, so that other groups can repeat every experiment.",
+    "Uncertainty is estimated from an ensemble of twenty samples for each missing column.",
+    "Les résultats confirment l'intérêt de la méthode pour la prévision à très courte échéance.",
+    "Ablation studies show that both the height prior and the station data are needed.",
+    "Compared with interpolation, the error falls by a third in the lowest kilometre.",
+    "Limitations remain for convective cells that form entirely inside the blind zone.",
+]
+
+
+def sample_passages(count: int = 48) -> list[str]:
+    passages = []
+    for n in range(count):
+        sentences = [_SENTENCES[(n * 5 + i * 7) % len(_SENTENCES)] for i in range(18)]
+        passages.append(f"Section {n + 1}. " + " ".join(sentences))
+    return passages