FastAPI service wrapping the Surya-OCR-2 model (datalab-to) served through vLLM: legacy /v1/api/ai/* endpoints, an OpenAI-compatible /v1/chat/completions endpoint, a coalescing request batcher, a local OCR CLI, Docker packaging, multilingual example outputs, and quantization/concurrency benchmarks. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
42 lines
1.3 KiB
Python
42 lines
1.3 KiB
Python
import json
|
|
|
|
from scripts.quant.capture import extract_capture, write_capture
|
|
|
|
|
|
def test_extract_capture_pulls_text_boxes_and_latency():
|
|
body = {
|
|
"data": {
|
|
"text_lines": "line one\nline two",
|
|
"elapsed_seconds": 5.4,
|
|
"ocr_text_json": {
|
|
"blocks": [
|
|
{"bbox": [1, 2, 3, 4], "html": "line one"},
|
|
{"bbox": [5, 6, 7, 8], "html": "line two"},
|
|
]
|
|
},
|
|
},
|
|
"message": "success",
|
|
"code": 200,
|
|
}
|
|
|
|
result = extract_capture(body)
|
|
|
|
assert result["text"] == "line one\nline two"
|
|
assert result["boxes"] == [[1, 2, 3, 4], [5, 6, 7, 8]]
|
|
assert result["elapsed_seconds"] == 5.4
|
|
|
|
|
|
def test_extract_capture_tolerates_missing_fields():
|
|
result = extract_capture({"data": {}})
|
|
assert result == {"text": "", "boxes": [], "elapsed_seconds": None}
|
|
|
|
|
|
def test_write_capture_writes_txt_and_json(tmp_path):
|
|
cap = {"text": "hello", "boxes": [[0, 0, 1, 1]], "elapsed_seconds": 2.0}
|
|
|
|
write_capture(cap, tmp_path, "page1")
|
|
|
|
assert (tmp_path / "page1.txt").read_text(encoding="utf-8") == "hello"
|
|
loaded = json.loads((tmp_path / "page1.json").read_text(encoding="utf-8"))
|
|
assert loaded == {"boxes": [[0, 0, 1, 1]], "elapsed_seconds": 2.0}
|