| """The on-prem GPU backend: open weights served on vLLM, for hosts with a GPU and no hosted API (#83). |
| |
| Two stages, two models, both reached through the provider seam so no call site changes: |
| |
| - Page transcription. `reocr.py` sends each page image to a vision model. olmOCR-2-7B is the default, |
| a 7B model trained for document transcription, served on its own vLLM endpoint. |
| - Extraction. `llm_extraction.py` forces one tool call whose schema `fields.py` derives. A larger |
| instruct model answers it, and vLLM constrains generation to that schema with xgrammar, so the output |
| cannot take a shape the checklist rejects. |
| |
| Configured with four environment variables. `VLLM_URL` and `VLLM_MODEL` name the extractor. |
| `VLLM_REOCR_URL` and `VLLM_REOCR_MODEL` name the transcription endpoint, each falling back to the |
| extractor's value and the olmOCR default, so a single server hosting one model still works. |
| |
| The backend reports unavailable until `VLLM_URL` and `VLLM_MODEL` are set, so the settings list disables |
| it on a host with no server rather than letting a switch fail at extraction (#94). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| from typing import Any, Optional |
|
|
| from endopath import llm |
|
|
| GPU_VLLM = "gpu_vllm" |
|
|
| URL_ENV = "VLLM_URL" |
| MODEL_ENV = "VLLM_MODEL" |
| REOCR_URL_ENV = "VLLM_REOCR_URL" |
| REOCR_MODEL_ENV = "VLLM_REOCR_MODEL" |
|
|
| |
| |
| DEFAULT_REOCR_MODEL = "allenai/olmOCR-2-7B-1025" |
|
|
| |
| |
| GUIDED_DECODING = "xgrammar" |
|
|
|
|
| def _env(name: str) -> str: |
| return (os.environ.get(name) or "").strip() |
|
|
|
|
| def _require(name: str) -> str: |
| value = _env(name) |
| if not value: |
| raise RuntimeError( |
| f"{name} is not set. The GPU backend serves open weights on vLLM: set {URL_ENV} and " |
| f"{MODEL_ENV} to the extractor endpoint, and optionally {REOCR_URL_ENV} and " |
| f"{REOCR_MODEL_ENV} for a separate transcription endpoint." |
| ) |
| return value |
|
|
|
|
| def build_extraction_client() -> Any: |
| """The extractor: a larger instruct model on vLLM, with xgrammar holding its output to the checklist |
| schema `llm_extraction.build_checklist_tool` derives from `fields.py`.""" |
| return llm.openai_compatible_client( |
| _require(URL_ENV), _require(MODEL_ENV), guided_decoding=GUIDED_DECODING |
| ) |
|
|
|
|
| def build_reocr_client() -> Any: |
| """The page-transcription model. Falls back to the extractor's endpoint when only one server runs, so |
| a single-server host still transcribes.""" |
| base_url = _env(REOCR_URL_ENV) or _require(URL_ENV) |
| model = _env(REOCR_MODEL_ENV) or DEFAULT_REOCR_MODEL |
| return llm.openai_compatible_client(base_url, model) |
|
|
|
|
| def _extraction_model() -> Optional[str]: |
| """The extractor model this backend serves, recorded on the activity a run writes (#151). None when |
| VLLM_MODEL is unset, because no extraction can have run under it.""" |
| return _env(MODEL_ENV) or None |
|
|
|
|
| def _available() -> bool: |
| """Whether a vLLM endpoint is configured. The settings list disables the backend until it is, rather |
| than offering a switch that fails at the first model call.""" |
| return bool(_env(URL_ENV) and _env(MODEL_ENV)) |
|
|
|
|
| llm.register_backend( |
| llm.Backend( |
| GPU_VLLM, |
| "GPU on-prem (open weights on vLLM)", |
| needs_key=False, |
| provider=llm.OPENAI_COMPATIBLE, |
| build=build_extraction_client, |
| build_reocr=build_reocr_client, |
| runs_llm_extraction=True, |
| extraction_model=_extraction_model, |
| |
| |
| page_reader=None, |
| available=_available, |
| ) |
| ) |
|
|
