diff --git a/README.md b/README.md index f133067..77f2fde 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,5 @@ # CHURRO logo CHURRO -CHURRO is an OCR toolkit for historical document transcription, built to make handwritten and printed sources readable at high accuracy and lower cost. - -It works with all major OCR proividers and vision-language models, and provides first-party support for the CHURRO 3B model and CHURRO-DS dataset. - [![Model](https://img.shields.io/badge/Model-CHURRO%203B-8A4FFF)](https://huggingface.co/stanford-oval/churro-3B) [![Dataset](https://img.shields.io/badge/Dataset-CHURRO--DS-0A7BBB)](https://huggingface.co/datasets/stanford-oval/churro-dataset) [![Paper](https://img.shields.io/badge/Paper-arXiv-B31B1B)](https://arxiv.org/abs/2509.19768) @@ -11,6 +7,10 @@ It works with all major OCR proividers and vision-language models, and provides [![Leaderboard](https://img.shields.io/badge/Leaderboard-Benchmark%20Snapshot-6B7280)](https://stanford-oval.github.io/Churro/leaderboard.html) [![GitHub Stars](https://img.shields.io/github/stars/stanford-oval/churro?style=social)](https://github.com/stanford-oval/churro/stargazers) +CHURRO is an OCR toolkit for historical document transcription, built to make handwritten and printed sources readable at high accuracy and lower cost. + +It works with all major OCR providers and vision-language models, and provides first-party support for the CHURRO 3B model and CHURRO-DS dataset. + - CHURRO 3B exceeds the accuracy of Gemini 2.5 Pro at 15.5x lower cost. - CHURRO-DS contains ~100K pages from 155 historical collections spanning 22 centuries and 46 language clusters. @@ -24,7 +24,8 @@ It works with all major OCR proividers and vision-language models, and provides ## Quick Try ```bash -pip install "churro-ocr[hf]" +uv tool install churro-ocr +churro-ocr install hf churro-ocr transcribe --image scan.png --backend hf --model stanford-oval/churro-3B ``` diff --git a/benchmark_results.json b/benchmark_results.json index 0fc66a2..0ccccbe 100644 --- a/benchmark_results.json +++ b/benchmark_results.json @@ -3,7 +3,7 @@ "modelName": "Churro", "modelId": "stanford-oval/churro-3B", "modelUrl": "https://huggingface.co/stanford-oval/churro-3B", - "hasIcon": true, + "iconPath": "_static/img/churro.png", "printed": 82.3309, "handwritten": 70.0965, "total": 75.7431 @@ -12,16 +12,25 @@ "modelName": "Gemini 3 Flash", "modelId": "gemini-3-flash", "modelUrl": "https://ai.google.dev/gemini-api/docs/models", - "hasIcon": false, + "iconPath": "_static/img/gemini-symbol.png", "printed": 82.6, "handwritten": 66.5, "total": 73.9 + }, + { + "modelName": "Gemini 3.1 Pro Preview", + "modelId": "gemini-3.1-pro-preview", + "modelUrl": "https://ai.google.dev/gemini-api/docs/models", + "iconPath": "_static/img/gemini-symbol.png", + "printed": 80.3, + "handwritten": 68.3, + "total": 73.8 }, { "modelName": "Gemini 3 Pro", "modelId": "gemini-3-pro", "modelUrl": "https://ai.google.dev/gemini-api/docs/models", - "hasIcon": false, + "iconPath": "_static/img/gemini-symbol.png", "printed": 78.2, "handwritten": 66.7, "total": 72.0 @@ -30,16 +39,25 @@ "modelName": "Gemini 2.5 Pro", "modelId": "gemini-2.5-pro", "modelUrl": "https://ai.google.dev/gemini-api/docs/models", - "hasIcon": false, + "iconPath": "_static/img/gemini-symbol.png", "printed": 80.8523, "handwritten": 63.6329, "total": 71.5803 }, + { + "modelName": "dots.mocr", + "modelId": "rednote-hilab/dots.mocr", + "modelUrl": "https://huggingface.co/rednote-hilab/dots.mocr", + "iconPath": null, + "printed": 81.2, + "handwritten": 55.0, + "total": 67.1 + }, { "modelName": "Gemini 2.5 Flash", "modelId": "gemini-2.5-flash", "modelUrl": "https://ai.google.dev/gemini-api/docs/models", - "hasIcon": false, + "iconPath": "_static/img/gemini-symbol.png", "printed": 73.719, "handwritten": 58.7283, "total": 65.6471 @@ -48,7 +66,7 @@ "modelName": "Qwen 3 VL (8B)", "modelId": "Qwen/Qwen3-VL-8B-Instruct", "modelUrl": "https://huggingface.co/Qwen/Qwen3-VL-8B-Instruct", - "hasIcon": false, + "iconPath": "_static/img/qwen-logo.svg", "printed": 76.6136, "handwritten": 48.4963, "total": 61.4735 @@ -57,7 +75,7 @@ "modelName": "Qwen 3 VL (30B-A3B)", "modelId": "Qwen/Qwen3-VL-30B-A3B-Instruct", "modelUrl": "https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct", - "hasIcon": false, + "iconPath": "_static/img/qwen-logo.svg", "printed": 74.8873, "handwritten": 49.558, "total": 61.2485 @@ -66,7 +84,7 @@ "modelName": "NuMarkdown", "modelId": "numind/NuMarkdown-8B-Thinking", "modelUrl": "https://huggingface.co/numind/NuMarkdown-8B-Thinking", - "hasIcon": false, + "iconPath": null, "printed": 72.7318, "handwritten": 51.2224, "total": 61.1498 @@ -75,7 +93,8 @@ "modelName": "GPT-4.1 Mini", "modelId": "gpt-4.1-mini-2025-04-14", "modelUrl": "https://platform.openai.com/docs/models/gpt-4.1-mini", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 73.1042, "handwritten": 50.2487, "total": 60.7974 @@ -84,7 +103,7 @@ "modelName": "Qwen 2.5 VL (72B)", "modelId": "Qwen/Qwen2.5-VL-72B-Instruct", "modelUrl": "https://huggingface.co/Qwen/Qwen2.5-VL-72B-Instruct", - "hasIcon": false, + "iconPath": "_static/img/qwen-logo.svg", "printed": 66.2783, "handwritten": 54.4717, "total": 59.9209 @@ -93,7 +112,7 @@ "modelName": "Azure OCR", "modelId": null, "modelUrl": null, - "hasIcon": false, + "iconPath": "_static/img/azure-logo.png", "printed": 71.8617, "handwritten": 47.7443, "total": 58.8754 @@ -102,7 +121,8 @@ "modelName": "GPT-5.2", "modelId": "gpt-5.2-2025-12-11", "modelUrl": "https://platform.openai.com/docs/models/gpt-5.2", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 71.4, "handwritten": 46.8, "total": 58.2 @@ -111,7 +131,8 @@ "modelName": "GPT-5 Mini", "modelId": "gpt-5-mini-2025-08-07", "modelUrl": "https://platform.openai.com/docs/models/gpt-5-mini", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 69.8107, "handwritten": 47.5121, "total": 57.8038 @@ -120,7 +141,7 @@ "modelName": "Claude Sonnet 3.7", "modelId": "claude-3-7-sonnet-20250219", "modelUrl": "https://docs.claude.com/en/docs/about-claude/models/overview", - "hasIcon": false, + "iconPath": "_static/img/claude-symbol.svg", "printed": 70.1866, "handwritten": 46.5904, "total": 57.481 @@ -129,7 +150,7 @@ "modelName": "RolmOCR", "modelId": "reducto/RolmOCR", "modelUrl": "https://huggingface.co/reducto/RolmOCR", - "hasIcon": false, + "iconPath": "_static/img/reducto-logomark.svg", "printed": 67.2281, "handwritten": 49.0075, "total": 57.417 @@ -138,7 +159,7 @@ "modelName": "Qwen 3 VL (4B)", "modelId": "Qwen/Qwen3-VL-4B-Instruct", "modelUrl": "https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct", - "hasIcon": false, + "iconPath": "_static/img/qwen-logo.svg", "printed": 74.9451, "handwritten": 40.6463, "total": 56.4765 @@ -147,7 +168,7 @@ "modelName": "Nanonets OCR", "modelId": "nanonets/Nanonets-OCR-s", "modelUrl": "https://huggingface.co/nanonets/Nanonets-OCR-s", - "hasIcon": false, + "iconPath": null, "printed": 69.7033, "handwritten": 43.1781, "total": 55.4205 @@ -156,16 +177,34 @@ "modelName": "olmOCR", "modelId": "allenai/olmOCR-7B-0825", "modelUrl": "https://huggingface.co/allenai/olmOCR-7B-0825", - "hasIcon": false, + "iconPath": "_static/img/ai2-symbol.svg", "printed": 69.766, "handwritten": 41.5165, "total": 54.5547 }, + { + "modelName": "olmOCR 2", + "modelId": "allenai/olmOCR-2-7B-1025", + "modelUrl": "https://huggingface.co/allenai/olmOCR-2-7B-1025", + "iconPath": "_static/img/ai2-symbol.svg", + "printed": 71.1, + "handwritten": 44.6, + "total": 56.8 + }, + { + "modelName": "LiquidAI LFM2.5 VL (1.6B)", + "modelId": "LiquidAI/LFM2.5-VL-1.6B", + "modelUrl": "https://huggingface.co/LiquidAI/LFM2.5-VL-1.6B", + "iconPath": null, + "printed": 40.2, + "handwritten": 25.5, + "total": 32.3 + }, { "modelName": "Qwen 2.5 VL (3B)", "modelId": "Qwen/Qwen2.5-VL-3B-Instruct", "modelUrl": "https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct", - "hasIcon": false, + "iconPath": "_static/img/qwen-logo.svg", "printed": 67.8366, "handwritten": 42.8576, "total": 54.3864 @@ -174,7 +213,8 @@ "modelName": "O4 Mini", "modelId": "o4-mini-2025-04-16", "modelUrl": "https://platform.openai.com/docs/models/o4-mini", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 65.5032, "handwritten": 41.4801, "total": 52.5677 @@ -183,7 +223,7 @@ "modelName": "Claude Opus 4.1", "modelId": "claude-opus-4-1-20250805", "modelUrl": "https://docs.claude.com/en/docs/about-claude/models/overview", - "hasIcon": false, + "iconPath": "_static/img/claude-symbol.svg", "printed": 66.6052, "handwritten": 40.1808, "total": 52.3767 @@ -192,7 +232,8 @@ "modelName": "GPT-4.1", "modelId": "gpt-4.1-2025-04-14", "modelUrl": "https://platform.openai.com/docs/models/gpt-4.1", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 64.4017, "handwritten": 41.4127, "total": 52.023 @@ -201,7 +242,7 @@ "modelName": "Claude Sonnet 4", "modelId": "claude-sonnet-4-20250514", "modelUrl": "https://docs.claude.com/en/docs/about-claude/models/overview", - "hasIcon": false, + "iconPath": "_static/img/claude-symbol.svg", "printed": 62.2358, "handwritten": 37.0701, "total": 48.685 @@ -210,34 +251,55 @@ "modelName": "O1", "modelId": "o1-2024-12-17", "modelUrl": "https://platform.openai.com/docs/models/o1", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 62.4692, "handwritten": 35.0806, "total": 47.7215 }, + { + "modelName": "PaddleOCR-VL 1.5", + "modelId": "PaddlePaddle/PaddleOCR-VL-1.5", + "modelUrl": "https://huggingface.co/PaddlePaddle/PaddleOCR-VL-1.5", + "iconPath": null, + "printed": 62.6, + "handwritten": 33.6, + "total": 47.0 + }, { "modelName": "O3", "modelId": "o3-2025-04-16", "modelUrl": "https://platform.openai.com/docs/models/o3", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 62.65, "handwritten": 30.9051, "total": 45.5566 }, { - "modelName": "Mistral OCR", - "modelId": null, + "modelName": "Mistral OCR 2", + "modelId": "mistral-ocr-2505", + "modelUrl": null, + "iconPath": "_static/img/mistral-logo.svg", + "printed": 64.3, + "handwritten": 30.4, + "total": 46.0 + }, + { + "modelName": "Mistral OCR 3", + "modelId": "mistral-ocr-2512", "modelUrl": null, - "hasIcon": false, - "printed": 64.0418, - "handwritten": 29.4388, - "total": 45.4094 + "iconPath": "_static/img/mistral-logo.svg", + "printed": 71.5, + "handwritten": 47.6, + "total": 58.6 }, { "modelName": "GPT-5", "modelId": "gpt-5-2025-08-07", "modelUrl": "https://platform.openai.com/docs/models/gpt-5", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 60.7442, "handwritten": 30.4073, "total": 44.409 @@ -246,7 +308,8 @@ "modelName": "GPT-4o", "modelId": "gpt-4o-2024-11-20", "modelUrl": "https://platform.openai.com/docs/models/gpt-4o", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 56.2807, "handwritten": 34.1932, "total": 44.3875 @@ -255,7 +318,7 @@ "modelName": "MiMo VL", "modelId": "XiaomiMiMo/MiMo-VL-7B-RL-2508", "modelUrl": "https://huggingface.co/XiaomiMiMo/MiMo-VL-7B-RL-2508", - "hasIcon": false, + "iconPath": null, "printed": 54.8469, "handwritten": 34.6299, "total": 43.9608 @@ -264,16 +327,26 @@ "modelName": "Gemma 3 (27B)", "modelId": "google/gemma-3-27b-it", "modelUrl": "https://huggingface.co/google/gemma-3-27b-it", - "hasIcon": false, + "iconPath": "_static/img/gemma-icon.svg", "printed": 55.3292, "handwritten": 34.1056, "total": 43.9011 }, + { + "modelName": "Qwen 3.5 (2B)", + "modelId": "Qwen/Qwen3.5-2B", + "modelUrl": "https://huggingface.co/Qwen/Qwen3.5-2B", + "iconPath": "_static/img/qwen-logo.svg", + "printed": 57.6, + "handwritten": 31.6, + "total": 43.6 + }, { "modelName": "GPT-4o Mini", "modelId": "gpt-4o-mini-2024-07-18", "modelUrl": "https://platform.openai.com/docs/models/gpt-4o-mini", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 52.8335, "handwritten": 29.7653, "total": 40.4121 @@ -282,7 +355,7 @@ "modelName": "MiniCPM-V 4.5", "modelId": "openbmb/MiniCPM-V-4_5", "modelUrl": "https://huggingface.co/openbmb/MiniCPM-V-4_5", - "hasIcon": false, + "iconPath": null, "printed": 49.7495, "handwritten": 32.0729, "total": 40.2314 @@ -291,16 +364,26 @@ "modelName": "GPT-4.1 Nano", "modelId": "gpt-4.1-nano-2025-04-14", "modelUrl": "https://platform.openai.com/docs/models/gpt-4.1-nano", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 51.8564, "handwritten": 28.3488, "total": 39.1984 }, + { + "modelName": "DeepSeek OCR 2", + "modelId": "deepseek-ai/DeepSeek-OCR-2", + "modelUrl": "https://huggingface.co/deepseek-ai/DeepSeek-OCR-2", + "iconPath": null, + "printed": 56.1, + "handwritten": 20.0, + "total": 36.7 + }, { "modelName": "Skywork R1V3", "modelId": "Skywork/Skywork-R1V3-38B", "modelUrl": "https://huggingface.co/Skywork/Skywork-R1V3-38B", - "hasIcon": false, + "iconPath": null, "printed": 42.4283, "handwritten": 25.619, "total": 33.3771 @@ -309,16 +392,25 @@ "modelName": "InternVL 3.5 (30B-A3B)", "modelId": "OpenGVLab/InternVL3_5-30B-A3B", "modelUrl": "https://huggingface.co/OpenGVLab/InternVL3_5-30B-A3B", - "hasIcon": false, + "iconPath": null, "printed": 35.8059, "handwritten": 26.3691, "total": 30.7246 }, + { + "modelName": "Qwen 3.5 (0.8B)", + "modelId": "Qwen/Qwen3.5-0.8B", + "modelUrl": "https://huggingface.co/Qwen/Qwen3.5-0.8B", + "iconPath": "_static/img/qwen-logo.svg", + "printed": 41.7, + "handwritten": 16.9, + "total": 28.4 + }, { "modelName": "R", "modelId": "YannQi/R-4B", "modelUrl": "https://huggingface.co/YannQi/R-4B", - "hasIcon": false, + "iconPath": null, "printed": 32.6874, "handwritten": 21.735, "total": 26.79 @@ -327,7 +419,8 @@ "modelName": "GPT-5 Nano", "modelId": "gpt-5-nano-2025-08-07", "modelUrl": "https://platform.openai.com/docs/models/gpt-5-nano", - "hasIcon": false, + "iconPath": "_static/img/openai-symbol.svg", + "iconInvertDark": true, "printed": 38.7162, "handwritten": 14.165, "total": 25.4963 @@ -336,16 +429,25 @@ "modelName": "Nemotron Nano VL", "modelId": "nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1", "modelUrl": "https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1", - "hasIcon": false, + "iconPath": "_static/img/nvidia-logo.svg", "printed": 29.3383, "handwritten": 15.2615, "total": 21.7585 }, + { + "modelName": "dots.ocr 1.5", + "modelId": "kristaller486/dots.ocr-1.5", + "modelUrl": "https://huggingface.co/kristaller486/dots.ocr-1.5", + "iconPath": null, + "printed": 22.2, + "handwritten": 10.4, + "total": 15.9 + }, { "modelName": "Phi 4 Multimodal", "modelId": "microsoft/Phi-4-multimodal-instruct", "modelUrl": "https://huggingface.co/microsoft/Phi-4-multimodal-instruct", - "hasIcon": false, + "iconPath": null, "printed": 9.0269, "handwritten": 5.3393, "total": 7.0413 diff --git a/docs/_static/css/custom.css b/docs/_static/css/custom.css index b750fbf..df6d205 100644 --- a/docs/_static/css/custom.css +++ b/docs/_static/css/custom.css @@ -112,6 +112,10 @@ html[data-theme="dark"] { z-index: 1; } +.benchmark-table tbody tr:nth-child(even) { + background: rgba(139, 69, 31, 0.03); +} + .benchmark-sort-button { align-items: center; appearance: none; @@ -150,11 +154,46 @@ html[data-theme="dark"] { gap: 0.8rem; } -.benchmark-model-icon { - border-radius: 0.4rem; - height: 1.65rem; +.benchmark-model-icon-frame { + align-items: center; + display: inline-flex; + flex: 0 0 auto; + height: 1.9rem; + justify-content: center; margin-top: 0.1rem; - width: 1.65rem; + width: auto; +} + +.benchmark-model-icon { + background: transparent; + border-radius: 0; + display: block; + height: 100%; + max-height: 100%; + max-width: 3rem; + width: auto; +} + +.benchmark-model-icon-placeholder { + align-items: center; + background: rgba(139, 69, 31, 0.06); + border: 1px dashed rgba(139, 69, 31, 0.22); + border-radius: 999px; + box-sizing: border-box; + color: var(--pst-color-text-muted); + display: inline-flex; + height: 1.9rem; + justify-content: center; + width: 1.9rem; +} + +.benchmark-model-icon-placeholder::before { + background: currentColor; + border-radius: 999px; + content: ""; + height: 0.34rem; + opacity: 0.45; + width: 0.34rem; } .benchmark-model-text { @@ -190,6 +229,30 @@ html[data-theme="dark"] .benchmark-table th { background: rgba(244, 178, 94, 0.06); } +html[data-theme="dark"] .benchmark-table tbody tr:nth-child(even) { + background: rgba(244, 178, 94, 0.05); +} + +html[data-theme="dark"] .benchmark-model-icon-placeholder, +[data-bs-theme="dark"] .benchmark-model-icon-placeholder { + background: rgba(244, 178, 94, 0.08); + border-color: rgba(244, 178, 94, 0.24); +} + +html[data-theme="dark"] .bd-content .benchmark-model-icon, +html[data-theme="dark"] .bd-content .benchmark-model-icon-frame, +[data-bs-theme="dark"] .bd-content .benchmark-model-icon, +[data-bs-theme="dark"] .bd-content .benchmark-model-icon-frame { + background: transparent; + box-shadow: none; + border-radius: 0; +} + +html[data-theme="dark"] .bd-content .benchmark-model-icon--invert-dark, +[data-bs-theme="dark"] .bd-content .benchmark-model-icon--invert-dark { + filter: invert(1); +} + @media (min-width: 1200px) { .bd-page-width { max-width: 96rem; diff --git a/docs/_static/img/ai2-symbol.svg b/docs/_static/img/ai2-symbol.svg new file mode 100644 index 0000000..cf2af97 --- /dev/null +++ b/docs/_static/img/ai2-symbol.svg @@ -0,0 +1,4 @@ + + + + diff --git a/docs/_static/img/azure-logo.png b/docs/_static/img/azure-logo.png new file mode 100644 index 0000000..23f6803 Binary files /dev/null and b/docs/_static/img/azure-logo.png differ diff --git a/docs/_static/img/claude-symbol.svg b/docs/_static/img/claude-symbol.svg new file mode 100644 index 0000000..879ad81 --- /dev/null +++ b/docs/_static/img/claude-symbol.svg @@ -0,0 +1,7 @@ + + + + + + + diff --git a/docs/_static/img/gemini-symbol.png b/docs/_static/img/gemini-symbol.png new file mode 100644 index 0000000..5141463 Binary files /dev/null and b/docs/_static/img/gemini-symbol.png differ diff --git a/docs/_static/img/gemma-icon.svg b/docs/_static/img/gemma-icon.svg new file mode 100644 index 0000000..ed81051 --- /dev/null +++ b/docs/_static/img/gemma-icon.svg @@ -0,0 +1 @@ +Gemma \ No newline at end of file diff --git a/docs/_static/img/mistral-logo.svg b/docs/_static/img/mistral-logo.svg new file mode 100644 index 0000000..9ea7556 --- /dev/null +++ b/docs/_static/img/mistral-logo.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/docs/_static/img/nvidia-logo.svg b/docs/_static/img/nvidia-logo.svg new file mode 100644 index 0000000..9393002 --- /dev/null +++ b/docs/_static/img/nvidia-logo.svg @@ -0,0 +1,32 @@ + + + + +generated by pstoedit version:3.44 from NVBadge_2D.eps + + + + diff --git a/docs/_static/img/openai-symbol.svg b/docs/_static/img/openai-symbol.svg new file mode 100644 index 0000000..74d9b1b --- /dev/null +++ b/docs/_static/img/openai-symbol.svg @@ -0,0 +1,5 @@ + + + + + \ No newline at end of file diff --git a/docs/_static/img/qwen-logo.svg b/docs/_static/img/qwen-logo.svg new file mode 100644 index 0000000..64be10b --- /dev/null +++ b/docs/_static/img/qwen-logo.svg @@ -0,0 +1,15 @@ + + + + + + + + + + + + + + + diff --git a/docs/_static/img/reducto-logomark.svg b/docs/_static/img/reducto-logomark.svg new file mode 100644 index 0000000..642c73e --- /dev/null +++ b/docs/_static/img/reducto-logomark.svg @@ -0,0 +1,3 @@ + + + diff --git a/docs/_static/js/benchmark-leaderboard.js b/docs/_static/js/benchmark-leaderboard.js index f0beee2..240a202 100644 --- a/docs/_static/js/benchmark-leaderboard.js +++ b/docs/_static/js/benchmark-leaderboard.js @@ -10,6 +10,10 @@ { key: "handwritten", label: "Handwritten", numeric: true }, { key: "total", label: "Total", numeric: true }, ]; + const NUMERIC_COLUMNS = new Set( + COLUMN_DEFINITIONS.filter((column) => column.numeric).map((column) => column.key), + ); + const DEFAULT_SORT_STATE = { key: "total", direction: "desc" }; function contentRoot() { return document.documentElement.dataset.content_root || ""; @@ -32,15 +36,42 @@ return SCORE_FORMATTER.format(value); } + function createIconFrame(row) { + const iconFrame = document.createElement("span"); + iconFrame.className = "benchmark-model-icon-frame"; + + if (!row.iconPath) { + const placeholder = document.createElement("span"); + placeholder.className = "benchmark-model-icon-placeholder"; + placeholder.setAttribute("aria-hidden", "true"); + iconFrame.append(placeholder); + return iconFrame; + } + + const icon = document.createElement("img"); + icon.className = "benchmark-model-icon"; + if (row.iconInvertDark) { + icon.classList.add("benchmark-model-icon--invert-dark"); + } + icon.alt = `${row.modelName} icon`; + icon.src = resolvePath(row.iconPath); + iconFrame.append(icon); + + return iconFrame; + } + function compareRows(left, right, sortState) { const { key, direction } = sortState; const multiplier = direction === "asc" ? 1 : -1; const leftValue = left[key]; const rightValue = right[key]; - if (typeof leftValue === "number" && typeof rightValue === "number") { - if (leftValue !== rightValue) { - return (leftValue - rightValue) * multiplier; + if (NUMERIC_COLUMNS.has(key)) { + const leftNumber = Number(leftValue); + const rightNumber = Number(rightValue); + + if (!Number.isNaN(leftNumber) && !Number.isNaN(rightNumber) && leftNumber !== rightNumber) { + return (leftNumber - rightNumber) * multiplier; } return left.modelName.localeCompare(right.modelName); } @@ -48,17 +79,11 @@ return String(leftValue || "").localeCompare(String(rightValue || "")) * multiplier; } - function createModelCell(row, logoPath) { + function createModelCell(row) { const wrapper = document.createElement("div"); wrapper.className = "benchmark-model"; - if (row.hasIcon) { - const icon = document.createElement("img"); - icon.className = "benchmark-model-icon"; - icon.alt = `${row.modelName} icon`; - icon.src = logoPath; - wrapper.append(icon); - } + wrapper.append(createIconFrame(row)); const textBlock = document.createElement("div"); textBlock.className = "benchmark-model-text"; @@ -104,7 +129,7 @@ return button; } - function renderLeaderboard(container, rows, sortState, logoPath) { + function renderLeaderboard(container, rows, sortState) { container.replaceChildren(); const sortedRows = [...rows].sort((left, right) => compareRows(left, right, sortState)); @@ -130,7 +155,7 @@ sortState.key = key; sortState.direction = key === "modelName" ? "asc" : "desc"; } - renderLeaderboard(container, rows, sortState, logoPath); + renderLeaderboard(container, rows, sortState); }; for (const column of COLUMN_DEFINITIONS) { @@ -147,7 +172,7 @@ sortedRows.forEach((row, index) => { const tr = document.createElement("tr"); - if (row.hasIcon) { + if (row.iconPath) { tr.classList.add("is-featured"); } @@ -157,7 +182,7 @@ tr.append(rankCell); const modelCell = document.createElement("td"); - modelCell.append(createModelCell(row, logoPath)); + modelCell.append(createModelCell(row)); tr.append(modelCell); for (const key of ["printed", "handwritten", "total"]) { @@ -184,8 +209,6 @@ } async function initializeLeaderboard(container) { - const logoPath = resolvePath("_static/img/churro.png"); - try { const response = await fetch(resolvePath("_static/data/benchmark_results.json")); if (!response.ok) { @@ -193,8 +216,8 @@ } const rows = await response.json(); - const sortState = { key: "total", direction: "desc" }; - renderLeaderboard(container, rows, sortState, logoPath); + const sortState = { ...DEFAULT_SORT_STATE }; + renderLeaderboard(container, rows, sortState); } catch (error) { renderError(container, "Unable to load the benchmark leaderboard data."); console.error("[benchmark-leaderboard] failed to initialize", error); diff --git a/docs/benchmarking.md b/docs/benchmarking.md index f6a8ea8..f21609a 100644 --- a/docs/benchmarking.md +++ b/docs/benchmarking.md @@ -1,10 +1,9 @@ # Benchmarking -This page is for reproducing CHURRO-DS benchmark runs from a repo checkout. +For the official leaderboard results, see the [Benchmark Leaderboard](leaderboard.md). -Run these commands from the repo root after `pixi install`. +This page describes how to benchmark your own model on [CHURRO-DS](https://huggingface.co/datasets/stanford-oval/churro-dataset). Please open a pull request if you would like to add your model to the official leaderboard. -For the current committed benchmark snapshot, see the [Benchmark Leaderboard](leaderboard.md). Contributor setup, test commands, and package checks live in [Contributing](contributing.md). ## Smallest Useful Run @@ -18,6 +17,7 @@ pixi run python -m tooling.benchmarking.benchmark \ ``` By default, results are written under `workdir/results//`. +The evaluation pipeline strips the default OCR wrapper tag, flattens supported XML-like OCR output, normalizes whitespace and punctuation, and applies additional Arabic normalization for languages with Arabic script (Arabic and Persian). ## Common Flags @@ -27,31 +27,31 @@ By default, results are written under `workdir/results//`. - `--language` and `--document-type`: filter the benchmark subset before slicing - `--output-dir PATH`: override the default results directory - `--max-concurrency N`: cap the number of in-flight OCR requests -- `--vllm-gpu-memory-utilization` and `--vllm-cpu-offload-gb`: pass through selected vLLM runtime knobs ## Output Files -Each benchmark run writes one result directory. The important files are: +Each benchmark run writes one result directory. The directory contains two JSON files: - `outputs.json`: one row per evaluated page with the raw predicted text, gold text, and page-level metrics - `all_metrics.json`: aggregate metrics grouped across the full run, by main language, by document type, and by the language/type combination -`outputs.json` stores page-level metric values as raw fractions. `all_metrics.json` converts aggregate values to percentages and rounds them to one decimal place. The evaluation pipeline strips the default OCR wrapper tag, flattens supported XML-like OCR output, normalizes whitespace and punctuation, and applies additional Arabic normalization for Arabic and Persian examples. ## Filtering And Slicing -Subset filters are applied before offset and limit: +You can run benchmarks on subsets of the data by combining `--language`, `--document-type`, `--offset`, and `--input-size`. The filters are applied in the following order: - `--language` filters on `main_language` - `--document-type` filters on `document_type` - `--offset` skips rows after filtering - `--input-size` limits rows after filtering and offset -That means `--language Arabic --offset 100 --input-size 50` selects rows 101 to 150 from the Arabic-only subset, not from the full split. +That means for example `--language Arabic --offset 100 --input-size 50` selects rows 101 to 150 from the Arabic-only subset, not from the full split. ## Example Commands +If you want to benchmark a model using vLLM, run a vLLM server separately and point `--backend openai-compatible` at its OpenAI-compatible endpoint. See the [official vLLM serving docs](https://docs.vllm.ai/en/stable/serving/openai_compatible_server.html). + | Model | Model ID | Backend | Full command | | --- | --- | --- | --- | | Gemini 2.5 Pro | `vertex_ai/gemini-2.5-pro` | `litellm` | `pixi run python -m tooling.benchmarking.benchmark --backend litellm --dataset-split test --model vertex_ai/gemini-2.5-pro --output-dir workdir/results/test/litellm_vertex_ai_gemini-2.5-pro` | -| Qwen 3.5-0.8B | `Qwen/Qwen3.5-0.8B` | `vllm` | `pixi run python -m tooling.benchmarking.benchmark --backend vllm --dataset-split test --model Qwen/Qwen3.5-0.8B --output-dir workdir/results/test/vllm_Qwen_Qwen3.5-0.8B` | +| Qwen 3.5-0.8B | `Qwen/Qwen3.5-0.8B` | `openai-compatible` | `pixi run python -m tooling.benchmarking.benchmark --backend openai-compatible --dataset-split test --model Qwen/Qwen3.5-0.8B --base-url http://127.0.0.1:8000/v1 --output-dir workdir/results/test/openai-compatible_Qwen_Qwen3.5-0.8B` | diff --git a/docs/cli.md b/docs/cli.md index d47f119..4c85737 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -3,14 +3,26 @@ Use the CLI when you want a quick sanity check before writing Python code. Use `churro-ocr --help` or `python -m churro_ocr --help` to inspect the top-level commands. +Install Churro in [Getting Started](getting-started.md), and use +[Providers And Configuration](guides/providers.md) +for backend-specific runtime setup. ## Command Summary | Command | Use it when | | --- | --- | +| `install` | you want Churro to install an optional runtime into the active UV environment | | `transcribe` | you want OCR text for one image | | `extract-pages` | you want page crops from an image or PDF | +## `install` Examples + +### Install Local Transformers OCR + +```bash +churro-ocr install hf +``` + ## `transcribe` Examples ### OCR One Image @@ -29,10 +41,11 @@ churro-ocr transcribe \ --image scan.png \ --backend openai-compatible \ --model local-model \ - --base-url http://127.0.0.1:8000/v1 \ - --api-key dummy + --base-url http://127.0.0.1:8000/v1 ``` +For vLLM, serve the model separately with its OpenAI-compatible server and then use this same `openai-compatible` route. See the [official vLLM serving docs](https://docs.vllm.ai/en/stable/serving/openai_compatible_server.html). + ## `extract-pages` Examples ### Extract Pages From An Image @@ -73,11 +86,10 @@ churro-ocr extract-pages \ | `--backend` value | Required flags | Notes | | --- | --- | --- | | `litellm` | `--model` | Uses LiteLLM credentials and routing. `--base-url`, `--api-key`, and `--api-version` are optional transport overrides. | -| `openai-compatible` | `--model`, `--base-url`, `--api-key` | For local or self-hosted OpenAI-style servers. | +| `openai-compatible` | `--model`, `--base-url` | For local or self-hosted OpenAI-style servers. `--api-key` is optional. | | `azure` | `--endpoint`, `--api-key` | `--model` is optional. | -| `mistral` | `--api-key` | `--model` defaults to `mistral-ocr-latest`. | +| `mistral` | `--api-key`, `--model` | `--model` must be either `mistral-ocr-2505` or `mistral-ocr-2512`. | | `hf` | `--model` | Local Transformers OCR. | -| `vllm` | `--model` | Local vLLM OCR. | ### `extract-pages` Detectors diff --git a/docs/getting-started.md b/docs/getting-started.md index 1dbf321..92cef8b 100644 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -11,36 +11,31 @@ | Run an end-to-end image or PDF OCR workflow | `DocumentOCRPipeline` | | Tune provider options directly | `build_ocr_backend(...)` + `OCRBackendSpec` | -## Install Only What You Need +## Install + +Use UV as the supported install path. ```bash -pip install churro-ocr -pip install "churro-ocr[llm]" -pip install "churro-ocr[local]" -pip install "churro-ocr[hf]" -pip install "churro-ocr[vllm]" -pip install "churro-ocr[azure]" -pip install "churro-ocr[mistral]" -pip install "churro-ocr[pdf]" -pip install "churro-ocr[all]" +uv tool install churro-ocr +# or, in a project: +uv add churro-ocr ``` -## Provider Extras +Then install the runtime for the backend you plan to use: -| Extra | Use it when | -| --- | --- | -| `llm` | you want hosted multimodal OCR and LLM-based page detection through LiteLLM | -| `local` | you have a local or self-hosted OpenAI-compatible server | -| `hf` | you want local Transformers inference in-process | -| `vllm` | you want higher-throughput local serving | -| `azure` | you want Azure Document Intelligence OCR or layout detection | -| `mistral` | you want Mistral OCR | -| `pdf` | you want PDF rasterization through `pypdfium2` | -| `all` | you want every supported backend and utility extra | +```bash +uv run churro-ocr install llm +uv run churro-ocr install hf +uv run churro-ocr install local +``` + +If you installed the CLI with `uv tool install churro-ocr`, drop the `uv run` prefix. +For the full provider/runtime matrix, use [Providers And Configuration](guides/providers.md). ## First OCR Example Use `OCRClient` when your input is already one page per image. +This example uses `provider="litellm"`, so install the `llm` runtime first. ```python from churro_ocr.ocr import OCRClient @@ -67,6 +62,8 @@ When an API accepts both `image` and `image_path`, pass exactly one of them. Use the CLI when you want to confirm a model or backend before writing Python code. ```bash +uv tool install churro-ocr +churro-ocr install hf churro-ocr transcribe \ --image scan.png \ --backend hf \ diff --git a/docs/guides/providers.md b/docs/guides/providers.md index 08e5b6f..1b061bb 100644 --- a/docs/guides/providers.md +++ b/docs/guides/providers.md @@ -15,14 +15,16 @@ backend = build_ocr_backend( ## Which OCR Backend Should You Use? -| Provider | Install extra | Good default when | +Install the base package first as shown in [Getting Started](../getting-started.md). +This page is the source of truth for matching providers to runtime targets. + +| Provider | Install command | Good default when | | --- | --- | --- | -| `litellm` | `llm` | you want hosted multimodal models routed through LiteLLM | -| `openai-compatible` | `local` | you have a local or self-hosted OpenAI-style server | -| `hf` | `hf` | you want local Transformers inference in-process | -| `vllm` | `vllm` | you want higher-throughput local serving | -| `azure` | `azure` | you want Azure Document Intelligence OCR | -| `mistral` | `mistral` | you want Mistral OCR | +| `litellm` | `churro-ocr install llm` | you want hosted multimodal models routed through LiteLLM | +| `openai-compatible` | `churro-ocr install local` | you have a local or self-hosted OpenAI-style server | +| `hf` | `churro-ocr install hf` | you want local Transformers inference in-process | +| `azure` | `churro-ocr install azure` | you want Azure Document Intelligence OCR | +| `mistral` | `churro-ocr install mistral` | you want Mistral OCR | ## Recommended Starting Points @@ -30,7 +32,8 @@ backend = build_ocr_backend( | --- | --- | --- | | hosted OCR | `litellm` + `vertex_ai/gemini-2.5-flash` | easiest hosted path with the standard builder interface | | local OCR | `hf` + `stanford-oval/churro-3B` | first-party local model support in-process | -| higher-throughput local serving | `vllm` + `stanford-oval/churro-3B` | better fit when you want a served local backend | +| layout-heavy local OCR | `hf` + `datalab-to/chandra-ocr-2` | built-in profile matches Chandra's layout-oriented prompt, scaling, and generation defaults | +| higher-throughput local serving | `openai-compatible` + your own OpenAI-style server | good when you already run a served local backend such as vLLM | ## Hosted Providers @@ -94,7 +97,7 @@ from churro_ocr.providers import MistralOptions, OCRBackendSpec, build_ocr_backe backend = build_ocr_backend( OCRBackendSpec( provider="mistral", - model="mistral-ocr-latest", + model="mistral-ocr-2512", options=MistralOptions(api_key=""), ) ) @@ -102,6 +105,8 @@ backend = build_ocr_backend( ## Local And Self-Hosted Providers +Before using a local or self-hosted provider, install the matching runtime from the table above. + ### OpenAI-compatible ```python @@ -117,12 +122,13 @@ backend = build_ocr_backend( model="local-model", transport=LiteLLMTransportConfig( api_base="http://127.0.0.1:8000/v1", - api_key="dummy", ), ) ) ``` +If you want to use vLLM, serve it separately and point this backend at that server's OpenAI-compatible endpoint. See the [official vLLM serving docs](https://docs.vllm.ai/en/stable/serving/openai_compatible_server.html). + ### Hugging Face ```python @@ -139,26 +145,14 @@ backend = build_ocr_backend( ) ``` -### vLLM - -```python -from churro_ocr.providers import OCRBackendSpec, VLLMOptions, build_ocr_backend - -backend = build_ocr_backend( - OCRBackendSpec( - provider="vllm", - model="stanford-oval/churro-3B", - options=VLLMOptions(), - ) -) -``` +Built-in model-specific profiles are resolved automatically for known models such as `stanford-oval/churro-3B`, `datalab-to/chandra-ocr-2`, `deepseek-ai/DeepSeek-OCR-2`, `kristaller486/dots.ocr-1.5`, `rednote-hilab/dots.mocr`, and the supported `olmOCR` checkpoints. ## `OCRBackendSpec` Reference | Field | Meaning | | --- | --- | -| `provider` | One of `litellm`, `openai-compatible`, `azure`, `mistral`, `hf`, or `vllm`. | -| `model` | Required for `litellm`, `openai-compatible`, `hf`, and `vllm`. Optional for `azure`. Defaults to `mistral-ocr-latest` when omitted for `mistral`. | +| `provider` | One of `litellm`, `openai-compatible`, `azure`, `mistral`, or `hf`. | +| `model` | Required for `litellm`, `openai-compatible`, `mistral`, and `hf`. Optional for `azure`. For `mistral`, use one of `mistral-ocr-2505` or `mistral-ocr-2512`. | | `profile` | `None`, a built-in profile name, or a custom `OCRModelProfile`. | | `transport` | Shared request transport config for LiteLLM-based providers. | | `options` | Provider-specific dataclass matching `provider`. | @@ -167,12 +161,11 @@ backend = build_ocr_backend( | Type | Used by | Required fields | Notes | | --- | --- | --- | --- | -| `LiteLLMTransportConfig` | `litellm`, `openai-compatible`, `LLMPageDetector` | None at the dataclass level | Use this for transport, credentials, and completion settings. | +| `LiteLLMTransportConfig` | `litellm`, `openai-compatible`, `LLMPageDetector` | None at the dataclass level | Use this for transport, credentials, and completion settings. `api_base` is required for `openai-compatible`; `api_key` is optional. | | `OpenAICompatibleOptions` | `openai-compatible` | None | Use `model_prefix` when your local server expects a provider prefix. | | `HuggingFaceOptions` | `hf` | None | Carries runtime, processor, generation, and template options. | -| `VLLMOptions` | `vllm` | None | Carries runtime and sampling settings for vLLM. | | `AzureDocumentIntelligenceOptions` | `azure` | `endpoint`, `api_key` | `model` is optional for Azure OCR in `OCRBackendSpec`. | -| `MistralOptions` | `mistral` | `api_key` | `model` defaults to `mistral-ocr-latest` when omitted. | +| `MistralOptions` | `mistral` | `api_key` | `model` is required and must be `mistral-ocr-2505` or `mistral-ocr-2512`. | ## Advanced Customization @@ -214,7 +207,10 @@ Useful public template exports: | `HFChatTemplate` | `churro_ocr.templates` | Build a Hugging Face chat-style multimodal prompt. | | `DEFAULT_OCR_TEMPLATE` | `churro_ocr.templates` | Generic OCR prompt template used by the default model profile. | | `CHURRO_3B_XML_TEMPLATE` | `churro_ocr.templates` | Built-in template for `stanford-oval/churro-3B`. | +| `CHANDRA_OCR_2_OCR_TEMPLATE` | `churro_ocr.templates` | Built-in template for `datalab-to/chandra-ocr-2`. | +| `DEEPSEEK_OCR_2_OCR_TEMPLATE` | `churro_ocr.templates` | Built-in template for `deepseek-ai/DeepSeek-OCR-2`. | | `DOTS_OCR_1_5_OCR_TEMPLATE` | `churro_ocr.templates` | Built-in template for `kristaller486/dots.ocr-1.5`. | +| `DOTS_MOCR_OCR_TEMPLATE` | `churro_ocr.templates` | Built-in template for `rednote-hilab/dots.mocr`. | | `OCRPromptTemplate` | `churro_ocr.templates` | Base protocol for custom profile integration. | Useful public prompt exports: @@ -224,6 +220,8 @@ Useful public prompt exports: | `DEFAULT_OCR_SYSTEM_PROMPT` | `churro_ocr.prompts` | Default system instruction for generic OCR prompting. | | `DEFAULT_OCR_USER_PROMPT` | `churro_ocr.prompts` | Default user prompt for plain OCR output. | | `DEFAULT_MARKDOWN_OCR_USER_PROMPT` | `churro_ocr.prompts` | Default user prompt when markdown-style OCR output is preferred. | +| `CHANDRA_OCR_LAYOUT_PROMPT` | `churro_ocr.prompts` | Upstream Chandra OCR 2 layout-block HTML prompt. | | `DEFAULT_OCR_OUTPUT_TAG` | `churro_ocr.prompts` | Shared tag name used by the default OCR postprocessor. | | `DEFAULT_BOUNDARY_DETECTION_PROMPT` | `churro_ocr.prompts` | Default prompt used by LLM-based page and text-block boundary detection helpers. | +| `parse_chandra_response(...)` | `churro_ocr.prompts` | Convert Chandra HTML-layout output to plain text and preserve raw HTML metadata. | | `strip_ocr_output_tag(...)` | `churro_ocr.prompts` | Remove the default OCR wrapper tag from model output. | diff --git a/docs/pypi.md b/docs/pypi.md index 6d55bf2..28e1359 100644 --- a/docs/pypi.md +++ b/docs/pypi.md @@ -6,20 +6,19 @@ Full documentation and project overview live at https://stanford-oval.github.io/ ## Install -Install only the pieces you need: +Use UV as the supported install path. ```bash -pip install churro-ocr -pip install "churro-ocr[llm]" -pip install "churro-ocr[local]" -pip install "churro-ocr[hf]" -pip install "churro-ocr[vllm]" -pip install "churro-ocr[azure]" -pip install "churro-ocr[mistral]" -pip install "churro-ocr[pdf]" -pip install "churro-ocr[all]" +uv tool install churro-ocr +# or, in a project: +uv add churro-ocr ``` +Runtime setup and provider-specific install commands are in +[Getting Started](https://stanford-oval.github.io/Churro/getting-started.html) +and +[Providers And Configuration](https://stanford-oval.github.io/Churro/guides/providers.html). + ## Which API Should You Use? | Goal | API | @@ -31,6 +30,8 @@ pip install "churro-ocr[all]" ## Quick Start +This example assumes you already installed the runtime for the provider you want to use. + ```python from churro_ocr.ocr import OCRClient from churro_ocr.providers import OCRBackendSpec, build_ocr_backend diff --git a/pixi.lock b/pixi.lock index edf4dd3..3beddc1 100644 --- a/pixi.lock +++ b/pixi.lock @@ -62,7 +62,6 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h366c992_103.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - - pypi: https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/8d/3f/95338030883d8c8b91223b4e21744b04d11b161a3ef117295d8241f50ab4/accessible_pygments-0.0.5-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0f/15/5bf3b99495fb160b63f95972b81750f18f7f4e02ad051373b669d17d44f2/aiohappyeyeballs-2.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/bd/c9/989f4034fb46841208de7aeeac2c6d8300745ab4f28c42f629ba77c2d916/aiohttp-3.13.5-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl @@ -86,9 +85,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/22/e5/06b1f88f42a5a99df42ce61208bdec3bddb3d261412874280a19796fc09c/coverage-7.13.5-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/34/71/1ea5a7352ae516d5512d17babe7e1b87d9db5150b21f794b1377eac1edc0/cryptography-46.0.6-cp311-abi3-manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/ee/5e/c0fe77a73aaefd3fff25ffaccaac69c5a63eafdf8b9a4c476626ef0ac703/cuda_bindings-13.2.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c4/74/8c66861b873d8eed51fde56d3091baa4906a56f0d4390cae991f2d41dda5/cuda_pathfinder-1.5.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b0/e5/247d094108e42ac26363ab8dc57f168840cf7c05774b40ffeb0d78868fcc/datasets-4.8.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl @@ -107,7 +103,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/df/9a/a24b26dc8a65f0ecc0fe5be981a19e61e7ca963b85e062c083f3a9100529/hf_xet-1.4.3-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/a8/af/48ac8483240de756d2438c380746e7130d1c6f75802ef22f3c6d49982787/huggingface_hub-0.36.2-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/73/37/0d15d16150e1829f3e90962c99f28257f6de9e526a680b4c6f5acdb54fd2/huggingface_hub-1.9.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/42/77/de194443bf38daed9452139e960c632b0ef9f9a5dd9ce605fdf18ca9f1b1/id-1.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/46/33/92ef41c6fad0233e41d3d84ba8e8ad18d1780f1e5d99b3c683e6d7f98b63/identify-2.6.18-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl @@ -134,30 +130,13 @@ environments: - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/c9/f9/98d825105c450b9c67c27026caa374112b7e466c18331601d02ca278a01b/mistralai-1.12.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d8/f4/5e52c7319b8087acef603ed6e50dc325c02eaa999355414830468611f13c/more_itertools-11.0.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fe/3b/8ec5074bcfc450fe84273713b4b0a0dd47c0249358f5d82eb8104ffe2520/multidict-6.7.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/a0/61/af9115673a5870fd885247e2f1b68c4f1197737da315b520a91c757a861a/multiprocess-0.70.19-py314-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/df/76d0321c3797b54b60fef9ec3bd6f4cfd124b9e422182156a1dd418722cf/myst_parser-4.0.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d6/b7/ec1cbc6b297a808c513f59f501656389623fc09ad6a58c640851289c7854/nh3-0.3.4-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/9d/91/04e965f8e717ba0ab4bdca5c112deeab11c9e750d94c4d4602f050295d39/nltk-3.9.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/98/7c/21252050676612625449b4807d6b695b9ce8a7c9e1c197ee6216c8a65c7c/numpy-2.4.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a3/22/0b4b932655d17a6da1b92fa92ab12844b053bb2ac2475e179ba6f043da1e/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/2a/9e/5bfa2270f902d5b92ab7d41ce0475b8630572e71e349b2a4996d14bdda93/openai-2.30.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/bf/93795954016c522008da367da292adceed71cca6ee1717e1d64c83089099/opentelemetry_api-1.40.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/8b/ca/8f122055c97a932311a3f640273f084e738008933503d0c2563cd5d591fc/opentelemetry_exporter_otlp_proto_common-1.40.0-py3-none-any.whl @@ -173,7 +152,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b2/f2/889ad4b2408f72fe1a4f6a19491177b30ea7bf1a0fd5f17050ca08cfc882/propcache-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/d7/64/0ea5be39e6a6515804cae8c280226d771f42750a08182f9d2e5f3b822694/PyArabic-0.6.15-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/36/2e/c0f017c405fcdc252dbccafbe05e36b0d0eb1ea9a958f081e01c6972927f/pyarrow-23.0.1-cp314-cp314-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl @@ -207,7 +185,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/ff/6b/a1548ac378a78332a4c3dcf4a134c2475a36d2a22ddfa272acd574140b50/ruff-0.15.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl @@ -224,14 +201,11 @@ environments: - pypi: https://files.pythonhosted.org/packages/27/83/859ecdd180cacc13b1f7e857abf8582a64552ea7a061057a6c716e790fce/sphinxcontrib_qthelp-2.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/52/a7/d2782e4e3f77c8450f727ba74a8f12756d5ba823d81b941f1b04da9d033a/sphinxcontrib_serializinghtml-2.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0b/c9/584bc9651441b4ba60cc4d557d8a547b5aff901af35bda3a4ee30c819b82/starlette-1.0.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/f5/de/9341a6d7a8f1b448573bbf3425fa57669ac58258a667eb48a25dfe916d70/tiktoken-0.12.0-cp314-cp314-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/47/e8/b98ca2d39b2e0e4730c0ee52537e488e7008025bc77ca89552ff91021f7c/torch-2.11.0-cp314-cp314-manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a4/21/a2266f7f1b0e58e624ff15fd6f01041f59182c49551ece0db9a183071329/torchvision-0.26.0-cp314-cp314-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl + - pypi: https://files.pythonhosted.org/packages/e7/28/35f7411ff80a3640c1f4fc907dcbb6a65061ebb82f66950e38bfc9f7f740/transformers-5.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/3a/7a/882d99539b19b1490cac5d77c67338d126e4122c8276bf640e411650c830/twine-6.2.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/89/29/8ac0281fc44c3297f0e58699ebf993c13621e32a0fab1025439d3ea8a2f1/ty-0.0.28-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl @@ -282,7 +256,6 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h366c992_103.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - - pypi: https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/8d/3f/95338030883d8c8b91223b4e21744b04d11b161a3ef117295d8241f50ab4/accessible_pygments-0.0.5-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0f/15/5bf3b99495fb160b63f95972b81750f18f7f4e02ad051373b669d17d44f2/aiohappyeyeballs-2.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/bd/c9/989f4034fb46841208de7aeeac2c6d8300745ab4f28c42f629ba77c2d916/aiohttp-3.13.5-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl @@ -306,9 +279,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/22/e5/06b1f88f42a5a99df42ce61208bdec3bddb3d261412874280a19796fc09c/coverage-7.13.5-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/34/71/1ea5a7352ae516d5512d17babe7e1b87d9db5150b21f794b1377eac1edc0/cryptography-46.0.6-cp311-abi3-manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/ee/5e/c0fe77a73aaefd3fff25ffaccaac69c5a63eafdf8b9a4c476626ef0ac703/cuda_bindings-13.2.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c4/74/8c66861b873d8eed51fde56d3091baa4906a56f0d4390cae991f2d41dda5/cuda_pathfinder-1.5.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b0/e5/247d094108e42ac26363ab8dc57f168840cf7c05774b40ffeb0d78868fcc/datasets-4.8.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/3f/27/4570e78fc0bf5ea0ca45eb1de3818a23787af9b390c0b0a0033a1b8236f9/diskcache-5.6.3-py3-none-any.whl @@ -327,7 +297,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/df/9a/a24b26dc8a65f0ecc0fe5be981a19e61e7ca963b85e062c083f3a9100529/hf_xet-1.4.3-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/a8/af/48ac8483240de756d2438c380746e7130d1c6f75802ef22f3c6d49982787/huggingface_hub-0.36.2-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/73/37/0d15d16150e1829f3e90962c99f28257f6de9e526a680b4c6f5acdb54fd2/huggingface_hub-1.9.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/42/77/de194443bf38daed9452139e960c632b0ef9f9a5dd9ce605fdf18ca9f1b1/id-1.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/46/33/92ef41c6fad0233e41d3d84ba8e8ad18d1780f1e5d99b3c683e6d7f98b63/identify-2.6.18-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl @@ -354,30 +324,13 @@ environments: - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/c9/f9/98d825105c450b9c67c27026caa374112b7e466c18331601d02ca278a01b/mistralai-1.12.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d8/f4/5e52c7319b8087acef603ed6e50dc325c02eaa999355414830468611f13c/more_itertools-11.0.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fe/3b/8ec5074bcfc450fe84273713b4b0a0dd47c0249358f5d82eb8104ffe2520/multidict-6.7.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/a0/61/af9115673a5870fd885247e2f1b68c4f1197737da315b520a91c757a861a/multiprocess-0.70.19-py314-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/df/76d0321c3797b54b60fef9ec3bd6f4cfd124b9e422182156a1dd418722cf/myst_parser-4.0.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d6/b7/ec1cbc6b297a808c513f59f501656389623fc09ad6a58c640851289c7854/nh3-0.3.4-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/9d/91/04e965f8e717ba0ab4bdca5c112deeab11c9e750d94c4d4602f050295d39/nltk-3.9.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/98/7c/21252050676612625449b4807d6b695b9ce8a7c9e1c197ee6216c8a65c7c/numpy-2.4.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a3/22/0b4b932655d17a6da1b92fa92ab12844b053bb2ac2475e179ba6f043da1e/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/2a/9e/5bfa2270f902d5b92ab7d41ce0475b8630572e71e349b2a4996d14bdda93/openai-2.30.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/bf/93795954016c522008da367da292adceed71cca6ee1717e1d64c83089099/opentelemetry_api-1.40.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/8b/ca/8f122055c97a932311a3f640273f084e738008933503d0c2563cd5d591fc/opentelemetry_exporter_otlp_proto_common-1.40.0-py3-none-any.whl @@ -393,7 +346,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b2/f2/889ad4b2408f72fe1a4f6a19491177b30ea7bf1a0fd5f17050ca08cfc882/propcache-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/16/92/d1e32e3e0d894fe00b15ce28ad4944ab692713f2e7f0a99787405e43533a/protobuf-6.33.6-cp39-abi3-manylinux2014_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/d7/64/0ea5be39e6a6515804cae8c280226d771f42750a08182f9d2e5f3b822694/PyArabic-0.6.15-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/36/2e/c0f017c405fcdc252dbccafbe05e36b0d0eb1ea9a958f081e01c6972927f/pyarrow-23.0.1-cp314-cp314-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl @@ -427,7 +379,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/ff/6b/a1548ac378a78332a4c3dcf4a134c2475a36d2a22ddfa272acd574140b50/ruff-0.15.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/e9/44/75a9c9421471a6c4805dbf2356f7c181a29c1879239abab1ea2cc8f38b40/sniffio-1.3.1-py3-none-any.whl @@ -444,14 +395,11 @@ environments: - pypi: https://files.pythonhosted.org/packages/27/83/859ecdd180cacc13b1f7e857abf8582a64552ea7a061057a6c716e790fce/sphinxcontrib_qthelp-2.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/52/a7/d2782e4e3f77c8450f727ba74a8f12756d5ba823d81b941f1b04da9d033a/sphinxcontrib_serializinghtml-2.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0b/c9/584bc9651441b4ba60cc4d557d8a547b5aff901af35bda3a4ee30c819b82/starlette-1.0.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/f5/de/9341a6d7a8f1b448573bbf3425fa57669ac58258a667eb48a25dfe916d70/tiktoken-0.12.0-cp314-cp314-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/47/e8/b98ca2d39b2e0e4730c0ee52537e488e7008025bc77ca89552ff91021f7c/torch-2.11.0-cp314-cp314-manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a4/21/a2266f7f1b0e58e624ff15fd6f01041f59182c49551ece0db9a183071329/torchvision-0.26.0-cp314-cp314-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl + - pypi: https://files.pythonhosted.org/packages/e7/28/35f7411ff80a3640c1f4fc907dcbb6a65061ebb82f66950e38bfc9f7f740/transformers-5.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/3a/7a/882d99539b19b1490cac5d77c67338d126e4122c8276bf640e411650c830/twine-6.2.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/89/29/8ac0281fc44c3297f0e58699ebf993c13621e32a0fab1025439d3ea8a2f1/ty-0.0.28-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl @@ -502,7 +450,6 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/tk-8.6.13-noxft_h366c992_103.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2025c-hc9c84f9_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/zstd-1.5.7-hb78ec9c_6.conda - - pypi: https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/8d/3f/95338030883d8c8b91223b4e21744b04d11b161a3ef117295d8241f50ab4/accessible_pygments-0.0.5-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0f/15/5bf3b99495fb160b63f95972b81750f18f7f4e02ad051373b669d17d44f2/aiohappyeyeballs-2.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/bd/c9/989f4034fb46841208de7aeeac2c6d8300745ab4f28c42f629ba77c2d916/aiohttp-3.13.5-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl @@ -523,9 +470,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/22/e5/06b1f88f42a5a99df42ce61208bdec3bddb3d261412874280a19796fc09c/coverage-7.13.5-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/34/71/1ea5a7352ae516d5512d17babe7e1b87d9db5150b21f794b1377eac1edc0/cryptography-46.0.6-cp311-abi3-manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/ee/5e/c0fe77a73aaefd3fff25ffaccaac69c5a63eafdf8b9a4c476626ef0ac703/cuda_bindings-13.2.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c4/74/8c66861b873d8eed51fde56d3091baa4906a56f0d4390cae991f2d41dda5/cuda_pathfinder-1.5.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b0/e5/247d094108e42ac26363ab8dc57f168840cf7c05774b40ffeb0d78868fcc/datasets-4.8.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/1e/77/dc8c558f7593132cf8fefec57c4f60c83b16941c574ac5f619abb3ae7933/dill-0.4.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl @@ -539,7 +483,7 @@ environments: - pypi: https://files.pythonhosted.org/packages/df/9a/a24b26dc8a65f0ecc0fe5be981a19e61e7ca963b85e062c083f3a9100529/hf_xet-1.4.3-cp37-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - pypi: https://files.pythonhosted.org/packages/7e/f5/f66802a942d491edb555dd61e3a9961140fd64c90bce1eafd741609d334d/httpcore-1.0.9-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2a/39/e50c7c3a983047577ee07d2a9e53faf5a69493943ec3f6a384bdc792deb2/httpx-0.28.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/a8/af/48ac8483240de756d2438c380746e7130d1c6f75802ef22f3c6d49982787/huggingface_hub-0.36.2-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/73/37/0d15d16150e1829f3e90962c99f28257f6de9e526a680b4c6f5acdb54fd2/huggingface_hub-1.9.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/42/77/de194443bf38daed9452139e960c632b0ef9f9a5dd9ce605fdf18ca9f1b1/id-1.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/46/33/92ef41c6fad0233e41d3d84ba8e8ad18d1780f1e5d99b3c683e6d7f98b63/identify-2.6.18-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl @@ -558,30 +502,13 @@ environments: - pypi: https://files.pythonhosted.org/packages/fb/86/dd6e5db36df29e76c7a7699123569a4a18c1623ce68d826ed96c62643cae/mdit_py_plugins-0.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d8/f4/5e52c7319b8087acef603ed6e50dc325c02eaa999355414830468611f13c/more_itertools-11.0.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/fe/3b/8ec5074bcfc450fe84273713b4b0a0dd47c0249358f5d82eb8104ffe2520/multidict-6.7.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/a0/61/af9115673a5870fd885247e2f1b68c4f1197737da315b520a91c757a861a/multiprocess-0.70.19-py314-none-any.whl - pypi: https://files.pythonhosted.org/packages/5f/df/76d0321c3797b54b60fef9ec3bd6f4cfd124b9e422182156a1dd418722cf/myst_parser-4.0.1-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/d6/b7/ec1cbc6b297a808c513f59f501656389623fc09ad6a58c640851289c7854/nh3-0.3.4-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/9d/91/04e965f8e717ba0ab4bdca5c112deeab11c9e750d94c4d4602f050295d39/nltk-3.9.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/98/7c/21252050676612625449b4807d6b695b9ce8a7c9e1c197ee6216c8a65c7c/numpy-2.4.4-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a3/22/0b4b932655d17a6da1b92fa92ab12844b053bb2ac2475e179ba6f043da1e/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl - pypi: https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/15/88/3cdd54fa279341afa10acf8d2b503556b1375245dccc9315659f795dd2e9/pandas-3.0.2-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/a8/2d/524f9318f6cbfcc79fbc004801ea6b607ec3f843977652fdee4857a7568b/pillow-11.3.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl @@ -589,7 +516,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/5d/19/fd3ef348460c80af7bb4669ea7926651d1f95c23ff2df18b9d24bab4f3fa/pre_commit-4.5.1-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b2/f2/889ad4b2408f72fe1a4f6a19491177b30ea7bf1a0fd5f17050ca08cfc882/propcache-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/d7/64/0ea5be39e6a6515804cae8c280226d771f42750a08182f9d2e5f3b822694/PyArabic-0.6.15-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/36/2e/c0f017c405fcdc252dbccafbe05e36b0d0eb1ea9a958f081e01c6972927f/pyarrow-23.0.1-cp314-cp314-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/5d/a0/7d793dce3fa811fe047d6ae2431c672364b462850c6235ae306c0efd025f/pyasn1-0.6.3-py3-none-any.whl @@ -617,7 +543,6 @@ environments: - pypi: https://files.pythonhosted.org/packages/ff/6b/a1548ac378a78332a4c3dcf4a134c2475a36d2a22ddfa272acd574140b50/ruff-0.15.9-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/c8/78/3565d011c61f5a43488987ee32b6f3f656e7f107ac2782dd57bdd7d91d9a/snowballstemmer-3.0.1-py3-none-any.whl @@ -633,13 +558,10 @@ environments: - pypi: https://files.pythonhosted.org/packages/27/83/859ecdd180cacc13b1f7e857abf8582a64552ea7a061057a6c716e790fce/sphinxcontrib_qthelp-2.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/52/a7/d2782e4e3f77c8450f727ba74a8f12756d5ba823d81b941f1b04da9d033a/sphinxcontrib_serializinghtml-2.0.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/0b/c9/584bc9651441b4ba60cc4d557d8a547b5aff901af35bda3a4ee30c819b82/starlette-1.0.0-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl + - pypi: https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2e/76/932be4b50ef6ccedf9d3c6639b056a967a86258c6d9200643f01269211ca/tokenizers-0.22.2-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/47/e8/b98ca2d39b2e0e4730c0ee52537e488e7008025bc77ca89552ff91021f7c/torch-2.11.0-cp314-cp314-manylinux_2_28_x86_64.whl - - pypi: https://files.pythonhosted.org/packages/a4/21/a2266f7f1b0e58e624ff15fd6f01041f59182c49551ece0db9a183071329/torchvision-0.26.0-cp314-cp314-manylinux_2_28_x86_64.whl - pypi: https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl - - pypi: https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl + - pypi: https://files.pythonhosted.org/packages/e7/28/35f7411ff80a3640c1f4fc907dcbb6a65061ebb82f66950e38bfc9f7f740/transformers-5.5.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/3a/7a/882d99539b19b1490cac5d77c67338d126e4122c8276bf640e411650c830/twine-6.2.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/89/29/8ac0281fc44c3297f0e58699ebf993c13621e32a0fab1025439d3ea8a2f1/ty-0.0.28-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl - pypi: https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl @@ -667,81 +589,6 @@ packages: purls: [] size: 28948 timestamp: 1770939786096 -- pypi: https://files.pythonhosted.org/packages/7e/46/02ac5e262d4af18054b3e922b2baedbb2a03289ee792162de60a865defc5/accelerate-1.13.0-py3-none-any.whl - name: accelerate - version: 1.13.0 - sha256: cf1a3efb96c18f7b152eb0fa7490f3710b19c3f395699358f08decca2b8b62e0 - requires_dist: - - numpy>=1.17 - - packaging>=20.0 - - psutil - - pyyaml - - torch>=2.0.0 - - huggingface-hub>=0.21.0 - - safetensors>=0.4.3 - - ruff==0.13.1 ; extra == 'quality' - - pytest>=7.2.0 ; extra == 'test-prod' - - pytest-xdist ; extra == 'test-prod' - - pytest-subtests ; extra == 'test-prod' - - parameterized ; extra == 'test-prod' - - pytest-order ; extra == 'test-prod' - - datasets ; extra == 'test-dev' - - diffusers ; extra == 'test-dev' - - evaluate ; extra == 'test-dev' - - torchdata>=0.8.0 ; extra == 'test-dev' - - torchpippy>=0.2.0 ; extra == 'test-dev' - - transformers ; extra == 'test-dev' - - scipy ; extra == 'test-dev' - - scikit-learn ; extra == 'test-dev' - - tqdm ; extra == 'test-dev' - - bitsandbytes ; extra == 'test-dev' - - timm ; extra == 'test-dev' - - pytest>=7.2.0 ; extra == 'testing' - - pytest-xdist ; extra == 'testing' - - pytest-subtests ; extra == 'testing' - - parameterized ; extra == 'testing' - - pytest-order ; extra == 'testing' - - datasets ; extra == 'testing' - - diffusers ; extra == 'testing' - - evaluate ; extra == 'testing' - - torchdata>=0.8.0 ; extra == 'testing' - - torchpippy>=0.2.0 ; extra == 'testing' - - transformers ; extra == 'testing' - - scipy ; extra == 'testing' - - scikit-learn ; extra == 'testing' - - tqdm ; extra == 'testing' - - bitsandbytes ; extra == 'testing' - - timm ; extra == 'testing' - - deepspeed ; extra == 'deepspeed' - - rich ; extra == 'rich' - - torchao ; extra == 'test-fp8' - - wandb ; extra == 'test-trackers' - - comet-ml ; extra == 'test-trackers' - - tensorboard ; extra == 'test-trackers' - - dvclive ; extra == 'test-trackers' - - matplotlib ; extra == 'test-trackers' - - swanlab[dashboard] ; extra == 'test-trackers' - - trackio ; extra == 'test-trackers' - - ruff==0.13.1 ; extra == 'dev' - - pytest>=7.2.0 ; extra == 'dev' - - pytest-xdist ; extra == 'dev' - - pytest-subtests ; extra == 'dev' - - parameterized ; extra == 'dev' - - pytest-order ; extra == 'dev' - - datasets ; extra == 'dev' - - diffusers ; extra == 'dev' - - evaluate ; extra == 'dev' - - torchdata>=0.8.0 ; extra == 'dev' - - torchpippy>=0.2.0 ; extra == 'dev' - - transformers ; extra == 'dev' - - scipy ; extra == 'dev' - - scikit-learn ; extra == 'dev' - - tqdm ; extra == 'dev' - - bitsandbytes ; extra == 'dev' - - timm ; extra == 'dev' - - rich ; extra == 'dev' - - sagemaker ; extra == 'sagemaker' - requires_python: '>=3.10.0' - pypi: https://files.pythonhosted.org/packages/8d/3f/95338030883d8c8b91223b4e21744b04d11b161a3ef117295d8241f50ab4/accessible_pygments-0.0.5-py3-none-any.whl name: accessible-pygments version: 0.0.5 @@ -932,21 +779,18 @@ packages: - pypi: ./ name: churro-ocr version: 0.2.0 - sha256: 6a6c02deb47eda1382303d34056d6b79ca06f84f75e8739d162c1f06eed4563d + sha256: 4aab922d00e91acd3427b8ddbb59341f7f5232ccc72fbc120a3754e7e05d0dbc requires_dist: - loguru>=0.7.2,<1 - pillow>=10.4.0,<12 - rich>=13.9.2,<14 + - tenacity>=9.1.2,<10 - typer>=0.12.3,<1 - google-auth>=2.41.1,<3 ; extra == 'llm' - litellm[caching]==1.82.3 ; extra == 'llm' - azure-ai-documentintelligence==1.0.2 ; extra == 'azure' - qwen-vl-utils ; extra == 'hf' - - transformers[torch]>=4.57.0,<5 ; extra == 'hf' - - torchvision ; extra == 'hf' - - transformers>=4.57.0,<5 ; extra == 'vllm' - - torchvision ; extra == 'vllm' - - vllm>=0.18,<1 ; extra == 'vllm' + - transformers>=5,<6 ; extra == 'hf' - mistralai>=1.6.0,<2 ; extra == 'mistral' - litellm[caching]==1.82.3 ; extra == 'local' - pypdfium2>=5,<6 ; extra == 'pdf' @@ -954,9 +798,7 @@ packages: - litellm[caching]==1.82.3 ; extra == 'all' - azure-ai-documentintelligence==1.0.2 ; extra == 'all' - qwen-vl-utils ; extra == 'all' - - transformers[torch]>=4.57.0,<5 ; extra == 'all' - - torchvision ; extra == 'all' - - vllm>=0.18,<1 ; extra == 'all' + - transformers>=5,<6 ; extra == 'all' - mistralai>=1.6.0,<2 ; extra == 'all' - pypdfium2>=5,<6 ; extra == 'all' requires_python: '>=3.12' @@ -1009,15 +851,6 @@ packages: - check-sdist ; extra == 'pep8test' - click>=8.0.1 ; extra == 'pep8test' requires_python: '>=3.8,!=3.9.0,!=3.9.1' -- pypi: https://files.pythonhosted.org/packages/ee/5e/c0fe77a73aaefd3fff25ffaccaac69c5a63eafdf8b9a4c476626ef0ac703/cuda_bindings-13.2.0-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl - name: cuda-bindings - version: 13.2.0 - sha256: f4af9f3e1be603fa12d5ad6cfca7844c9d230befa9792b5abdf7dd79979c3626 - requires_dist: - - cuda-pathfinder~=1.1 - - cuda-toolkit[nvfatbin,nvjitlink,nvrtc,nvvm]==13.* ; extra == 'all' - - cuda-toolkit[cufile]==13.* ; sys_platform == 'linux' and extra == 'all' - requires_python: '>=3.10' - conda: https://conda.anaconda.org/conda-forge/linux-64/cuda-cudart-12.9.79-h5888daf_0.conda sha256: 57d1294ecfaf9dc8cdb5fc4be3e63ebc7614538bddb5de53cfd9b1b7de43aed5 md5: cb15315d19b58bd9cd424084e58ad081 @@ -1086,66 +919,6 @@ packages: purls: [] size: 30747 timestamp: 1746192810479 -- pypi: https://files.pythonhosted.org/packages/c4/74/8c66861b873d8eed51fde56d3091baa4906a56f0d4390cae991f2d41dda5/cuda_pathfinder-1.5.1-py3-none-any.whl - name: cuda-pathfinder - version: 1.5.1 - sha256: b3718097fb57cf9e8a904dd072d806f2c9a27627e35c020b06ab9454bcec08c0 - requires_python: '>=3.10' -- pypi: https://files.pythonhosted.org/packages/57/b2/453099f5f3b698d7d0eab38916aac44c7f76229f451709e2eb9db6615dcd/cuda_toolkit-13.0.2-py2.py3-none-any.whl - name: cuda-toolkit - version: 13.0.2 - sha256: b198824cf2f54003f50d64ada3a0f184b42ca0846c1c94192fa269ecd97a66eb - requires_dist: - - nvidia-cublas==13.1.0.3.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-cccl==13.0.85.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-crt==13.0.88.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-culibos==13.0.85.* ; sys_platform == 'linux' and extra == 'all' - - nvidia-cuda-cupti==13.0.85.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-cuxxfilt==13.0.85.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-nvcc==13.0.88.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-nvrtc==13.0.88.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-opencl==13.0.85.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-profiler-api==13.0.85.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-runtime==13.0.96.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-sanitizer-api==13.0.85.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cufft==12.0.0.61.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cufile==1.15.1.6.* ; sys_platform == 'linux' and extra == 'all' - - nvidia-curand==10.4.0.35.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cusolver==12.0.4.66.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cusparse==12.6.3.3.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-npp==13.0.1.2.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-nvfatbin==13.0.85.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-nvjitlink==13.0.88.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-nvjpeg==13.0.1.86.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-nvml-dev==13.0.87.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-nvptxcompiler==13.0.88.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-nvtx==13.0.85.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-nvvm==13.0.88.* ; (sys_platform == 'linux' and extra == 'all') or (sys_platform == 'win32' and extra == 'all') - - nvidia-cuda-cccl==13.0.85.* ; (sys_platform == 'linux' and extra == 'cccl') or (sys_platform == 'win32' and extra == 'cccl') - - nvidia-cuda-crt==13.0.88.* ; (sys_platform == 'linux' and extra == 'crt') or (sys_platform == 'win32' and extra == 'crt') - - nvidia-cublas==13.1.0.3.* ; (sys_platform == 'linux' and extra == 'cublas') or (sys_platform == 'win32' and extra == 'cublas') - - nvidia-cuda-runtime==13.0.96.* ; (sys_platform == 'linux' and extra == 'cudart') or (sys_platform == 'win32' and extra == 'cudart') - - nvidia-cufft==12.0.0.61.* ; (sys_platform == 'linux' and extra == 'cufft') or (sys_platform == 'win32' and extra == 'cufft') - - nvidia-cufile==1.15.1.6.* ; sys_platform == 'linux' and extra == 'cufile' - - nvidia-cuda-culibos==13.0.85.* ; sys_platform == 'linux' and extra == 'culibos' - - nvidia-cuda-cupti==13.0.85.* ; (sys_platform == 'linux' and extra == 'cupti') or (sys_platform == 'win32' and extra == 'cupti') - - nvidia-curand==10.4.0.35.* ; (sys_platform == 'linux' and extra == 'curand') or (sys_platform == 'win32' and extra == 'curand') - - nvidia-cusolver==12.0.4.66.* ; (sys_platform == 'linux' and extra == 'cusolver') or (sys_platform == 'win32' and extra == 'cusolver') - - nvidia-cusparse==12.6.3.3.* ; (sys_platform == 'linux' and extra == 'cusparse') or (sys_platform == 'win32' and extra == 'cusparse') - - nvidia-cuda-cuxxfilt==13.0.85.* ; (sys_platform == 'linux' and extra == 'cuxxfilt') or (sys_platform == 'win32' and extra == 'cuxxfilt') - - nvidia-npp==13.0.1.2.* ; (sys_platform == 'linux' and extra == 'npp') or (sys_platform == 'win32' and extra == 'npp') - - nvidia-cuda-nvcc==13.0.88.* ; (sys_platform == 'linux' and extra == 'nvcc') or (sys_platform == 'win32' and extra == 'nvcc') - - nvidia-nvfatbin==13.0.85.* ; (sys_platform == 'linux' and extra == 'nvfatbin') or (sys_platform == 'win32' and extra == 'nvfatbin') - - nvidia-nvjitlink==13.0.88.* ; (sys_platform == 'linux' and extra == 'nvjitlink') or (sys_platform == 'win32' and extra == 'nvjitlink') - - nvidia-nvjpeg==13.0.1.86.* ; (sys_platform == 'linux' and extra == 'nvjpeg') or (sys_platform == 'win32' and extra == 'nvjpeg') - - nvidia-nvml-dev==13.0.87.* ; (sys_platform == 'linux' and extra == 'nvml') or (sys_platform == 'win32' and extra == 'nvml') - - nvidia-nvptxcompiler==13.0.88.* ; (sys_platform == 'linux' and extra == 'nvptxcompiler') or (sys_platform == 'win32' and extra == 'nvptxcompiler') - - nvidia-cuda-nvrtc==13.0.88.* ; (sys_platform == 'linux' and extra == 'nvrtc') or (sys_platform == 'win32' and extra == 'nvrtc') - - nvidia-nvtx==13.0.85.* ; (sys_platform == 'linux' and extra == 'nvtx') or (sys_platform == 'win32' and extra == 'nvtx') - - nvidia-nvvm==13.0.88.* ; (sys_platform == 'linux' and extra == 'nvvm') or (sys_platform == 'win32' and extra == 'nvvm') - - nvidia-cuda-opencl==13.0.85.* ; (sys_platform == 'linux' and extra == 'opencl') or (sys_platform == 'win32' and extra == 'opencl') - - nvidia-cuda-profiler-api==13.0.85.* ; (sys_platform == 'linux' and extra == 'profiler') or (sys_platform == 'win32' and extra == 'profiler') - - nvidia-cuda-sanitizer-api==13.0.85.* ; (sys_platform == 'linux' and extra == 'sanitizer') or (sys_platform == 'win32' and extra == 'sanitizer') - conda: https://conda.anaconda.org/conda-forge/noarch/cuda-version-12.9-h4f385c5_3.conda sha256: 5f5f428031933f117ff9f7fcc650e6ea1b3fef5936cf84aa24af79167513b656 md5: b6d5d7f1c171cbd228ea06b556cfa859 @@ -1647,49 +1420,38 @@ packages: - socksio==1.* ; extra == 'socks' - zstandard>=0.18.0 ; extra == 'zstd' requires_python: '>=3.8' -- pypi: https://files.pythonhosted.org/packages/a8/af/48ac8483240de756d2438c380746e7130d1c6f75802ef22f3c6d49982787/huggingface_hub-0.36.2-py3-none-any.whl +- pypi: https://files.pythonhosted.org/packages/73/37/0d15d16150e1829f3e90962c99f28257f6de9e526a680b4c6f5acdb54fd2/huggingface_hub-1.9.0-py3-none-any.whl name: huggingface-hub - version: 0.36.2 - sha256: 48f0c8eac16145dfce371e9d2d7772854a4f591bcb56c9cf548accf531d54270 + version: 1.9.0 + sha256: 2999328c058d39fd19ab748dd09bd4da2fbaa4f4c1ddea823eab103051e14a1f requires_dist: - - filelock + - filelock>=3.10.0 - fsspec>=2023.5.0 - - hf-xet>=1.1.3,<2.0.0 ; platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' + - hf-xet>=1.4.3,<2.0.0 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' + - httpx>=0.23.0,<1 - packaging>=20.9 - pyyaml>=5.1 - - requests - tqdm>=4.42.1 - - typing-extensions>=3.7.4.3 - - inquirerpy==0.3.4 ; extra == 'cli' - - aiohttp ; extra == 'inference' + - typer + - typing-extensions>=4.1.0 - authlib>=1.3.2 ; extra == 'oauth' - fastapi ; extra == 'oauth' - httpx ; extra == 'oauth' - itsdangerous ; extra == 'oauth' - torch ; extra == 'torch' - safetensors[torch] ; extra == 'torch' - - hf-transfer>=0.1.4 ; extra == 'hf-transfer' - toml ; extra == 'fastai' - fastai>=2.4 ; extra == 'fastai' - fastcore>=1.3.27 ; extra == 'fastai' - - tensorflow ; extra == 'tensorflow' - - pydot ; extra == 'tensorflow' - - graphviz ; extra == 'tensorflow' - - tensorflow ; extra == 'tensorflow-testing' - - keras<3.0 ; extra == 'tensorflow-testing' - - hf-xet>=1.1.2,<2.0.0 ; extra == 'hf-xet' + - hf-xet>=1.4.3,<2.0.0 ; extra == 'hf-xet' - mcp>=1.8.0 ; extra == 'mcp' - - typer ; extra == 'mcp' - - aiohttp ; extra == 'mcp' - - inquirerpy==0.3.4 ; extra == 'testing' - - aiohttp ; extra == 'testing' - authlib>=1.3.2 ; extra == 'testing' - fastapi ; extra == 'testing' - httpx ; extra == 'testing' - itsdangerous ; extra == 'testing' - jedi ; extra == 'testing' - jinja2 ; extra == 'testing' - - pytest>=8.1.1,<8.2.2 ; extra == 'testing' + - pytest>=8.4.2 ; extra == 'testing' - pytest-cov ; extra == 'testing' - pytest-env ; extra == 'testing' - pytest-xdist ; extra == 'testing' @@ -1700,30 +1462,28 @@ packages: - urllib3<2.0 ; extra == 'testing' - soundfile ; extra == 'testing' - pillow ; extra == 'testing' - - gradio>=4.0.0 ; extra == 'testing' - numpy ; extra == 'testing' + - duckdb ; extra == 'testing' - fastapi ; extra == 'testing' + - gradio>=5.0.0 ; extra == 'gradio' + - requests ; extra == 'gradio' - typing-extensions>=4.8.0 ; extra == 'typing' - types-pyyaml ; extra == 'typing' - - types-requests ; extra == 'typing' - types-simplejson ; extra == 'typing' - types-toml ; extra == 'typing' - types-tqdm ; extra == 'typing' - types-urllib3 ; extra == 'typing' - ruff>=0.9.0 ; extra == 'quality' - - mypy>=1.14.1,<1.15.0 ; python_full_version == '3.8.*' and extra == 'quality' - - mypy==1.15.0 ; python_full_version >= '3.9' and extra == 'quality' + - mypy==1.15.0 ; extra == 'quality' - libcst>=1.4.0 ; extra == 'quality' - ty ; extra == 'quality' - - inquirerpy==0.3.4 ; extra == 'all' - - aiohttp ; extra == 'all' - authlib>=1.3.2 ; extra == 'all' - fastapi ; extra == 'all' - httpx ; extra == 'all' - itsdangerous ; extra == 'all' - jedi ; extra == 'all' - jinja2 ; extra == 'all' - - pytest>=8.1.1,<8.2.2 ; extra == 'all' + - pytest>=8.4.2 ; extra == 'all' - pytest-cov ; extra == 'all' - pytest-env ; extra == 'all' - pytest-xdist ; extra == 'all' @@ -1734,30 +1494,26 @@ packages: - urllib3<2.0 ; extra == 'all' - soundfile ; extra == 'all' - pillow ; extra == 'all' - - gradio>=4.0.0 ; extra == 'all' - numpy ; extra == 'all' + - duckdb ; extra == 'all' - fastapi ; extra == 'all' - ruff>=0.9.0 ; extra == 'all' - - mypy>=1.14.1,<1.15.0 ; python_full_version == '3.8.*' and extra == 'all' - - mypy==1.15.0 ; python_full_version >= '3.9' and extra == 'all' + - mypy==1.15.0 ; extra == 'all' - libcst>=1.4.0 ; extra == 'all' - ty ; extra == 'all' - typing-extensions>=4.8.0 ; extra == 'all' - types-pyyaml ; extra == 'all' - - types-requests ; extra == 'all' - types-simplejson ; extra == 'all' - types-toml ; extra == 'all' - types-tqdm ; extra == 'all' - types-urllib3 ; extra == 'all' - - inquirerpy==0.3.4 ; extra == 'dev' - - aiohttp ; extra == 'dev' - authlib>=1.3.2 ; extra == 'dev' - fastapi ; extra == 'dev' - httpx ; extra == 'dev' - itsdangerous ; extra == 'dev' - jedi ; extra == 'dev' - jinja2 ; extra == 'dev' - - pytest>=8.1.1,<8.2.2 ; extra == 'dev' + - pytest>=8.4.2 ; extra == 'dev' - pytest-cov ; extra == 'dev' - pytest-env ; extra == 'dev' - pytest-xdist ; extra == 'dev' @@ -1768,22 +1524,20 @@ packages: - urllib3<2.0 ; extra == 'dev' - soundfile ; extra == 'dev' - pillow ; extra == 'dev' - - gradio>=4.0.0 ; extra == 'dev' - numpy ; extra == 'dev' + - duckdb ; extra == 'dev' - fastapi ; extra == 'dev' - ruff>=0.9.0 ; extra == 'dev' - - mypy>=1.14.1,<1.15.0 ; python_full_version == '3.8.*' and extra == 'dev' - - mypy==1.15.0 ; python_full_version >= '3.9' and extra == 'dev' + - mypy==1.15.0 ; extra == 'dev' - libcst>=1.4.0 ; extra == 'dev' - ty ; extra == 'dev' - typing-extensions>=4.8.0 ; extra == 'dev' - types-pyyaml ; extra == 'dev' - - types-requests ; extra == 'dev' - types-simplejson ; extra == 'dev' - types-toml ; extra == 'dev' - types-tqdm ; extra == 'dev' - types-urllib3 ; extra == 'dev' - requires_python: '>=3.8.0' + requires_python: '>=3.10.0' - conda: https://conda.anaconda.org/conda-forge/linux-64/icu-78.3-h33c6efd_0.conda sha256: fbf86c4a59c2ed05bbffb2ba25c7ed94f6185ec30ecb691615d42342baa1a16a md5: c80d8a3b84358cb967fa81e7075fbc8a @@ -2539,19 +2293,6 @@ packages: version: 11.0.1 sha256: eaf287826069452a8f61026c597eae2428b2d1ba2859083abbf240b46842ce6d requires_python: '>=3.10' -- pypi: https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl - name: mpmath - version: 1.3.0 - sha256: a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c - requires_dist: - - pytest>=4.6 ; extra == 'develop' - - pycodestyle ; extra == 'develop' - - pytest-cov ; extra == 'develop' - - codecov ; extra == 'develop' - - wheel ; extra == 'develop' - - sphinx ; extra == 'docs' - - gmpy2>=2.1.0a4 ; platform_python_implementation != 'PyPy' and extra == 'gmpy' - - pytest>=4.6 ; extra == 'tests' - pypi: https://files.pythonhosted.org/packages/fe/3b/8ec5074bcfc450fe84273713b4b0a0dd47c0249358f5d82eb8104ffe2520/multidict-6.7.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl name: multidict version: 6.7.1 @@ -2613,49 +2354,6 @@ packages: purls: [] size: 891641 timestamp: 1738195959188 -- pypi: https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl - name: networkx - version: 3.6.1 - sha256: d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762 - requires_dist: - - asv ; extra == 'benchmarking' - - virtualenv ; extra == 'benchmarking' - - numpy>=1.25 ; extra == 'default' - - scipy>=1.11.2 ; extra == 'default' - - matplotlib>=3.8 ; extra == 'default' - - pandas>=2.0 ; extra == 'default' - - pre-commit>=4.1 ; extra == 'developer' - - mypy>=1.15 ; extra == 'developer' - - sphinx>=8.0 ; extra == 'doc' - - pydata-sphinx-theme>=0.16 ; extra == 'doc' - - sphinx-gallery>=0.18 ; extra == 'doc' - - numpydoc>=1.8.0 ; extra == 'doc' - - pillow>=10 ; extra == 'doc' - - texext>=0.6.7 ; extra == 'doc' - - myst-nb>=1.1 ; extra == 'doc' - - intersphinx-registry ; extra == 'doc' - - osmnx>=2.0.0 ; extra == 'example' - - momepy>=0.7.2 ; extra == 'example' - - contextily>=1.6 ; extra == 'example' - - seaborn>=0.13 ; extra == 'example' - - cairocffi>=1.7 ; extra == 'example' - - igraph>=0.11 ; extra == 'example' - - scikit-learn>=1.5 ; extra == 'example' - - iplotx>=0.9.0 ; extra == 'example' - - lxml>=4.6 ; extra == 'extra' - - pygraphviz>=1.14 ; extra == 'extra' - - pydot>=3.0.1 ; extra == 'extra' - - sympy>=1.10 ; extra == 'extra' - - build>=0.10 ; extra == 'release' - - twine>=4.0 ; extra == 'release' - - wheel>=0.40 ; extra == 'release' - - changelist==0.5 ; extra == 'release' - - pytest>=7.2 ; extra == 'test' - - pytest-cov>=4.0 ; extra == 'test' - - pytest-xdist>=3.0 ; extra == 'test' - - pytest-mpl ; extra == 'test-extras' - - pytest-randomly ; extra == 'test-extras' - requires_python: '>=3.11,!=3.14.1' - pypi: https://files.pythonhosted.org/packages/d6/b7/ec1cbc6b297a808c513f59f501656389623fc09ad6a58c640851289c7854/nh3-0.3.4-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl name: nh3 version: 0.3.4 @@ -2697,90 +2395,6 @@ packages: version: 2.4.4 sha256: 27a8d92cd10f1382a67d7cf4db7ce18341b66438bdd9f691d7b0e48d104c2a9d requires_python: '>=3.11' -- pypi: https://files.pythonhosted.org/packages/e7/44/423ac00af4dd95a5aeb27207e2c0d9b7118702149bf4704c3ddb55bb7429/nvidia_cublas-13.1.0.3-py3-none-manylinux_2_27_x86_64.whl - name: nvidia-cublas - version: 13.1.0.3 - sha256: ee8722c1f0145ab246bccb9e452153b5e0515fd094c3678df50b2a0888b8b171 - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl - name: nvidia-cuda-cupti - version: 13.0.85 - sha256: 4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8 - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl - name: nvidia-cuda-nvrtc - version: 13.0.88 - sha256: ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575 - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - name: nvidia-cuda-runtime - version: 13.0.96 - sha256: 7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548 - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/a3/22/0b4b932655d17a6da1b92fa92ab12844b053bb2ac2475e179ba6f043da1e/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_x86_64.whl - name: nvidia-cudnn-cu13 - version: 9.19.0.56 - sha256: d20e1734305e9d68889a96e3f35094d733ff1f83932ebe462753973e53a572bf - requires_dist: - - nvidia-cublas - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - name: nvidia-cufft - version: 12.0.0.61 - sha256: 6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3 - requires_dist: - - nvidia-nvjitlink - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - name: nvidia-cufile - version: 1.15.1.6 - sha256: 08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44 - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl - name: nvidia-curand - version: 10.4.0.35 - sha256: 1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl - name: nvidia-cusolver - version: 12.0.4.66 - sha256: 0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112 - requires_dist: - - nvidia-cublas - - nvidia-nvjitlink - - nvidia-cusparse - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - name: nvidia-cusparse - version: 12.6.3.3 - sha256: 2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b - requires_dist: - - nvidia-nvjitlink - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/fd/53/43b0d71f4e702fa9733f8b4571fdca50a8813f1e450b656c239beff12315/nvidia_cusparselt_cu13-0.8.0-py3-none-manylinux2014_x86_64.whl - name: nvidia-cusparselt-cu13 - version: 0.8.0 - sha256: 25e30a8a7323935d4ad0340b95a0b69926eee755767e8e0b1cf8dd85b197d3fd -- pypi: https://files.pythonhosted.org/packages/b0/b4/878fefaad5b2bcc6fcf8d474a25e3e3774bc5133e4b58adff4d0bca238bc/nvidia_nccl_cu13-2.28.9-py3-none-manylinux_2_18_x86_64.whl - name: nvidia-nccl-cu13 - version: 2.28.9 - sha256: e4553a30f34195f3fa1da02a6da3d6337d28f2003943aa0a3d247bbc25fefc42 - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/56/7a/123e033aaff487c77107195fa5a2b8686795ca537935a24efae476c41f05/nvidia_nvjitlink-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl - name: nvidia-nvjitlink - version: 13.0.88 - sha256: 13a74f429e23b921c1109976abefacc69835f2f433ebd323d3946e11d804e47b - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl - name: nvidia-nvshmem-cu13 - version: 3.4.5 - sha256: 290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80 - requires_python: '>=3' -- pypi: https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl - name: nvidia-nvtx - version: 13.0.85 - sha256: 4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4 - requires_python: '>=3' - conda: https://conda.anaconda.org/conda-forge/linux-64/ocl-icd-2.3.3-hb9d3cd8_0.conda sha256: 2254dae821b286fb57c61895f2b40e3571a070910fdab79a948ff703e1ea807b md5: 56f8947aa9d5cf37b0b3d43b83f34192 @@ -3066,50 +2680,6 @@ packages: version: 6.33.6 sha256: e9db7e292e0ab79dd108d7f1a94fe31601ce1ee3f7b79e0692043423020b0593 requires_python: '>=3.9' -- pypi: https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl - name: psutil - version: 7.2.2 - sha256: 076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9 - requires_dist: - - psleak ; extra == 'dev' - - pytest ; extra == 'dev' - - pytest-instafail ; extra == 'dev' - - pytest-xdist ; extra == 'dev' - - setuptools ; extra == 'dev' - - abi3audit ; extra == 'dev' - - black ; extra == 'dev' - - check-manifest ; extra == 'dev' - - coverage ; extra == 'dev' - - packaging ; extra == 'dev' - - pylint ; extra == 'dev' - - pyperf ; extra == 'dev' - - pypinfo ; extra == 'dev' - - pytest-cov ; extra == 'dev' - - requests ; extra == 'dev' - - rstcheck ; extra == 'dev' - - ruff ; extra == 'dev' - - sphinx ; extra == 'dev' - - sphinx-rtd-theme ; extra == 'dev' - - toml-sort ; extra == 'dev' - - twine ; extra == 'dev' - - validate-pyproject[all] ; extra == 'dev' - - virtualenv ; extra == 'dev' - - vulture ; extra == 'dev' - - wheel ; extra == 'dev' - - colorama ; os_name == 'nt' and extra == 'dev' - - pyreadline3 ; os_name == 'nt' and extra == 'dev' - - pywin32 ; implementation_name != 'pypy' and os_name == 'nt' and extra == 'dev' - - wheel ; implementation_name != 'pypy' and os_name == 'nt' and extra == 'dev' - - wmi ; implementation_name != 'pypy' and os_name == 'nt' and extra == 'dev' - - psleak ; extra == 'test' - - pytest ; extra == 'test' - - pytest-instafail ; extra == 'test' - - pytest-xdist ; extra == 'test' - - setuptools ; extra == 'test' - - pywin32 ; implementation_name != 'pypy' and os_name == 'nt' and extra == 'test' - - wheel ; implementation_name != 'pypy' and os_name == 'nt' and extra == 'test' - - wmi ; implementation_name != 'pypy' and os_name == 'nt' and extra == 'test' - requires_python: '>=3.6' - pypi: https://files.pythonhosted.org/packages/d7/64/0ea5be39e6a6515804cae8c280226d771f42750a08182f9d2e5f3b822694/PyArabic-0.6.15-py3-none-any.whl name: pyarabic version: 0.6.15 @@ -3527,63 +3097,6 @@ packages: - cryptography>=2.0 - jeepney>=0.6 requires_python: '>=3.10' -- pypi: https://files.pythonhosted.org/packages/e1/e3/c164c88b2e5ce7b24d667b9bd83589cf4f3520d97cad01534cd3c4f55fdb/setuptools-81.0.0-py3-none-any.whl - name: setuptools - version: 81.0.0 - sha256: fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6 - requires_dist: - - pytest>=6,!=8.1.* ; extra == 'test' - - virtualenv>=13.0.0 ; extra == 'test' - - wheel>=0.44.0 ; extra == 'test' - - pip>=19.1 ; extra == 'test' - - packaging>=24.2 ; extra == 'test' - - jaraco-envs>=2.2 ; extra == 'test' - - pytest-xdist>=3 ; extra == 'test' - - jaraco-path>=3.7.2 ; extra == 'test' - - build[virtualenv]>=1.0.3 ; extra == 'test' - - filelock>=3.4.0 ; extra == 'test' - - ini2toml[lite]>=0.14 ; extra == 'test' - - tomli-w>=1.0.0 ; extra == 'test' - - pytest-timeout ; extra == 'test' - - pytest-perf ; sys_platform != 'cygwin' and extra == 'test' - - jaraco-develop>=7.21 ; python_full_version >= '3.9' and sys_platform != 'cygwin' and extra == 'test' - - pytest-home>=0.5 ; extra == 'test' - - pytest-subprocess ; extra == 'test' - - pyproject-hooks!=1.1 ; extra == 'test' - - jaraco-test>=5.5 ; extra == 'test' - - sphinx>=3.5 ; extra == 'doc' - - jaraco-packaging>=9.3 ; extra == 'doc' - - rst-linker>=1.9 ; extra == 'doc' - - furo ; extra == 'doc' - - sphinx-lint ; extra == 'doc' - - jaraco-tidelift>=1.4 ; extra == 'doc' - - pygments-github-lexers==0.0.5 ; extra == 'doc' - - sphinx-favicon ; extra == 'doc' - - sphinx-inline-tabs ; extra == 'doc' - - sphinx-reredirects ; extra == 'doc' - - sphinxcontrib-towncrier ; extra == 'doc' - - sphinx-notfound-page>=1,<2 ; extra == 'doc' - - pyproject-hooks!=1.1 ; extra == 'doc' - - towncrier<24.7 ; extra == 'doc' - - packaging>=24.2 ; extra == 'core' - - more-itertools>=8.8 ; extra == 'core' - - jaraco-text>=3.7 ; extra == 'core' - - importlib-metadata>=6 ; python_full_version < '3.10' and extra == 'core' - - tomli>=2.0.1 ; python_full_version < '3.11' and extra == 'core' - - wheel>=0.43.0 ; extra == 'core' - - platformdirs>=4.2.2 ; extra == 'core' - - jaraco-functools>=4 ; extra == 'core' - - more-itertools ; extra == 'core' - - pytest-checkdocs>=2.4 ; extra == 'check' - - pytest-ruff>=0.2.1 ; sys_platform != 'cygwin' and extra == 'check' - - ruff>=0.13.0 ; sys_platform != 'cygwin' and extra == 'check' - - pytest-cov ; extra == 'cover' - - pytest-enabler>=2.2 ; extra == 'enabler' - - pytest-mypy ; extra == 'type' - - mypy==1.18.* ; extra == 'type' - - importlib-metadata>=7.0.2 ; python_full_version < '3.10' and extra == 'type' - - jaraco-develop>=7.21 ; sys_platform != 'cygwin' and extra == 'type' - requires_python: '>=3.9' - pypi: https://files.pythonhosted.org/packages/e0/f9/0595336914c5619e5f28a1fb793285925a8cd4b432c9da0a987836c7f822/shellingham-1.5.4-py2.py3-none-any.whl name: shellingham version: 1.5.4 @@ -3782,15 +3295,17 @@ packages: - python-multipart>=0.0.18 ; extra == 'full' - pyyaml ; extra == 'full' requires_python: '>=3.10' -- pypi: https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl - name: sympy - version: 1.14.0 - sha256: e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5 +- pypi: https://files.pythonhosted.org/packages/d7/c1/eb8f9debc45d3b7918a32ab756658a0904732f75e555402972246b0b8e71/tenacity-9.1.4-py3-none-any.whl + name: tenacity + version: 9.1.4 + sha256: 6095a360c919085f28c6527de529e76a06ad89b23659fa881ae0649b867a9d55 requires_dist: - - mpmath>=1.1.0,<1.4 - - pytest>=7.1.0 ; extra == 'dev' - - hypothesis>=6.70.0 ; extra == 'dev' - requires_python: '>=3.9' + - reno ; extra == 'doc' + - sphinx ; extra == 'doc' + - pytest ; extra == 'test' + - tornado>=4.5 ; extra == 'test' + - typeguard ; extra == 'test' + requires_python: '>=3.10' - pypi: https://files.pythonhosted.org/packages/f5/de/9341a6d7a8f1b448573bbf3425fa57669ac58258a667eb48a25dfe916d70/tiktoken-0.12.0-cp314-cp314-manylinux_2_28_x86_64.whl name: tiktoken version: 0.12.0 @@ -3832,40 +3347,6 @@ packages: - setuptools-rust ; extra == 'docs' - tokenizers[testing] ; extra == 'dev' requires_python: '>=3.9' -- pypi: https://files.pythonhosted.org/packages/47/e8/b98ca2d39b2e0e4730c0ee52537e488e7008025bc77ca89552ff91021f7c/torch-2.11.0-cp314-cp314-manylinux_2_28_x86_64.whl - name: torch - version: 2.11.0 - sha256: 4dc8b3809469b6c30b411bb8c4cad3828efd26236153d9beb6a3ec500f211a60 - requires_dist: - - filelock - - typing-extensions>=4.10.0 - - setuptools<82 - - sympy>=1.13.3 - - networkx>=2.5.1 - - jinja2 - - fsspec>=0.8.5 - - cuda-toolkit[cublas,cudart,cufft,cufile,cupti,curand,cusolver,cusparse,nvjitlink,nvrtc,nvtx]==13.0.2 ; sys_platform == 'linux' - - cuda-bindings>=13.0.3,<14 ; sys_platform == 'linux' - - nvidia-cudnn-cu13==9.19.0.56 ; sys_platform == 'linux' - - nvidia-cusparselt-cu13==0.8.0 ; sys_platform == 'linux' - - nvidia-nccl-cu13==2.28.9 ; sys_platform == 'linux' - - nvidia-nvshmem-cu13==3.4.5 ; sys_platform == 'linux' - - triton==3.6.0 ; sys_platform == 'linux' - - optree>=0.13.0 ; extra == 'optree' - - opt-einsum>=3.3 ; extra == 'opt-einsum' - - pyyaml ; extra == 'pyyaml' - requires_python: '>=3.10' -- pypi: https://files.pythonhosted.org/packages/a4/21/a2266f7f1b0e58e624ff15fd6f01041f59182c49551ece0db9a183071329/torchvision-0.26.0-cp314-cp314-manylinux_2_28_x86_64.whl - name: torchvision - version: 0.26.0 - sha256: 0f3e572efe62ad645017ea847e0b5e4f2f638d4e39f05bc011d1eb9ac68d4806 - requires_dist: - - numpy - - torch==2.11.0 - - pillow>=5.3.0,!=8.3.* - - gdown>=4.7.3 ; extra == 'gdown' - - scipy ; extra == 'scipy' - requires_python: '>=3.10,!=3.14.1' - pypi: https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl name: tqdm version: 4.67.3 @@ -3883,489 +3364,265 @@ packages: - requests ; extra == 'telegram' - ipywidgets>=6 ; extra == 'notebook' requires_python: '>=3.7' -- pypi: https://files.pythonhosted.org/packages/03/b8/e484ef633af3887baeeb4b6ad12743363af7cce68ae51e938e00aaa0529d/transformers-4.57.6-py3-none-any.whl +- pypi: https://files.pythonhosted.org/packages/e7/28/35f7411ff80a3640c1f4fc907dcbb6a65061ebb82f66950e38bfc9f7f740/transformers-5.5.0-py3-none-any.whl name: transformers - version: 4.57.6 - sha256: 4c9e9de11333ddfe5114bc872c9f370509198acf0b87a832a0ab9458e2bd0550 + version: 5.5.0 + sha256: 821a9ff0961abbb29eb1eb686d78df1c85929fdf213a3fe49dc6bd94f9efa944 requires_dist: - - filelock - - huggingface-hub>=0.34.0,<1.0 + - huggingface-hub>=1.5.0,<2.0 - numpy>=1.17 - packaging>=20.0 - pyyaml>=5.1 - - regex!=2019.12.17 - - requests + - regex>=2025.10.22 - tokenizers>=0.22.0,<=0.23.0 + - typer - safetensors>=0.4.3 - tqdm>=4.27 - - fugashi>=1.0 ; extra == 'ja' - - ipadic>=1.0.0,<2.0 ; extra == 'ja' - - unidic-lite>=1.0.7 ; extra == 'ja' - - unidic>=1.0.2 ; extra == 'ja' - - sudachipy>=0.6.6 ; extra == 'ja' - - sudachidict-core>=20220729 ; extra == 'ja' - - rhoknp>=1.1.0,<1.3.1 ; extra == 'ja' + - torch>=2.4 ; extra == 'torch' + - accelerate>=1.1.0 ; extra == 'torch' + - torchvision ; extra == 'vision' + - pillow>=10.0.1,<=15.0 ; extra == 'vision' + - torchaudio ; extra == 'audio' + - librosa ; extra == 'audio' + - pyctcdecode>=0.4.0 ; extra == 'audio' + - phonemizer ; extra == 'audio' + - av ; extra == 'video' + - timm>=1.0.23 ; extra == 'timm' + - datasets>=2.15.0 ; extra == 'quality' + - ruff==0.14.10 ; extra == 'quality' + - gitpython<3.1.19 ; extra == 'quality' + - urllib3<2.0.0 ; extra == 'quality' + - libcst ; extra == 'quality' + - rich ; extra == 'quality' + - ty==0.0.20 ; extra == 'quality' + - tomli ; extra == 'quality' + - hf-doc-builder ; extra == 'docs' + - kernels>=0.12.0,<0.13 ; extra == 'kernels' + - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'sentencepiece' + - protobuf ; extra == 'sentencepiece' + - tiktoken ; extra == 'tiktoken' + - blobfile ; extra == 'tiktoken' + - mistral-common[image]>=1.10.0 ; extra == 'mistral-common' + - jinja2>=3.1.0 ; extra == 'chat-template' + - jmespath>=1.0.1 ; extra == 'chat-template' - scikit-learn ; extra == 'sklearn' - - tensorflow>2.9,<2.16 ; extra == 'tf' - - onnxconverter-common ; extra == 'tf' - - tf2onnx ; extra == 'tf' - - tensorflow-text<2.16 ; extra == 'tf' - - keras-nlp>=0.3.1,<0.14.0 ; extra == 'tf' - - keras>2.9,<2.16 ; extra == 'tf-cpu' - - tensorflow-cpu>2.9,<2.16 ; extra == 'tf-cpu' - - onnxconverter-common ; extra == 'tf-cpu' - - tf2onnx ; extra == 'tf-cpu' - - tensorflow-text<2.16 ; extra == 'tf-cpu' - - keras-nlp>=0.3.1,<0.14.0 ; extra == 'tf-cpu' - - tensorflow-probability<0.24 ; extra == 'tf-cpu' - - torch>=2.2 ; extra == 'torch' - - accelerate>=0.26.0 ; extra == 'torch' - - accelerate>=0.26.0 ; extra == 'accelerate' - - hf-xet ; extra == 'hf-xet' + - accelerate>=1.1.0 ; extra == 'accelerate' - faiss-cpu ; extra == 'retrieval' - datasets>=2.15.0 ; extra == 'retrieval' - - jax>=0.4.1,<=0.4.13 ; extra == 'flax' - - jaxlib>=0.4.1,<=0.4.13 ; extra == 'flax' - - flax>=0.4.1,<=0.7.0 ; extra == 'flax' - - optax>=0.0.8,<=0.1.4 ; extra == 'flax' - - scipy<1.13.0 ; extra == 'flax' - - tokenizers>=0.22.0,<=0.23.0 ; extra == 'tokenizers' - - ftfy ; extra == 'ftfy' - - onnxruntime>=1.4.0 ; extra == 'onnxruntime' - - onnxruntime-tools>=1.4.2 ; extra == 'onnxruntime' - - onnxconverter-common ; extra == 'onnx' - - tf2onnx ; extra == 'onnx' - - onnxruntime>=1.4.0 ; extra == 'onnx' - - onnxruntime-tools>=1.4.2 ; extra == 'onnx' - - cookiecutter==1.7.3 ; extra == 'modelcreation' - sagemaker>=2.31.0 ; extra == 'sagemaker' - deepspeed>=0.9.3 ; extra == 'deepspeed' - - accelerate>=0.26.0 ; extra == 'deepspeed' + - accelerate>=1.1.0 ; extra == 'deepspeed' - optuna ; extra == 'optuna' - - ray[tune]>=2.7.0 ; extra == 'ray' - - sigopt ; extra == 'sigopt' - - kernels>=0.6.1,<=0.9 ; extra == 'hub-kernels' - - kernels>=0.6.1,<=0.9 ; extra == 'integrations' + - kernels>=0.12.0,<0.13 ; extra == 'integrations' - optuna ; extra == 'integrations' + - codecarbon>=2.8.1 ; extra == 'integrations' - ray[tune]>=2.7.0 ; extra == 'integrations' + - ray[tune]>=2.7.0 ; extra == 'ray' + - codecarbon>=2.8.1 ; extra == 'codecarbon' - openai>=1.98.0 ; extra == 'serving' - pydantic>=2 ; extra == 'serving' - uvicorn ; extra == 'serving' - fastapi ; extra == 'serving' - starlette ; extra == 'serving' - - torch>=2.2 ; extra == 'serving' - - accelerate>=0.26.0 ; extra == 'serving' - - librosa ; extra == 'audio' - - pyctcdecode>=0.4.0 ; extra == 'audio' - - phonemizer ; extra == 'audio' - - kenlm ; extra == 'audio' - - torchaudio ; extra == 'speech' - - librosa ; extra == 'speech' - - pyctcdecode>=0.4.0 ; extra == 'speech' - - phonemizer ; extra == 'speech' - - kenlm ; extra == 'speech' - - torchaudio ; extra == 'torch-speech' - - librosa ; extra == 'torch-speech' - - pyctcdecode>=0.4.0 ; extra == 'torch-speech' - - phonemizer ; extra == 'torch-speech' - - kenlm ; extra == 'torch-speech' - - librosa ; extra == 'tf-speech' - - pyctcdecode>=0.4.0 ; extra == 'tf-speech' - - phonemizer ; extra == 'tf-speech' - - kenlm ; extra == 'tf-speech' - - librosa ; extra == 'flax-speech' - - pyctcdecode>=0.4.0 ; extra == 'flax-speech' - - phonemizer ; extra == 'flax-speech' - - kenlm ; extra == 'flax-speech' - - pillow>=10.0.1,<=15.0 ; extra == 'vision' - - timm!=1.0.18,<=1.0.19 ; extra == 'timm' - - torchvision ; extra == 'torch-vision' - - pillow>=10.0.1,<=15.0 ; extra == 'torch-vision' - - natten>=0.14.6,<0.15.0 ; extra == 'natten' - - codecarbon>=2.8.1 ; extra == 'codecarbon' - - av ; extra == 'video' + - rich ; extra == 'serving' + - torch>=2.4 ; extra == 'serving' + - accelerate>=1.1.0 ; extra == 'serving' - num2words ; extra == 'num2words' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'sentencepiece' - - protobuf ; extra == 'sentencepiece' - - tiktoken ; extra == 'tiktoken' - - blobfile ; extra == 'tiktoken' - - mistral-common[opencv]>=1.6.3 ; extra == 'mistral-common' - - jinja2>=3.1.0 ; extra == 'chat-template' - - pytest>=7.2.0 ; extra == 'testing' - - pytest-asyncio ; extra == 'testing' + - optimum-benchmark>=0.3.0 ; extra == 'benchmark' + - fugashi>=1.0 ; extra == 'ja' + - ipadic>=1.0.0,<2.0 ; extra == 'ja' + - unidic-lite>=1.0.7 ; extra == 'ja' + - unidic>=1.0.2 ; extra == 'ja' + - rhoknp>=1.1.0,<1.3.1 ; extra == 'ja' + - sudachipy>=0.6.6 ; extra == 'ja' + - sudachidict-core>=20220729 ; extra == 'ja' + - opentelemetry-api ; extra == 'open-telemetry' + - opentelemetry-exporter-otlp ; extra == 'open-telemetry' + - opentelemetry-sdk ; extra == 'open-telemetry' + - pytest>=7.2.0,<9.0.0 ; extra == 'testing' + - pytest-asyncio>=1.2.0 ; extra == 'testing' + - pytest-random-order ; extra == 'testing' - pytest-rich ; extra == 'testing' - pytest-xdist ; extra == 'testing' - pytest-order ; extra == 'testing' - pytest-rerunfailures<16.0 ; extra == 'testing' + - pytest-timeout ; extra == 'testing' + - pytest-env ; extra == 'testing' - timeout-decorator ; extra == 'testing' - parameterized>=0.9 ; extra == 'testing' - psutil ; extra == 'testing' - - datasets>=2.15.0 ; extra == 'testing' - dill<0.3.5 ; extra == 'testing' - - evaluate>=0.2.0 ; extra == 'testing' - - pytest-timeout ; extra == 'testing' - - ruff==0.13.1 ; extra == 'testing' + - evaluate>=0.4.6 ; extra == 'testing' - rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1 ; extra == 'testing' - nltk<=3.8.1 ; extra == 'testing' - - gitpython<3.1.19 ; extra == 'testing' - sacremoses ; extra == 'testing' - rjieba ; extra == 'testing' - beautifulsoup4 ; extra == 'testing' - tensorboard ; extra == 'testing' - - pydantic>=2 ; extra == 'testing' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'testing' - sacrebleu>=1.4.12,<2.0.0 ; extra == 'testing' + - filelock ; extra == 'testing' + - hf-doc-builder ; extra == 'testing' + - datasets>=2.15.0 ; extra == 'testing' + - ruff==0.14.10 ; extra == 'testing' + - gitpython<3.1.19 ; extra == 'testing' + - urllib3<2.0.0 ; extra == 'testing' - libcst ; extra == 'testing' + - rich ; extra == 'testing' + - ty==0.0.20 ; extra == 'testing' + - tomli ; extra == 'testing' - faiss-cpu ; extra == 'testing' - datasets>=2.15.0 ; extra == 'testing' - - cookiecutter==1.7.3 ; extra == 'testing' - - mistral-common[opencv]>=1.6.3 ; extra == 'testing' + - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'testing' + - protobuf ; extra == 'testing' - openai>=1.98.0 ; extra == 'testing' - pydantic>=2 ; extra == 'testing' - uvicorn ; extra == 'testing' - fastapi ; extra == 'testing' - starlette ; extra == 'testing' - - torch>=2.2 ; extra == 'testing' - - accelerate>=0.26.0 ; extra == 'testing' + - rich ; extra == 'testing' + - torch>=2.4 ; extra == 'testing' + - accelerate>=1.1.0 ; extra == 'testing' + - mistral-common[image]>=1.10.0 ; extra == 'testing' - deepspeed>=0.9.3 ; extra == 'deepspeed-testing' - - accelerate>=0.26.0 ; extra == 'deepspeed-testing' - - pytest>=7.2.0 ; extra == 'deepspeed-testing' - - pytest-asyncio ; extra == 'deepspeed-testing' + - accelerate>=1.1.0 ; extra == 'deepspeed-testing' + - pytest>=7.2.0,<9.0.0 ; extra == 'deepspeed-testing' + - pytest-asyncio>=1.2.0 ; extra == 'deepspeed-testing' + - pytest-random-order ; extra == 'deepspeed-testing' - pytest-rich ; extra == 'deepspeed-testing' - pytest-xdist ; extra == 'deepspeed-testing' - pytest-order ; extra == 'deepspeed-testing' - pytest-rerunfailures<16.0 ; extra == 'deepspeed-testing' + - pytest-timeout ; extra == 'deepspeed-testing' + - pytest-env ; extra == 'deepspeed-testing' - timeout-decorator ; extra == 'deepspeed-testing' - parameterized>=0.9 ; extra == 'deepspeed-testing' - psutil ; extra == 'deepspeed-testing' - - datasets>=2.15.0 ; extra == 'deepspeed-testing' - dill<0.3.5 ; extra == 'deepspeed-testing' - - evaluate>=0.2.0 ; extra == 'deepspeed-testing' - - pytest-timeout ; extra == 'deepspeed-testing' - - ruff==0.13.1 ; extra == 'deepspeed-testing' + - evaluate>=0.4.6 ; extra == 'deepspeed-testing' - rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1 ; extra == 'deepspeed-testing' - nltk<=3.8.1 ; extra == 'deepspeed-testing' - - gitpython<3.1.19 ; extra == 'deepspeed-testing' - sacremoses ; extra == 'deepspeed-testing' - rjieba ; extra == 'deepspeed-testing' - beautifulsoup4 ; extra == 'deepspeed-testing' - tensorboard ; extra == 'deepspeed-testing' - - pydantic>=2 ; extra == 'deepspeed-testing' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'deepspeed-testing' - sacrebleu>=1.4.12,<2.0.0 ; extra == 'deepspeed-testing' + - filelock ; extra == 'deepspeed-testing' + - hf-doc-builder ; extra == 'deepspeed-testing' + - datasets>=2.15.0 ; extra == 'deepspeed-testing' + - ruff==0.14.10 ; extra == 'deepspeed-testing' + - gitpython<3.1.19 ; extra == 'deepspeed-testing' + - urllib3<2.0.0 ; extra == 'deepspeed-testing' - libcst ; extra == 'deepspeed-testing' + - rich ; extra == 'deepspeed-testing' + - ty==0.0.20 ; extra == 'deepspeed-testing' + - tomli ; extra == 'deepspeed-testing' - faiss-cpu ; extra == 'deepspeed-testing' - datasets>=2.15.0 ; extra == 'deepspeed-testing' - - cookiecutter==1.7.3 ; extra == 'deepspeed-testing' - - mistral-common[opencv]>=1.6.3 ; extra == 'deepspeed-testing' + - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'deepspeed-testing' + - protobuf ; extra == 'deepspeed-testing' - openai>=1.98.0 ; extra == 'deepspeed-testing' - pydantic>=2 ; extra == 'deepspeed-testing' - uvicorn ; extra == 'deepspeed-testing' - fastapi ; extra == 'deepspeed-testing' - starlette ; extra == 'deepspeed-testing' - - torch>=2.2 ; extra == 'deepspeed-testing' - - accelerate>=0.26.0 ; extra == 'deepspeed-testing' + - rich ; extra == 'deepspeed-testing' + - torch>=2.4 ; extra == 'deepspeed-testing' + - accelerate>=1.1.0 ; extra == 'deepspeed-testing' + - mistral-common[image]>=1.10.0 ; extra == 'deepspeed-testing' - optuna ; extra == 'deepspeed-testing' - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'deepspeed-testing' - protobuf ; extra == 'deepspeed-testing' - - ruff==0.13.1 ; extra == 'ruff' - - datasets>=2.15.0 ; extra == 'quality' - - ruff==0.13.1 ; extra == 'quality' - - gitpython<3.1.19 ; extra == 'quality' - - urllib3<2.0.0 ; extra == 'quality' - - libcst ; extra == 'quality' - - rich ; extra == 'quality' - - pandas<2.3.0 ; extra == 'quality' - - tensorflow>2.9,<2.16 ; extra == 'all' - - onnxconverter-common ; extra == 'all' - - tf2onnx ; extra == 'all' - - tensorflow-text<2.16 ; extra == 'all' - - keras-nlp>=0.3.1,<0.14.0 ; extra == 'all' - - torch>=2.2 ; extra == 'all' - - accelerate>=0.26.0 ; extra == 'all' - - jax>=0.4.1,<=0.4.13 ; extra == 'all' - - jaxlib>=0.4.1,<=0.4.13 ; extra == 'all' - - flax>=0.4.1,<=0.7.0 ; extra == 'all' - - optax>=0.0.8,<=0.1.4 ; extra == 'all' - - scipy<1.13.0 ; extra == 'all' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'all' - - protobuf ; extra == 'all' - - tokenizers>=0.22.0,<=0.23.0 ; extra == 'all' + - torch>=2.4 ; extra == 'all' + - accelerate>=1.1.0 ; extra == 'all' + - torchvision ; extra == 'all' + - pillow>=10.0.1,<=15.0 ; extra == 'all' - torchaudio ; extra == 'all' - librosa ; extra == 'all' - pyctcdecode>=0.4.0 ; extra == 'all' - phonemizer ; extra == 'all' - - kenlm ; extra == 'all' - - pillow>=10.0.1,<=15.0 ; extra == 'all' - - kernels>=0.6.1,<=0.9 ; extra == 'all' - - optuna ; extra == 'all' - - ray[tune]>=2.7.0 ; extra == 'all' - - timm!=1.0.18,<=1.0.19 ; extra == 'all' - - torchvision ; extra == 'all' - - pillow>=10.0.1,<=15.0 ; extra == 'all' - - codecarbon>=2.8.1 ; extra == 'all' - - accelerate>=0.26.0 ; extra == 'all' - av ; extra == 'all' - - num2words ; extra == 'all' - - mistral-common[opencv]>=1.6.3 ; extra == 'all' + - kernels>=0.12.0,<0.13 ; extra == 'all' + - timm>=1.0.23 ; extra == 'all' + - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'all' + - protobuf ; extra == 'all' + - tiktoken ; extra == 'all' + - blobfile ; extra == 'all' - jinja2>=3.1.0 ; extra == 'all' - - pytest>=7.2.0 ; extra == 'dev-torch' - - pytest-asyncio ; extra == 'dev-torch' - - pytest-rich ; extra == 'dev-torch' - - pytest-xdist ; extra == 'dev-torch' - - pytest-order ; extra == 'dev-torch' - - pytest-rerunfailures<16.0 ; extra == 'dev-torch' - - timeout-decorator ; extra == 'dev-torch' - - parameterized>=0.9 ; extra == 'dev-torch' - - psutil ; extra == 'dev-torch' - - datasets>=2.15.0 ; extra == 'dev-torch' - - dill<0.3.5 ; extra == 'dev-torch' - - evaluate>=0.2.0 ; extra == 'dev-torch' - - pytest-timeout ; extra == 'dev-torch' - - ruff==0.13.1 ; extra == 'dev-torch' - - rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1 ; extra == 'dev-torch' - - nltk<=3.8.1 ; extra == 'dev-torch' - - gitpython<3.1.19 ; extra == 'dev-torch' - - sacremoses ; extra == 'dev-torch' - - rjieba ; extra == 'dev-torch' - - beautifulsoup4 ; extra == 'dev-torch' - - tensorboard ; extra == 'dev-torch' - - pydantic>=2 ; extra == 'dev-torch' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'dev-torch' - - sacrebleu>=1.4.12,<2.0.0 ; extra == 'dev-torch' - - libcst ; extra == 'dev-torch' - - faiss-cpu ; extra == 'dev-torch' - - datasets>=2.15.0 ; extra == 'dev-torch' - - cookiecutter==1.7.3 ; extra == 'dev-torch' - - mistral-common[opencv]>=1.6.3 ; extra == 'dev-torch' - - openai>=1.98.0 ; extra == 'dev-torch' - - pydantic>=2 ; extra == 'dev-torch' - - uvicorn ; extra == 'dev-torch' - - fastapi ; extra == 'dev-torch' - - starlette ; extra == 'dev-torch' - - torch>=2.2 ; extra == 'dev-torch' - - accelerate>=0.26.0 ; extra == 'dev-torch' - - torch>=2.2 ; extra == 'dev-torch' - - accelerate>=0.26.0 ; extra == 'dev-torch' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'dev-torch' - - protobuf ; extra == 'dev-torch' - - tokenizers>=0.22.0,<=0.23.0 ; extra == 'dev-torch' - - torchaudio ; extra == 'dev-torch' - - librosa ; extra == 'dev-torch' - - pyctcdecode>=0.4.0 ; extra == 'dev-torch' - - phonemizer ; extra == 'dev-torch' - - kenlm ; extra == 'dev-torch' - - pillow>=10.0.1,<=15.0 ; extra == 'dev-torch' - - kernels>=0.6.1,<=0.9 ; extra == 'dev-torch' - - optuna ; extra == 'dev-torch' - - ray[tune]>=2.7.0 ; extra == 'dev-torch' - - timm!=1.0.18,<=1.0.19 ; extra == 'dev-torch' - - torchvision ; extra == 'dev-torch' - - pillow>=10.0.1,<=15.0 ; extra == 'dev-torch' - - codecarbon>=2.8.1 ; extra == 'dev-torch' - - datasets>=2.15.0 ; extra == 'dev-torch' - - ruff==0.13.1 ; extra == 'dev-torch' - - gitpython<3.1.19 ; extra == 'dev-torch' - - urllib3<2.0.0 ; extra == 'dev-torch' - - libcst ; extra == 'dev-torch' - - rich ; extra == 'dev-torch' - - pandas<2.3.0 ; extra == 'dev-torch' - - fugashi>=1.0 ; extra == 'dev-torch' - - ipadic>=1.0.0,<2.0 ; extra == 'dev-torch' - - unidic-lite>=1.0.7 ; extra == 'dev-torch' - - unidic>=1.0.2 ; extra == 'dev-torch' - - sudachipy>=0.6.6 ; extra == 'dev-torch' - - sudachidict-core>=20220729 ; extra == 'dev-torch' - - rhoknp>=1.1.0,<1.3.1 ; extra == 'dev-torch' - - scikit-learn ; extra == 'dev-torch' - - cookiecutter==1.7.3 ; extra == 'dev-torch' - - onnxruntime>=1.4.0 ; extra == 'dev-torch' - - onnxruntime-tools>=1.4.2 ; extra == 'dev-torch' - - num2words ; extra == 'dev-torch' - - pytest>=7.2.0 ; extra == 'dev-tensorflow' - - pytest-asyncio ; extra == 'dev-tensorflow' - - pytest-rich ; extra == 'dev-tensorflow' - - pytest-xdist ; extra == 'dev-tensorflow' - - pytest-order ; extra == 'dev-tensorflow' - - pytest-rerunfailures<16.0 ; extra == 'dev-tensorflow' - - timeout-decorator ; extra == 'dev-tensorflow' - - parameterized>=0.9 ; extra == 'dev-tensorflow' - - psutil ; extra == 'dev-tensorflow' - - datasets>=2.15.0 ; extra == 'dev-tensorflow' - - dill<0.3.5 ; extra == 'dev-tensorflow' - - evaluate>=0.2.0 ; extra == 'dev-tensorflow' - - pytest-timeout ; extra == 'dev-tensorflow' - - ruff==0.13.1 ; extra == 'dev-tensorflow' - - rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1 ; extra == 'dev-tensorflow' - - nltk<=3.8.1 ; extra == 'dev-tensorflow' - - gitpython<3.1.19 ; extra == 'dev-tensorflow' - - sacremoses ; extra == 'dev-tensorflow' - - rjieba ; extra == 'dev-tensorflow' - - beautifulsoup4 ; extra == 'dev-tensorflow' - - tensorboard ; extra == 'dev-tensorflow' - - pydantic>=2 ; extra == 'dev-tensorflow' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'dev-tensorflow' - - sacrebleu>=1.4.12,<2.0.0 ; extra == 'dev-tensorflow' - - libcst ; extra == 'dev-tensorflow' - - faiss-cpu ; extra == 'dev-tensorflow' - - datasets>=2.15.0 ; extra == 'dev-tensorflow' - - cookiecutter==1.7.3 ; extra == 'dev-tensorflow' - - mistral-common[opencv]>=1.6.3 ; extra == 'dev-tensorflow' - - openai>=1.98.0 ; extra == 'dev-tensorflow' - - pydantic>=2 ; extra == 'dev-tensorflow' - - uvicorn ; extra == 'dev-tensorflow' - - fastapi ; extra == 'dev-tensorflow' - - starlette ; extra == 'dev-tensorflow' - - torch>=2.2 ; extra == 'dev-tensorflow' - - accelerate>=0.26.0 ; extra == 'dev-tensorflow' - - tensorflow>2.9,<2.16 ; extra == 'dev-tensorflow' - - onnxconverter-common ; extra == 'dev-tensorflow' - - tf2onnx ; extra == 'dev-tensorflow' - - tensorflow-text<2.16 ; extra == 'dev-tensorflow' - - keras-nlp>=0.3.1,<0.14.0 ; extra == 'dev-tensorflow' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'dev-tensorflow' - - protobuf ; extra == 'dev-tensorflow' - - tokenizers>=0.22.0,<=0.23.0 ; extra == 'dev-tensorflow' - - pillow>=10.0.1,<=15.0 ; extra == 'dev-tensorflow' - - datasets>=2.15.0 ; extra == 'dev-tensorflow' - - ruff==0.13.1 ; extra == 'dev-tensorflow' - - gitpython<3.1.19 ; extra == 'dev-tensorflow' - - urllib3<2.0.0 ; extra == 'dev-tensorflow' - - libcst ; extra == 'dev-tensorflow' - - rich ; extra == 'dev-tensorflow' - - pandas<2.3.0 ; extra == 'dev-tensorflow' - - scikit-learn ; extra == 'dev-tensorflow' - - cookiecutter==1.7.3 ; extra == 'dev-tensorflow' - - onnxconverter-common ; extra == 'dev-tensorflow' - - tf2onnx ; extra == 'dev-tensorflow' - - onnxruntime>=1.4.0 ; extra == 'dev-tensorflow' - - onnxruntime-tools>=1.4.2 ; extra == 'dev-tensorflow' - - librosa ; extra == 'dev-tensorflow' - - pyctcdecode>=0.4.0 ; extra == 'dev-tensorflow' - - phonemizer ; extra == 'dev-tensorflow' - - kenlm ; extra == 'dev-tensorflow' - - tensorflow>2.9,<2.16 ; extra == 'dev' - - onnxconverter-common ; extra == 'dev' - - tf2onnx ; extra == 'dev' - - tensorflow-text<2.16 ; extra == 'dev' - - keras-nlp>=0.3.1,<0.14.0 ; extra == 'dev' - - torch>=2.2 ; extra == 'dev' - - accelerate>=0.26.0 ; extra == 'dev' - - jax>=0.4.1,<=0.4.13 ; extra == 'dev' - - jaxlib>=0.4.1,<=0.4.13 ; extra == 'dev' - - flax>=0.4.1,<=0.7.0 ; extra == 'dev' - - optax>=0.0.8,<=0.1.4 ; extra == 'dev' - - scipy<1.13.0 ; extra == 'dev' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'dev' - - protobuf ; extra == 'dev' - - tokenizers>=0.22.0,<=0.23.0 ; extra == 'dev' + - jmespath>=1.0.1 ; extra == 'all' + - num2words ; extra == 'all' + - mistral-common[image]>=1.10.0 ; extra == 'all' + - torch>=2.4 ; extra == 'dev' + - accelerate>=1.1.0 ; extra == 'dev' + - torchvision ; extra == 'dev' + - pillow>=10.0.1,<=15.0 ; extra == 'dev' - torchaudio ; extra == 'dev' - librosa ; extra == 'dev' - pyctcdecode>=0.4.0 ; extra == 'dev' - phonemizer ; extra == 'dev' - - kenlm ; extra == 'dev' - - pillow>=10.0.1,<=15.0 ; extra == 'dev' - - kernels>=0.6.1,<=0.9 ; extra == 'dev' - - optuna ; extra == 'dev' - - ray[tune]>=2.7.0 ; extra == 'dev' - - timm!=1.0.18,<=1.0.19 ; extra == 'dev' - - torchvision ; extra == 'dev' - - pillow>=10.0.1,<=15.0 ; extra == 'dev' - - codecarbon>=2.8.1 ; extra == 'dev' - - accelerate>=0.26.0 ; extra == 'dev' - av ; extra == 'dev' - - num2words ; extra == 'dev' - - mistral-common[opencv]>=1.6.3 ; extra == 'dev' + - kernels>=0.12.0,<0.13 ; extra == 'dev' + - timm>=1.0.23 ; extra == 'dev' + - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'dev' + - protobuf ; extra == 'dev' + - tiktoken ; extra == 'dev' + - blobfile ; extra == 'dev' - jinja2>=3.1.0 ; extra == 'dev' - - pytest>=7.2.0 ; extra == 'dev' - - pytest-asyncio ; extra == 'dev' + - jmespath>=1.0.1 ; extra == 'dev' + - num2words ; extra == 'dev' + - mistral-common[image]>=1.10.0 ; extra == 'dev' + - pytest>=7.2.0,<9.0.0 ; extra == 'dev' + - pytest-asyncio>=1.2.0 ; extra == 'dev' + - pytest-random-order ; extra == 'dev' - pytest-rich ; extra == 'dev' - pytest-xdist ; extra == 'dev' - pytest-order ; extra == 'dev' - pytest-rerunfailures<16.0 ; extra == 'dev' + - pytest-timeout ; extra == 'dev' + - pytest-env ; extra == 'dev' - timeout-decorator ; extra == 'dev' - parameterized>=0.9 ; extra == 'dev' - psutil ; extra == 'dev' - - datasets>=2.15.0 ; extra == 'dev' - dill<0.3.5 ; extra == 'dev' - - evaluate>=0.2.0 ; extra == 'dev' - - pytest-timeout ; extra == 'dev' - - ruff==0.13.1 ; extra == 'dev' + - evaluate>=0.4.6 ; extra == 'dev' - rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1 ; extra == 'dev' - nltk<=3.8.1 ; extra == 'dev' - - gitpython<3.1.19 ; extra == 'dev' - sacremoses ; extra == 'dev' - rjieba ; extra == 'dev' - beautifulsoup4 ; extra == 'dev' - tensorboard ; extra == 'dev' - - pydantic>=2 ; extra == 'dev' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'dev' - sacrebleu>=1.4.12,<2.0.0 ; extra == 'dev' + - filelock ; extra == 'dev' + - hf-doc-builder ; extra == 'dev' + - datasets>=2.15.0 ; extra == 'dev' + - ruff==0.14.10 ; extra == 'dev' + - gitpython<3.1.19 ; extra == 'dev' + - urllib3<2.0.0 ; extra == 'dev' - libcst ; extra == 'dev' + - rich ; extra == 'dev' + - ty==0.0.20 ; extra == 'dev' + - tomli ; extra == 'dev' - faiss-cpu ; extra == 'dev' - datasets>=2.15.0 ; extra == 'dev' - - cookiecutter==1.7.3 ; extra == 'dev' - - mistral-common[opencv]>=1.6.3 ; extra == 'dev' + - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'dev' + - protobuf ; extra == 'dev' - openai>=1.98.0 ; extra == 'dev' - pydantic>=2 ; extra == 'dev' - uvicorn ; extra == 'dev' - fastapi ; extra == 'dev' - starlette ; extra == 'dev' - - torch>=2.2 ; extra == 'dev' - - accelerate>=0.26.0 ; extra == 'dev' - - datasets>=2.15.0 ; extra == 'dev' - - ruff==0.13.1 ; extra == 'dev' - - gitpython<3.1.19 ; extra == 'dev' - - urllib3<2.0.0 ; extra == 'dev' - - libcst ; extra == 'dev' - rich ; extra == 'dev' - - pandas<2.3.0 ; extra == 'dev' + - torch>=2.4 ; extra == 'dev' + - accelerate>=1.1.0 ; extra == 'dev' + - mistral-common[image]>=1.10.0 ; extra == 'dev' - fugashi>=1.0 ; extra == 'dev' - ipadic>=1.0.0,<2.0 ; extra == 'dev' - unidic-lite>=1.0.7 ; extra == 'dev' - unidic>=1.0.2 ; extra == 'dev' + - rhoknp>=1.1.0,<1.3.1 ; extra == 'dev' - sudachipy>=0.6.6 ; extra == 'dev' - sudachidict-core>=20220729 ; extra == 'dev' - - rhoknp>=1.1.0,<1.3.1 ; extra == 'dev' - scikit-learn ; extra == 'dev' - - cookiecutter==1.7.3 ; extra == 'dev' - - filelock ; extra == 'torchhub' - - huggingface-hub>=0.34.0,<1.0 ; extra == 'torchhub' - - importlib-metadata ; extra == 'torchhub' - - numpy>=1.17 ; extra == 'torchhub' - - packaging>=20.0 ; extra == 'torchhub' - - protobuf ; extra == 'torchhub' - - regex!=2019.12.17 ; extra == 'torchhub' - - requests ; extra == 'torchhub' - - sentencepiece>=0.1.91,!=0.1.92 ; extra == 'torchhub' - - torch>=2.2 ; extra == 'torchhub' - - tokenizers>=0.22.0,<=0.23.0 ; extra == 'torchhub' - - tqdm>=4.27 ; extra == 'torchhub' - - optimum-benchmark>=0.3.0 ; extra == 'benchmark' - - opentelemetry-api ; extra == 'open-telemetry' - - opentelemetry-exporter-otlp ; extra == 'open-telemetry' - - opentelemetry-sdk ; extra == 'open-telemetry' - requires_python: '>=3.9.0' -- pypi: https://files.pythonhosted.org/packages/df/3d/9e7eee57b37c80cec63322c0231bb6da3cfe535a91d7a4d64896fcb89357/triton-3.6.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl - name: triton - version: 3.6.0 - sha256: a17a5d5985f0ac494ed8a8e54568f092f7057ef60e1b0fa09d3fd1512064e803 - requires_dist: - - importlib-metadata ; python_full_version < '3.10' - - cmake>=3.20,<4.0 ; extra == 'build' - - lit ; extra == 'build' - - autopep8 ; extra == 'tests' - - isort ; extra == 'tests' - - numpy ; extra == 'tests' - - pytest ; extra == 'tests' - - pytest-forked ; extra == 'tests' - - pytest-xdist ; extra == 'tests' - - scipy>=1.7.1 ; extra == 'tests' - - llnl-hatchet ; extra == 'tests' - - matplotlib ; extra == 'tutorials' - - pandas ; extra == 'tutorials' - - tabulate ; extra == 'tutorials' - requires_python: '>=3.10,<3.15' + requires_python: '>=3.10.0' - pypi: https://files.pythonhosted.org/packages/3a/7a/882d99539b19b1490cac5d77c67338d126e4122c8276bf640e411650c830/twine-6.2.0-py3-none-any.whl name: twine version: 6.2.0 diff --git a/pyproject.toml b/pyproject.toml index 013c748..3ef6d83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "churro-ocr" version = "0.2.0" -description = "Library-first OCR and layout detection for historical documents" +description = "OCR for historical documents" readme = { file = "docs/pypi.md", content-type = "text/markdown" } requires-python = ">=3.12" license = "Apache-2.0" @@ -20,6 +20,7 @@ dependencies = [ "loguru>=0.7.2,<1", "Pillow>=10.4.0,<12", "rich>=13.9.2,<14", + "tenacity>=9.1.2,<10", "typer>=0.12.3,<1", ] @@ -33,13 +34,7 @@ azure = [ ] hf = [ "qwen-vl-utils", - "transformers[torch]>=4.57.0,<5", - "torchvision", -] -vllm = [ - "transformers>=4.57.0,<5", - "torchvision", - "vllm>=0.18,<1", + "transformers>=5,<6", ] mistral = [ "mistralai>=1.6.0,<2", @@ -55,9 +50,7 @@ all = [ "litellm[caching]==1.82.3", "azure-ai-documentintelligence==1.0.2", "qwen-vl-utils", - "transformers[torch]>=4.57.0,<5", - "torchvision", - "vllm>=0.18,<1", + "transformers>=5,<6", "mistralai>=1.6.0,<2", "pypdfium2>=5,<6", ] @@ -89,7 +82,7 @@ dev-test = [ dev-tooling = [ "build", "pre-commit>=4,<5", - "ruff", + "ruff==0.15.9", "twine", "ty>=0.0.28,<0.0.29", ] diff --git a/ruff.toml b/ruff.toml index 3fb075f..869de7c 100644 --- a/ruff.toml +++ b/ruff.toml @@ -1,5 +1,5 @@ line-length = 110 -required-version = "==0.15.7" +required-version = "==0.15.9" [lint] select = ["E", "F", "I", "B", "UP", "ASYNC", "SIM", "D"] diff --git a/scripts/package_check.py b/scripts/package_check.py index 8e23c23..6dd41cd 100644 --- a/scripts/package_check.py +++ b/scripts/package_check.py @@ -35,7 +35,6 @@ "local", "mistral", "pdf", - "vllm", } FORBIDDEN_ARTIFACT_SEGMENTS = ("/tests/", "/tooling/", "/scripts/") FORBIDDEN_ARTIFACT_SUFFIXES = ("PYPI_AUDIT.md",) @@ -125,11 +124,7 @@ def _assert_metadata(metadata_message: Message, entry_points_text: str) -> None: if provides_extra != EXPECTED_EXTRAS: raise RuntimeError(f"Unexpected extras set: {sorted(provides_extra)!r}.") - requires_dist = metadata_message.get_all("Requires-Dist", []) - if not any( - requirement.startswith("vllm") and 'extra == "all"' in requirement for requirement in requires_dist - ): - raise RuntimeError("The all extra does not include vllm.") + _assert_local_runtime_packaging_policy(metadata_message) if ( "[console_scripts]" not in entry_points_text @@ -138,6 +133,34 @@ def _assert_metadata(metadata_message: Message, entry_points_text: str) -> None: raise RuntimeError("Console script entry point is missing or incorrect.") +def _iter_requirements_for_extra(metadata_message: Message, extra: str) -> list[Requirement]: + requirements: list[Requirement] = [] + for requirement_text in metadata_message.get_all("Requires-Dist", []): + parsed = Requirement(requirement_text) + marker_text = str(parsed.marker) if parsed.marker is not None else "" + if f'extra == "{extra}"' in marker_text: + requirements.append(parsed) + return requirements + + +def _assert_local_runtime_packaging_policy(metadata_message: Message) -> None: + disallowed_runtime_reqs: list[str] = [] + for extra in EXPECTED_EXTRAS: + for requirement in _iter_requirements_for_extra(metadata_message, extra): + normalized_name = requirement.name.replace("_", "-").lower() + if normalized_name in {"torch", "torchvision", "vllm"}: + disallowed_runtime_reqs.append(f"{extra}:{requirement}") + continue + if normalized_name == "transformers" and "torch" in requirement.extras: + disallowed_runtime_reqs.append(f"{extra}:{requirement}") + if disallowed_runtime_reqs: + formatted = ", ".join(sorted(disallowed_runtime_reqs)) + raise RuntimeError( + "PyPI extras for active-environment runtimes must not pin local PyTorch or vLLM runtimes. " + f"Found disallowed requirements: {formatted}." + ) + + def _assert_runtime_only_artifacts(wheel: Path, sdist: Path) -> None: with zipfile.ZipFile(wheel) as zip_file: wheel_names = zip_file.namelist() diff --git a/src/churro_ocr/_internal/install.py b/src/churro_ocr/_internal/install.py new file mode 100644 index 0000000..e1e7ff0 --- /dev/null +++ b/src/churro_ocr/_internal/install.py @@ -0,0 +1,145 @@ +"""UV-backed runtime installation helpers.""" + +from __future__ import annotations + +import shutil +import subprocess +import sys +from dataclasses import dataclass +from importlib import metadata +from typing import Final + +from churro_ocr.errors import ConfigurationError + +PROJECT_DISTRIBUTION_NAME: Final[str] = "churro-ocr" +INSTALL_TARGETS: Final[tuple[str, ...]] = ( + "llm", + "local", + "hf", + "azure", + "mistral", + "pdf", + "all", +) +_CURRENT_ENV_TARGET_EXTRAS: Final[dict[str, tuple[str, ...]]] = { + "llm": ("llm",), + "local": ("local",), + "hf": ("hf",), + "azure": ("azure",), + "mistral": ("mistral",), + "pdf": ("pdf",), + "all": ("llm", "local", "hf", "azure", "mistral", "pdf"), +} +_PYTORCH_TARGETS: Final[frozenset[str]] = frozenset({"hf", "all"}) +_PYTORCH_PACKAGES: Final[tuple[str, ...]] = ("torch", "torchvision") + + +@dataclass(frozen=True, slots=True) +class RuntimeInstallResult: + """Summary of a completed runtime installation.""" + + target: str + executed_commands: tuple[tuple[str, ...], ...] + notes: tuple[str, ...] = () + + +def install_command_hint(target: str) -> str: + """Return a short user-facing install hint for a runtime target.""" + return f"Run `churro-ocr install {target}`." + + +def install_runtime_dependencies( + *, + target: str, + torch_backend: str = "auto", +) -> RuntimeInstallResult: + """Install one of the supported optional runtime targets with UV.""" + normalized_target = target.strip().lower() + if normalized_target not in INSTALL_TARGETS: + supported = ", ".join(INSTALL_TARGETS) + raise ConfigurationError(f"Unknown install target '{target}'. Choose one of: {supported}.") + + uv_executable = _require_uv_executable() + executed_commands: list[tuple[str, ...]] = [] + notes: list[str] = [] + + current_env_extras = _CURRENT_ENV_TARGET_EXTRAS.get(normalized_target, ()) + requirements = _requirements_for_extra(current_env_extras) + if requirements: + executed_commands.append( + _run_command( + [ + uv_executable, + "pip", + "install", + "--python", + sys.executable, + "--upgrade", + *requirements, + ] + ) + ) + + if normalized_target in _PYTORCH_TARGETS: + executed_commands.append( + _run_command( + [ + uv_executable, + "pip", + "install", + "--python", + sys.executable, + "--upgrade", + f"--torch-backend={torch_backend}", + *_PYTORCH_PACKAGES, + ] + ) + ) + + return RuntimeInstallResult( + target=normalized_target, + executed_commands=tuple(executed_commands), + notes=tuple(notes), + ) + + +def _require_uv_executable() -> str: + uv_executable = shutil.which("uv") + if uv_executable is None: + raise ConfigurationError( + "`churro-ocr install` requires `uv` on PATH. Install uv and rerun the command." + ) + return uv_executable + + +def _distribution_requirements() -> list[str]: + try: + distribution = metadata.distribution(PROJECT_DISTRIBUTION_NAME) + except metadata.PackageNotFoundError as exc: # pragma: no cover - depends on install mode + raise ConfigurationError( + "The Churro installer must run from an installed `churro-ocr` environment." + ) from exc + return list(distribution.requires or []) + + +def _requirements_for_extra(extras: tuple[str, ...]) -> list[str]: + requirements = _distribution_requirements() + selected: list[str] = [] + for extra in extras: + for requirement_text in requirements: + requirement, _, marker = requirement_text.partition(";") + if f'extra == "{extra}"' not in marker: + continue + normalized_requirement = requirement.strip() + if normalized_requirement and normalized_requirement not in selected: + selected.append(normalized_requirement) + return selected + + +def _run_command(command: list[str]) -> tuple[str, ...]: + try: + subprocess.run(command, check=True) + except (OSError, subprocess.CalledProcessError) as exc: + rendered_command = " ".join(command) + raise ConfigurationError(f"Command failed: {rendered_command}") from exc + return tuple(command) diff --git a/src/churro_ocr/_internal/litellm.py b/src/churro_ocr/_internal/litellm.py index 967d3c3..eaed37c 100644 --- a/src/churro_ocr/_internal/litellm.py +++ b/src/churro_ocr/_internal/litellm.py @@ -12,6 +12,8 @@ from PIL import Image from churro_ocr._internal.image import image_to_base64 +from churro_ocr._internal.install import install_command_hint +from churro_ocr._internal.retry import retry_api_call from churro_ocr.errors import ConfigurationError, ProviderError from churro_ocr.providers.specs import LiteLLMTransportConfig from churro_ocr.templates import OCRConversation @@ -28,8 +30,7 @@ def _ensure_initialized() -> None: import litellm except ImportError as exc: # pragma: no cover - optional extra path raise ConfigurationError( - "LiteLLM-backed providers require the 'llm' extra. " - 'Install with `pip install "churro-ocr[llm]"`.' + f"LiteLLM-backed providers require the `llm` runtime. {install_command_hint('llm')}" ) from exc litellm_any = cast(Any, litellm) @@ -161,6 +162,7 @@ async def complete_text( messages: list[dict[str, Any]], timeout_seconds: int = 600, output_json: bool = False, + allow_empty: bool = False, ) -> str: """Run a LiteLLM completion and return the text content.""" if self._config.cache_dir is not None: @@ -186,15 +188,26 @@ async def complete_text( kwargs.update(self._config.completion_kwargs) try: - response = await acompletion(**kwargs) + response = await retry_api_call( + lambda: acompletion(**kwargs), + operation_name="LiteLLM request", + context=f"for model '{model}'", + ) except Exception as exc: # pragma: no cover - provider-specific failure path raise ProviderError(f"LiteLLM request failed for model '{model}': {exc}") from exc self._record_response_cost(model=model, response=response) answer = response.choices[0].message.content - if not isinstance(answer, str) or not answer.strip(): + if isinstance(answer, str): + if answer.strip(): + return answer + if allow_empty: + return "" + elif answer is None and allow_empty: + return "" + if not isinstance(answer, str): raise ProviderError(f"LiteLLM returned empty output for model '{model}'.") - return answer + raise ProviderError(f"LiteLLM returned empty output for model '{model}'.") def _resolved_image_detail(self) -> str | None: return "high" if self._config.image_detail is None else self._config.image_detail @@ -328,6 +341,7 @@ async def complete_text( api_version: str | None = None, timeout_seconds: int = 600, output_json: bool = False, + allow_empty: bool = False, completion_kwargs: dict[str, object] | None = None, ) -> str: """Run a LiteLLM completion and return the text content.""" @@ -344,6 +358,7 @@ async def complete_text( messages=messages, timeout_seconds=timeout_seconds, output_json=output_json, + allow_empty=allow_empty, ) diff --git a/src/churro_ocr/_internal/pdf.py b/src/churro_ocr/_internal/pdf.py index 4f89194..770819d 100644 --- a/src/churro_ocr/_internal/pdf.py +++ b/src/churro_ocr/_internal/pdf.py @@ -6,6 +6,7 @@ from PIL import Image +from churro_ocr._internal.install import install_command_hint from churro_ocr.errors import ConfigurationError @@ -15,7 +16,7 @@ def rasterize_pdf(path: str | Path, *, dpi: int = 300) -> list[Image.Image]: import pypdfium2 except ImportError as exc: # pragma: no cover - depends on optional extra raise ConfigurationError( - "PDF support requires the 'pdf' extra. Install with `pip install \"churro-ocr[pdf]\"`." + f"PDF support requires the `pdf` runtime. {install_command_hint('pdf')}" ) from exc resolved = Path(path) diff --git a/src/churro_ocr/_internal/retry.py b/src/churro_ocr/_internal/retry.py new file mode 100644 index 0000000..4595461 --- /dev/null +++ b/src/churro_ocr/_internal/retry.py @@ -0,0 +1,202 @@ +"""Shared retry helpers for provider API calls.""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Awaitable, Callable, Mapping +from typing import cast + +from tenacity import AsyncRetrying, RetryCallState, retry_if_exception, stop_after_attempt + +logger = logging.getLogger(__name__) + +DEFAULT_MAX_ATTEMPTS = 6 +DEFAULT_INITIAL_BACKOFF_SECONDS = 1.0 +DEFAULT_MAX_BACKOFF_SECONDS = 16.0 +TRANSIENT_STATUS_CODES = frozenset({408, 429, 500, 502, 503, 504, 520, 521, 522, 524}) +RETRYABLE_EXCEPTION_CLASS_NAMES = frozenset( + { + "APIConnectionError", + "APITimeoutError", + "ConnectError", + "ConnectTimeout", + "PoolTimeout", + "RateLimitError", + "ReadTimeout", + "RemoteProtocolError", + "ServiceRequestError", + "ServiceResponseError", + "WriteTimeout", + } +) +RETRYABLE_EXCEPTION_MODULE_PREFIXES = ("httpcore", "httpx") + +retry_sleep = asyncio.sleep + +type RetryPredicate = Callable[[BaseException], bool] + + +def _coerce_status_code(value: object) -> int | None: + if isinstance(value, int): + return value + if isinstance(value, float) and value.is_integer(): + return int(value) + return None + + +def _headers_from_candidate(value: object) -> Mapping[str, object] | None: + if isinstance(value, Mapping): + return cast("Mapping[str, object]", value) + return None + + +def get_error_status_code(exc: BaseException) -> int | None: + """Extract an HTTP-like status code from a provider exception when available.""" + status_code = _coerce_status_code(getattr(exc, "status_code", None)) + if status_code is not None: + return status_code + + for attribute_name in ("response", "raw_response"): + response = getattr(exc, attribute_name, None) + status_code = _coerce_status_code(getattr(response, "status_code", None)) + if status_code is not None: + return status_code + return None + + +def get_error_retry_after_seconds(exc: BaseException) -> float | None: + """Extract a retry delay from response headers when one is available.""" + headers = _headers_from_candidate(getattr(exc, "headers", None)) + if headers is None: + for attribute_name in ("response", "raw_response"): + response = getattr(exc, attribute_name, None) + headers = _headers_from_candidate(getattr(response, "headers", None)) + if headers is not None: + break + if headers is None: + return None + + retry_after = headers.get("retry-after") + if retry_after is None: + retry_after = headers.get("Retry-After") + if retry_after is None: + return None + if not isinstance(retry_after, str | int | float): + return None + try: + return max(0.0, float(retry_after)) + except ValueError: + return None + + +def compute_retry_delay_seconds( + exc: BaseException, + *, + attempt_number: int, + initial_backoff_seconds: float = DEFAULT_INITIAL_BACKOFF_SECONDS, + max_backoff_seconds: float = DEFAULT_MAX_BACKOFF_SECONDS, +) -> float: + """Compute the retry delay for a failed provider request.""" + retry_after = get_error_retry_after_seconds(exc) + if retry_after is not None: + return retry_after + return min( + initial_backoff_seconds * (2 ** max(0, attempt_number - 1)), + max_backoff_seconds, + ) + + +def is_retryable_api_error(exc: BaseException) -> bool: + """Return whether a provider exception should be retried.""" + if isinstance(exc, TimeoutError): + return True + + status_code = get_error_status_code(exc) + if status_code is not None: + return status_code in TRANSIENT_STATUS_CODES + + exc_type = exc.__class__ + if exc_type.__name__ in RETRYABLE_EXCEPTION_CLASS_NAMES: + return True + + module_name = exc_type.__module__ + return any(module_name.startswith(prefix) for prefix in RETRYABLE_EXCEPTION_MODULE_PREFIXES) + + +def _build_before_sleep_callback( + *, + operation_name: str, + context: str | None, + max_attempts: int, +) -> Callable[[RetryCallState], None]: + message_context = f" {context}" if context else "" + + def _before_sleep(retry_state: RetryCallState) -> None: + outcome = retry_state.outcome + exc = outcome.exception() if outcome is not None and outcome.failed else None + if exc is None: + return + delay_seconds = retry_state.next_action.sleep if retry_state.next_action is not None else 0.0 + logger.warning( + "Transient %s failure%s (status=%s, attempt=%s/%s); retrying in %.1fs.", + operation_name, + message_context, + get_error_status_code(exc) or "unknown", + retry_state.attempt_number, + max_attempts, + delay_seconds, + ) + + return _before_sleep + + +async def retry_api_call[T]( + fn: Callable[[], Awaitable[T]], + *, + operation_name: str, + context: str | None = None, + max_attempts: int = DEFAULT_MAX_ATTEMPTS, + retry_filter: RetryPredicate = is_retryable_api_error, + initial_backoff_seconds: float = DEFAULT_INITIAL_BACKOFF_SECONDS, + max_backoff_seconds: float = DEFAULT_MAX_BACKOFF_SECONDS, +) -> T: + """Execute an async provider request with shared retry behavior.""" + retrying = AsyncRetrying( + reraise=True, + stop=stop_after_attempt(max_attempts), + retry=retry_if_exception(retry_filter), + wait=lambda retry_state: ( + 0.0 + if retry_state.outcome is None or not retry_state.outcome.failed + else compute_retry_delay_seconds( + retry_state.outcome.exception(), + attempt_number=retry_state.attempt_number, + initial_backoff_seconds=initial_backoff_seconds, + max_backoff_seconds=max_backoff_seconds, + ) + ), + before_sleep=_build_before_sleep_callback( + operation_name=operation_name, + context=context, + max_attempts=max_attempts, + ), + sleep=retry_sleep, + ) + + async for attempt in retrying: + with attempt: + return await fn() + + raise AssertionError("AsyncRetrying exited without returning or raising.") + + +__all__ = [ + "DEFAULT_MAX_ATTEMPTS", + "compute_retry_delay_seconds", + "get_error_retry_after_seconds", + "get_error_status_code", + "is_retryable_api_error", + "retry_api_call", + "retry_sleep", +] diff --git a/src/churro_ocr/cli.py b/src/churro_ocr/cli.py index 8715bcd..9a00b21 100644 --- a/src/churro_ocr/cli.py +++ b/src/churro_ocr/cli.py @@ -6,6 +6,11 @@ import typer +from churro_ocr._internal.install import ( + INSTALL_TARGETS, + install_runtime_dependencies, +) +from churro_ocr.errors import ConfigurationError from churro_ocr.ocr import OCRClient from churro_ocr.page_detection import DocumentPage, DocumentPageDetector, PageDetectionRequest from churro_ocr.providers import ( @@ -17,12 +22,18 @@ MistralOptions, OCRBackendSpec, OpenAICompatibleOptions, - VLLMOptions, build_ocr_backend, ) +from churro_ocr.providers.specs import MISTRAL_OCR_MODEL_IDS, validate_mistral_ocr_model app = typer.Typer(help="churro-ocr library-first CLI") +_INSTALL_TARGET_METAVAR = "{" + "|".join(INSTALL_TARGETS) + "}" +_MISTRAL_MODEL_OPTION_ERROR = ( + "--model is required for backend=mistral and must be one of: " + + ", ".join(MISTRAL_OCR_MODEL_IDS) +) + def _build_ocr_backend( *, @@ -48,10 +59,8 @@ def _build_ocr_backend( ) ) if backend == "openai-compatible": - if not model or not base_url or not api_key: - raise typer.BadParameter( - "--model, --base-url, and --api-key are required for backend=openai-compatible" - ) + if not model or not base_url: + raise typer.BadParameter("--model and --base-url are required for backend=openai-compatible") return build_ocr_backend( OCRBackendSpec( provider="openai-compatible", @@ -80,10 +89,14 @@ def _build_ocr_backend( if backend == "mistral": if not api_key: raise typer.BadParameter("--api-key is required for backend=mistral") + try: + mistral_model = validate_mistral_ocr_model(model) + except ConfigurationError as exc: + raise typer.BadParameter(_MISTRAL_MODEL_OPTION_ERROR) from exc return build_ocr_backend( OCRBackendSpec( provider="mistral", - model=model or "mistral-ocr-latest", + model=mistral_model, options=MistralOptions(api_key=api_key), ) ) @@ -97,16 +110,6 @@ def _build_ocr_backend( options=HuggingFaceOptions(model_kwargs={"device_map": "auto", "torch_dtype": "auto"}), ) ) - if backend == "vllm": - if not model: - raise typer.BadParameter("--model is required for backend=vllm") - return build_ocr_backend( - OCRBackendSpec( - provider="vllm", - model=model, - options=VLLMOptions(), - ) - ) raise typer.BadParameter(f"Unsupported backend: {backend}") @@ -230,6 +233,33 @@ def extract_pages_command( typer.echo(str(output_path)) +@app.command("install") +def install_command( + target: str = typer.Argument( + ..., + metavar=_INSTALL_TARGET_METAVAR, + help="Runtime target to install into the active environment with uv.", + ), + torch_backend: str = typer.Option( + "auto", + help="PyTorch backend passed through to uv when a local runtime needs torch.", + ), +) -> None: + """Install optional runtime dependencies with uv.""" + try: + result = install_runtime_dependencies( + target=target, + torch_backend=torch_backend, + ) + except ConfigurationError as exc: + typer.echo(str(exc), err=True) + raise typer.Exit(code=1) from exc + + typer.echo(f"Installed runtime target: {result.target}") + for note in result.notes: + typer.echo(note) + + def main() -> None: """Console-script entrypoint.""" app() diff --git a/src/churro_ocr/prompts/__init__.py b/src/churro_ocr/prompts/__init__.py index 578679d..19db07e 100644 --- a/src/churro_ocr/prompts/__init__.py +++ b/src/churro_ocr/prompts/__init__.py @@ -4,18 +4,26 @@ DEFAULT_BOUNDARY_DETECTION_PROMPT, ) from churro_ocr.prompts.ocr import ( + CHANDRA_OCR_LAYOUT_PROMPT, DEFAULT_MARKDOWN_OCR_USER_PROMPT, DEFAULT_OCR_OUTPUT_TAG, DEFAULT_OCR_SYSTEM_PROMPT, DEFAULT_OCR_USER_PROMPT, + OLMOCR_V4_YAML_PROMPT, + parse_chandra_response, + parse_olmocr_response, strip_ocr_output_tag, ) __all__ = [ + "CHANDRA_OCR_LAYOUT_PROMPT", "DEFAULT_BOUNDARY_DETECTION_PROMPT", "DEFAULT_MARKDOWN_OCR_USER_PROMPT", "DEFAULT_OCR_OUTPUT_TAG", "DEFAULT_OCR_SYSTEM_PROMPT", "DEFAULT_OCR_USER_PROMPT", + "OLMOCR_V4_YAML_PROMPT", + "parse_chandra_response", + "parse_olmocr_response", "strip_ocr_output_tag", ] diff --git a/src/churro_ocr/prompts/ocr.py b/src/churro_ocr/prompts/ocr.py index b96db4e..a888d4a 100644 --- a/src/churro_ocr/prompts/ocr.py +++ b/src/churro_ocr/prompts/ocr.py @@ -2,7 +2,9 @@ from __future__ import annotations +import html import re +from typing import Any DEFAULT_OCR_OUTPUT_TAG = "output" @@ -39,6 +41,64 @@ "tables, and line breaks when they are visible." ) +CHANDRA_OCR_LAYOUT_PROMPT = ( + "OCR this image to HTML, arranged as layout blocks. Each layout block should be a div " + "with the data-bbox attribute representing the bounding box of the block in x0 y0 x1 y1 " + "format. Bboxes are normalized 0-1000. The data-label attribute is the label for the block.\n\n" + "Use the following labels:\n" + "- Caption\n" + "- Footnote\n" + "- Equation-Block\n" + "- List-Group\n" + "- Page-Header\n" + "- Page-Footer\n" + "- Image\n" + "- Section-Header\n" + "- Table\n" + "- Text\n" + "- Complex-Block\n" + "- Code-Block\n" + "- Form\n" + "- Table-Of-Contents\n" + "- Figure\n" + "- Chemical-Block\n" + "- Diagram\n" + "- Bibliography\n" + "- Blank-Page\n\n" + "Only use these tags ['math', 'br', 'i', 'b', 'u', 'del', 'sup', 'sub', 'table', 'tr', " + "'td', 'p', 'th', 'div', 'pre', 'h1', 'h2', 'h3', 'h4', 'h5', 'ul', 'ol', 'li', 'input', " + "'a', 'span', 'img', 'hr', 'tbody', 'small', 'caption', 'strong', 'thead', 'big', 'code', " + "'chem'], and these attributes ['class', 'colspan', 'rowspan', 'display', 'checked', " + "'type', 'border', 'value', 'style', 'href', 'alt', 'align', 'data-bbox', 'data-label'].\n\n" + "Guidelines:\n" + "* Inline math: Surround math with ... tags. Math expressions should be rendered " + "in KaTeX-compatible LaTeX. Use display for block math.\n" + "* Tables: Use colspan and rowspan attributes to match table structure.\n" + "* Formatting: Maintain consistent formatting with the image, including spacing, indentation, " + "subscripts/superscripts, and special characters.\n" + "* Images: Include a description of any images in the alt attribute of an tag. Do not " + "fill out the src property. Describe in detail inside the div tag. Also convert charts to high " + "fidelity data, and convert diagrams to mermaid.\n" + "* Forms: Mark checkboxes and radio buttons properly.\n" + "* Text: join lines together properly into paragraphs using

...

tags. Use
tags for " + "line breaks within paragraphs, but only when absolutely necessary to maintain meaning.\n" + "* Chemistry: Use ... tags for chemical formulas with reactive SMILES.\n" + "* Lists: Preserve indents and proper list markers.\n" + "* Use the simplest possible HTML structure that accurately represents the content of the block.\n" + "* Make sure the text is accurate and easy for a human to read and interpret. Reading order " + "should be correct and natural." +) + +OLMOCR_V4_YAML_PROMPT = ( + "Attached is one page of a document that you must process. " + "Just return the plain text representation of this document as if you were reading it naturally.\n" + "Convert equations to LateX and tables to HTML.\n" + "If there are any figures or charts, label them with the following markdown syntax " + "![Alt text describing the contents of the figure](page_startx_starty_width_height.png)\n" + "Return your output as markdown, with a front matter section on top specifying values for the " + "primary_language, is_rotation_valid, rotation_correction, is_table, and is_diagram parameters." +) + def strip_ocr_output_tag(text: str, *, output_tag: str = DEFAULT_OCR_OUTPUT_TAG) -> str: """Remove outer OCR output tags and any stray tag tokens when present. @@ -57,3 +117,108 @@ def strip_ocr_output_tag(text: str, *, output_tag: str = DEFAULT_OCR_OUTPUT_TAG) stray_tag_pattern = re.compile(rf"]*>", flags=re.IGNORECASE) return stray_tag_pattern.sub("", text).strip() + + +def _extract_yaml_front_matter(text: str) -> tuple[dict[str, object], str]: + """Return YAML front matter fields and the remaining markdown body.""" + stripped = text.strip() + if not stripped.startswith("---\n"): + return {}, stripped + + end_index = stripped.find("\n---", 4) + if end_index == -1: + return {}, stripped + + front_matter_block = stripped[4:end_index] + body = stripped[end_index + 4 :].strip() + front_matter: dict[str, object] = {} + for line in front_matter_block.splitlines(): + if ":" not in line: + continue + key, raw_value = line.split(":", 1) + key = key.strip() + value = raw_value.strip() + lower = value.lower() + if lower == "null": + parsed: object = None + elif lower == "true": + parsed = True + elif lower == "false": + parsed = False + elif re.fullmatch(r"-?\d+", value): + parsed = int(value) + else: + parsed = value + front_matter[key] = parsed + return front_matter, body + + +def _strip_rich_ocr_markup_to_plain_text(text: str) -> str: + """Best-effort plain-text conversion for OCR markdown/HTML output.""" + cleaned = text.strip() + if not cleaned: + return "" + + cleaned = re.sub(r"!\[[^\]]*]\([^)]+\)", "", cleaned) + cleaned = re.sub(r"\[([^\]]+)]\([^)]+\)", r"\1", cleaned) + + html_replacements = ( + (r"(?is)<\s*input\b[^>]*\bchecked(?:=(?:\"[^\"]*\"|'[^']*'|[^\s>]+))?[^>]*>", "[x]"), + (r"(?is)<\s*input\b[^>]*>", "[ ]"), + (r"(?i)<\s*br\s*/?\s*>", "\n"), + (r"(?i)<\s*hr\s*/?\s*>", "\n"), + (r"(?i)", "\n"), + (r"(?i)", " | "), + (r"(?i)<\s*li\b[^>]*>", ""), + ( + r"(?i)]*>", + "", + ), + (r"(?i)<\s*img\b[^>]*>", ""), + ) + for pattern, replacement in html_replacements: + cleaned = re.sub(pattern, replacement, cleaned) + + cleaned = html.unescape(cleaned) + cleaned = re.sub(r"(?is)]*>", "", cleaned) + cleaned = re.sub(r"(?m)^\s{0,3}#{1,6}\s*", "", cleaned) + cleaned = re.sub(r"(?m)^\s*[-+*]\s+", "", cleaned) + cleaned = re.sub(r"(?m)^\s*>\s?", "", cleaned) + cleaned = cleaned.replace("```", "") + cleaned = cleaned.replace("**", "") + cleaned = cleaned.replace("__", "") + cleaned = cleaned.replace("`", "") + cleaned = re.sub(r"[ \t]+\n", "\n", cleaned) + cleaned = re.sub(r"[ \t]{2,}", " ", cleaned) + + normalized_lines: list[str] = [] + saw_content = False + for raw_line in cleaned.splitlines(): + line = re.sub(r"\s*\|\s*$", "", raw_line.strip()) + line = re.sub(r"^\|\s*", "", line) + line = re.sub(r"\s*\|\s*", " | ", line) + if line: + normalized_lines.append(line) + saw_content = True + continue + if saw_content and normalized_lines and normalized_lines[-1] != "": + normalized_lines.append("") + return "\n".join(normalized_lines).strip() + + +def parse_olmocr_response(text: str) -> tuple[str, dict[str, Any]]: + """Extract plain text and metadata from an olmOCR YAML-front-matter response.""" + front_matter, markdown_body = _extract_yaml_front_matter(text) + return _strip_rich_ocr_markup_to_plain_text(markdown_body), { + "front_matter": front_matter, + "raw_markdown": markdown_body, + } + + +def parse_chandra_response(text: str) -> tuple[str, dict[str, Any]]: + """Extract plain text and metadata from a Chandra HTML-layout response.""" + raw_html = text.strip() + return _strip_rich_ocr_markup_to_plain_text(raw_html), { + "raw_html": raw_html, + } diff --git a/src/churro_ocr/providers/__init__.py b/src/churro_ocr/providers/__init__.py index ecab2bd..0a96528 100644 --- a/src/churro_ocr/providers/__init__.py +++ b/src/churro_ocr/providers/__init__.py @@ -24,7 +24,6 @@ OCRBackendSpec, OCRModelProfile, OpenAICompatibleOptions, - VLLMOptions, resolve_ocr_profile, ) @@ -53,7 +52,6 @@ "OCRModelProfile": ("churro_ocr.providers.specs", "OCRModelProfile"), "OpenAICompatibleOptions": ("churro_ocr.providers.specs", "OpenAICompatibleOptions"), "resolve_ocr_profile": ("churro_ocr.providers.specs", "resolve_ocr_profile"), - "VLLMOptions": ("churro_ocr.providers.specs", "VLLMOptions"), } __all__ = [ @@ -72,7 +70,6 @@ "OCRModelProfile", "OpenAICompatibleOptions", "resolve_ocr_profile", - "VLLMOptions", ] diff --git a/src/churro_ocr/providers/_shared.py b/src/churro_ocr/providers/_shared.py index 9d72254..42f7f01 100644 --- a/src/churro_ocr/providers/_shared.py +++ b/src/churro_ocr/providers/_shared.py @@ -77,11 +77,20 @@ def build_ocr_result( metadata: dict[str, Any] | None = None, ) -> OCRResult: """Build a normalized OCR result after postprocessing.""" + processed = text_postprocessor(text) + postprocessor_metadata: dict[str, Any] = {} + if isinstance(processed, tuple): + processed_text, postprocessor_metadata = processed + else: + processed_text = processed + + combined_metadata = dict(metadata or {}) + combined_metadata.update(postprocessor_metadata) return OCRResult( - text=text_postprocessor(text), + text=processed_text, provider_name=provider_name, model_name=model_name, - metadata=dict(metadata or {}), + metadata=combined_metadata, ) diff --git a/src/churro_ocr/providers/builder.py b/src/churro_ocr/providers/builder.py index 7f7550b..09975ed 100644 --- a/src/churro_ocr/providers/builder.py +++ b/src/churro_ocr/providers/builder.py @@ -6,8 +6,13 @@ from churro_ocr.errors import ConfigurationError from churro_ocr.ocr import OCRBackend from churro_ocr.providers.hf import ( + ChandraOCR2OCRBackend, + DeepSeekOCR2OCRBackend, + DotsMOCROCRBackend, DotsOCR15OCRBackend, HuggingFaceVisionOCRBackend, + LFM25VLOCRBackend, + PaddleOCRVL15OCRBackend, _default_dots_ocr_1_5_model_kwargs, ) from churro_ocr.providers.ocr import ( @@ -24,10 +29,9 @@ OCRBackendSpec, OCRModelProfile, OpenAICompatibleOptions, - VLLMOptions, resolve_ocr_profile, + validate_mistral_ocr_model, ) -from churro_ocr.providers.vllm import VLLMVisionOCRBackend def _merge_mapping( @@ -92,32 +96,6 @@ def _merge_huggingface_options( ) -def _merge_vllm_options( - base: VLLMOptions, - override: VLLMOptions | None, -) -> VLLMOptions: - if override is None: - return VLLMOptions( - trust_remote_code=base.trust_remote_code, - processor_kwargs=dict(base.processor_kwargs), - llm_kwargs=dict(base.llm_kwargs), - sampling_kwargs=dict(base.sampling_kwargs), - limit_mm_per_prompt=dict(base.limit_mm_per_prompt), - ) - return VLLMOptions( - trust_remote_code=( - override.trust_remote_code if override.trust_remote_code is not None else base.trust_remote_code - ), - processor_kwargs=_merge_mapping(base.processor_kwargs, override.processor_kwargs), - llm_kwargs=_merge_mapping(base.llm_kwargs, override.llm_kwargs), - sampling_kwargs=_merge_mapping(base.sampling_kwargs, override.sampling_kwargs), - limit_mm_per_prompt={ - **base.limit_mm_per_prompt, - **override.limit_mm_per_prompt, - }, - ) - - def _merge_openai_options( override: OpenAICompatibleOptions | None, ) -> OpenAICompatibleOptions: @@ -166,10 +144,8 @@ def _build_openai_compatible_backend(spec: OCRBackendSpec, profile: OCRModelProf _ensure_options_type(spec.options, OpenAICompatibleOptions, provider=spec.provider) ) transport_config = _merge_transport_config(profile.transport, spec.transport) - if not transport_config.api_base or not transport_config.api_key: - raise ConfigurationError( - "OCR provider 'openai-compatible' requires `transport.api_base` and `transport.api_key`." - ) + if not transport_config.api_base: + raise ConfigurationError("OCR provider 'openai-compatible' requires `transport.api_base`.") return OpenAICompatibleOCRBackend( model=spec.model, model_prefix=options.model_prefix or "openai", @@ -190,9 +166,19 @@ def _build_huggingface_backend(spec: OCRBackendSpec, profile: OCRModelProfile) - ) backend_cls: type[HuggingFaceVisionOCRBackend] = HuggingFaceVisionOCRBackend model_kwargs = dict(options.model_kwargs) - if options.backend_variant == "dots-ocr-1.5": + if options.backend_variant in {"dots-ocr-1.5", "dots-mocr"}: backend_cls = DotsOCR15OCRBackend + if options.backend_variant == "dots-mocr": + backend_cls = DotsMOCROCRBackend model_kwargs = _merge_mapping(_default_dots_ocr_1_5_model_kwargs(), model_kwargs) + elif options.backend_variant == "deepseek-ocr-2": + backend_cls = DeepSeekOCR2OCRBackend + elif options.backend_variant == "chandra-ocr-2": + backend_cls = ChandraOCR2OCRBackend + elif options.backend_variant == "paddleocr-vl-1.5": + backend_cls = PaddleOCRVL15OCRBackend + elif options.backend_variant == "lfm2.5-vl": + backend_cls = LFM25VLOCRBackend return backend_cls( model_id=spec.model, template=profile.template, @@ -207,27 +193,6 @@ def _build_huggingface_backend(spec: OCRBackendSpec, profile: OCRModelProfile) - ) -def _build_vllm_backend(spec: OCRBackendSpec, profile: OCRModelProfile) -> OCRBackend: - if spec.model is None: - raise ConfigurationError("OCR provider 'vllm' requires `model`.") - options = _merge_vllm_options( - profile.vllm, - _ensure_options_type(spec.options, VLLMOptions, provider=spec.provider), - ) - return VLLMVisionOCRBackend( - model_id=spec.model, - template=profile.template, - model_name=_resolve_model_name(profile, spec.model, fallback=spec.model), - trust_remote_code=bool(options.trust_remote_code), - processor_kwargs=dict(options.processor_kwargs), - llm_kwargs=dict(options.llm_kwargs), - sampling_kwargs=dict(options.sampling_kwargs), - limit_mm_per_prompt=dict(options.limit_mm_per_prompt) or {"image": 1}, - image_preprocessor=profile.image_preprocessor, - text_postprocessor=profile.text_postprocessor, - ) - - def _build_azure_backend(spec: OCRBackendSpec, profile: OCRModelProfile) -> OCRBackend: options = _ensure_options_type(spec.options, AzureDocumentIntelligenceOptions, provider=spec.provider) if options is None or not options.endpoint or not options.api_key: @@ -249,7 +214,7 @@ def _build_mistral_backend(spec: OCRBackendSpec, profile: OCRModelProfile) -> OC options = _ensure_options_type(spec.options, MistralOptions, provider=spec.provider) if options is None or not options.api_key: raise ConfigurationError("OCR provider 'mistral' requires MistralOptions(api_key=...).") - model = spec.model or "mistral-ocr-latest" + model = validate_mistral_ocr_model(spec.model) return MistralOCRBackend( api_key=options.api_key, model=model, @@ -275,8 +240,6 @@ def build_ocr_backend(spec: OCRBackendSpec) -> OCRBackend: return _build_openai_compatible_backend(spec, profile) if spec.provider == "hf": return _build_huggingface_backend(spec, profile) - if spec.provider == "vllm": - return _build_vllm_backend(spec, profile) if spec.provider == "azure": return _build_azure_backend(spec, profile) if spec.provider == "mistral": diff --git a/src/churro_ocr/providers/hf.py b/src/churro_ocr/providers/hf.py index 7832b29..567f291 100644 --- a/src/churro_ocr/providers/hf.py +++ b/src/churro_ocr/providers/hf.py @@ -5,11 +5,15 @@ import asyncio import threading from dataclasses import dataclass, field +from importlib import import_module from pathlib import Path -from typing import Any +from tempfile import TemporaryDirectory +from types import MethodType +from typing import Any, cast +from churro_ocr._internal.install import install_command_hint from churro_ocr._internal.prompt_logging import log_prompt_payload_once -from churro_ocr.errors import ConfigurationError +from churro_ocr.errors import ConfigurationError, ProviderError from churro_ocr.ocr import OCRBackend, OCRResult from churro_ocr.page_detection import DocumentPage from churro_ocr.providers._shared import ( @@ -27,12 +31,28 @@ identity_text_postprocessor, ) from churro_ocr.templates import ( + CHANDRA_OCR_2_MODEL_ID, + CHANDRA_OCR_2_OCR_TEMPLATE, CHURRO_3B_MODEL_ID, CHURRO_3B_XML_TEMPLATE, + DEEPSEEK_OCR_2_MODEL_ID, + DEEPSEEK_OCR_2_OCR_TEMPLATE, + DOTS_MOCR_MODEL_ID, + DOTS_MOCR_OCR_TEMPLATE, DOTS_OCR_1_5_MODEL_ID, DOTS_OCR_1_5_OCR_TEMPLATE, + LFM2_5_VL_1_6B_MODEL_ID, + LFM2_5_VL_1_6B_OCR_TEMPLATE, + PADDLEOCR_VL_1_5_MODEL_ID, + PADDLEOCR_VL_1_5_OCR_TEMPLATE, OCRConversation, OCRPromptTemplateLike, + build_ocr_conversation, +) + +_HF_EXTRA_INSTALL_HINT = install_command_hint("hf") +_HF_TORCH_INSTALL_HINT = ( + f"Hugging Face OCR requires a separately installed PyTorch runtime. {_HF_EXTRA_INSTALL_HINT}" ) @@ -43,13 +63,21 @@ class _HFRuntime: process_vision_info: Any +def _ensure_hf_torch_runtime() -> None: + try: + import_module("torch") + except ImportError as exc: # pragma: no cover - optional extra path + raise ConfigurationError(_HF_TORCH_INSTALL_HINT) from exc + + def _load_hf_runtime() -> _HFRuntime: + _ensure_hf_torch_runtime() try: from qwen_vl_utils import process_vision_info from transformers import AutoModelForImageTextToText, AutoProcessor except ImportError as exc: # pragma: no cover - optional extra path raise ConfigurationError( - "Hugging Face OCR requires the 'hf' extra. Install with `pip install \"churro-ocr[hf]\"`." + f"Hugging Face OCR requires the `hf` runtime. {_HF_EXTRA_INSTALL_HINT}" ) from exc return _HFRuntime( @@ -60,12 +88,13 @@ def _load_hf_runtime() -> _HFRuntime: def _load_hf_causal_runtime() -> _HFRuntime: + _ensure_hf_torch_runtime() try: from qwen_vl_utils import process_vision_info from transformers import AutoModelForCausalLM, AutoProcessor except ImportError as exc: # pragma: no cover - optional extra path raise ConfigurationError( - "Hugging Face OCR requires the 'hf' extra. Install with `pip install \"churro-ocr[hf]\"`." + f"Hugging Face OCR requires the `hf` runtime. {_HF_EXTRA_INSTALL_HINT}" ) from exc return _HFRuntime( @@ -75,6 +104,33 @@ def _load_hf_causal_runtime() -> _HFRuntime: ) +def _load_hf_auto_model_runtime() -> _HFRuntime: + _ensure_hf_torch_runtime() + try: + from transformers import AutoModel, AutoTokenizer + except ImportError as exc: # pragma: no cover - optional extra path + raise ConfigurationError( + f"Hugging Face OCR requires the `hf` runtime. {_HF_EXTRA_INSTALL_HINT}" + ) from exc + + return _HFRuntime( + processor_cls=AutoTokenizer, + model_cls=AutoModel, + process_vision_info=None, + ) + + +def _ensure_deepseek_ocr_2_cuda_runtime() -> Any: + _ensure_hf_torch_runtime() + torch = import_module("torch") + if not torch.cuda.is_available(): + raise ConfigurationError( + "DeepSeek-OCR-2 HF backend requires a CUDA-capable PyTorch runtime because " + "the upstream `infer(...)` implementation moves inputs to CUDA." + ) + return torch + + _DOTS_OCR_1_5_LOCAL_DIRNAME = "DotsOCR_1_5" _DOTS_FLASH_ATTN_IMPORT = "from flash_attn import flash_attn_varlen_func" _DOTS_FLASH_ATTN_FALLBACK = """try: @@ -129,7 +185,7 @@ def _prepare_dots_ocr_model_dir(model_id: str) -> str: from huggingface_hub import snapshot_download except ImportError as exc: # pragma: no cover - transitively provided by transformers raise ConfigurationError( - 'Hugging Face OCR requires huggingface_hub. Install with `pip install "churro-ocr[hf]"`.' + f"Hugging Face OCR requires the `hf` runtime. {_HF_EXTRA_INSTALL_HINT}" ) from exc model_dir = ( @@ -145,11 +201,69 @@ def _prepare_dots_ocr_model_dir(model_id: str) -> str: return str(model_dir) +def _patch_dots_ocr_prepare_inputs_for_generation(model: Any) -> None: + prepare_inputs_for_generation = getattr(model, "prepare_inputs_for_generation", None) + if not callable(prepare_inputs_for_generation): + return + if getattr(model, "_churro_dots_prepare_inputs_patched", False): + return + + original_prepare_inputs = getattr(prepare_inputs_for_generation, "__func__", None) + base_prepare_inputs_for_generation = prepare_inputs_for_generation + if original_prepare_inputs is not None: + for candidate in type(model).__mro__[1:]: + candidate_prepare_inputs = candidate.__dict__.get("prepare_inputs_for_generation") + if candidate_prepare_inputs is None or candidate_prepare_inputs is original_prepare_inputs: + continue + base_prepare_inputs_for_generation = cast("Any", candidate_prepare_inputs).__get__( + model, + type(model), + ) + break + if not callable(base_prepare_inputs_for_generation): + return + + def _patched_prepare_inputs_for_generation( + self: Any, + input_ids: object, + past_key_values: object = None, + inputs_embeds: object = None, + pixel_values: object = None, + attention_mask: object = None, + cache_position: object = None, + num_logits_to_keep: object = None, + **kwargs: object, + ) -> Any: + model_inputs = base_prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + num_logits_to_keep=num_logits_to_keep, + **kwargs, + ) + + first_cache_position: int | None = None + if cache_position is not None: + try: + first_cache_position = int(cast("Any", cache_position)[0]) + except Exception: + first_cache_position = None + if first_cache_position in (None, 0): + model_inputs["pixel_values"] = pixel_values + + return model_inputs + + model.prepare_inputs_for_generation = MethodType(_patched_prepare_inputs_for_generation, model) + model._churro_dots_prepare_inputs_patched = True + + def _default_dots_ocr_1_5_model_kwargs() -> dict[str, object]: model_kwargs: dict[str, object] = {"dtype": "auto"} try: - import torch - except ImportError: # pragma: no cover - torch comes from the hf extra + torch = import_module("torch") + except ImportError: # pragma: no cover - torch is installed separately for local HF use return model_kwargs if not torch.cuda.is_available(): @@ -165,6 +279,104 @@ def _default_dots_ocr_1_5_model_kwargs() -> dict[str, object]: return model_kwargs +def _default_chandra_ocr_2_model_kwargs() -> dict[str, object]: + model_kwargs: dict[str, object] = { + "device_map": "auto", + "dtype": "auto", + } + try: + torch = import_module("torch") + except ImportError: # pragma: no cover - torch is installed separately for local HF use + return model_kwargs + + if torch.cuda.is_available(): + model_kwargs["dtype"] = torch.bfloat16 + return model_kwargs + + +def _deepseek_ocr_2_prompt_from_conversation(conversation: OCRConversation) -> str: + prompt_lines: list[str] = [] + has_image = False + for message in conversation: + if message.get("role") == "system": + content_items = cast("list[dict[str, object]]", message["content"]) + system_text = "\n".join( + cast("str", item["text"]).strip() + for item in content_items + if item.get("type") == "text" and isinstance(item.get("text"), str) + ).strip() + if system_text: + raise ConfigurationError("DeepSeek-OCR-2 does not support system prompts in the HF backend.") + continue + if message.get("role") != "user": + continue + content_items = cast("list[dict[str, object]]", message["content"]) + for item in content_items: + if item.get("type") == "image": + has_image = True + continue + if item.get("type") == "text" and isinstance(item.get("text"), str): + text = cast("str", item["text"]).strip() + if text: + prompt_lines.append(text) + prompt_text = "\n".join(prompt_lines).strip() + if not prompt_text: + raise ConfigurationError("DeepSeek-OCR-2 requires a non-empty OCR prompt.") + if has_image: + return f"\n{prompt_text}" + return prompt_text + + +def _move_batch_to_model(batch: Any, model: Any) -> Any: + model_device = getattr(model, "device", None) + if hasattr(batch, "to") and model_device is not None: + batch = batch.to(model_device) + model_dtype = getattr(model, "dtype", None) + if model_dtype is not None: + batch_mapping = cast(dict[str, object], batch) + for key, value in batch_mapping.items(): + if hasattr(value, "dtype") and getattr(value.dtype, "is_floating_point", False): + batch_mapping[key] = cast(Any, value).to(dtype=model_dtype) + return batch + + +def _decode_completion_texts(processor: Any, batch: Any, generated_ids: Any) -> list[str]: + batch_mapping = cast(dict[str, object], batch) + attention_mask = batch_mapping.get("attention_mask") + if attention_mask is not None and hasattr(attention_mask, "sum"): + prompt_lengths = cast(Any, attention_mask).sum(dim=1).tolist() + completion_ids = [ + output_ids[int(prompt_length) :] + for prompt_length, output_ids in zip(prompt_lengths, generated_ids, strict=True) + ] + else: + prompt_length = cast(Any, batch_mapping["input_ids"]).shape[1] + completion_ids = generated_ids[:, prompt_length:] + return processor.batch_decode( + completion_ids, + skip_special_tokens=True, + clean_up_tokenization_spaces=False, + ) + + +def _paddleocr_vl_processor_kwargs(*, processor: Any, padding: bool) -> dict[str, object]: + processor_kwargs: dict[str, object] = { + "text_kwargs": { + "padding": padding, + "return_mm_token_type_ids": True, + } + } + image_processor = getattr(processor, "image_processor", None) + images_kwargs: dict[str, int] = {} + for key in ("min_pixels", "max_pixels"): + value = getattr(image_processor, key, None) + if isinstance(value, int): + images_kwargs[key] = value + if images_kwargs: + processor_kwargs["images_kwargs"] = images_kwargs + return processor_kwargs + + @dataclass(slots=True) class HuggingFaceVisionOCRBackend(OCRBackend): """OCR backend for local Hugging Face multimodal models with custom templates. @@ -254,23 +466,12 @@ def _ocr_sync(self, page: DocumentPage) -> OCRResult: if normalized_video_inputs is not None: batch_kwargs["videos"] = normalized_video_inputs batch = processor(**batch_kwargs) - model_device = getattr(model, "device", None) - if hasattr(batch, "to") and model_device is not None: - batch = batch.to(model_device) - model_dtype = getattr(model, "dtype", None) - if model_dtype is not None: - for key, value in batch.items(): - if hasattr(value, "dtype") and getattr(value.dtype, "is_floating_point", False): - batch[key] = value.to(dtype=model_dtype) - - generated_ids = model.generate(**batch, **self.generation_kwargs) - prompt_length = batch["input_ids"].shape[1] - completion_ids = generated_ids[:, prompt_length:] - text = processor.batch_decode( - completion_ids, - skip_special_tokens=True, - clean_up_tokenization_spaces=False, - )[0] + batch = _move_batch_to_model(batch, model) + generated_ids = model.generate( + **self._generation_inputs(batch), + **self.generation_kwargs, + ) + text = _decode_completion_texts(processor, batch, generated_ids)[0] return build_ocr_result( text, provider_name=self.provider_name, @@ -324,26 +525,12 @@ def _ocr_batch_sync(self, pages: list[DocumentPage]) -> list[OCRResult]: if has_videos: batch_kwargs["videos"] = video_batch batch = processor(**batch_kwargs) - model_device = getattr(model, "device", None) - if hasattr(batch, "to") and model_device is not None: - batch = batch.to(model_device) - model_dtype = getattr(model, "dtype", None) - if model_dtype is not None: - for key, value in batch.items(): - if hasattr(value, "dtype") and getattr(value.dtype, "is_floating_point", False): - batch[key] = value.to(dtype=model_dtype) - - generated_ids = model.generate(**batch, **self.generation_kwargs) - prompt_lengths = batch["attention_mask"].sum(dim=1).tolist() - completion_ids = [ - output_ids[int(prompt_length) :] - for prompt_length, output_ids in zip(prompt_lengths, generated_ids, strict=True) - ] - texts = processor.batch_decode( - completion_ids, - skip_special_tokens=True, - clean_up_tokenization_spaces=False, + batch = _move_batch_to_model(batch, model) + generated_ids = model.generate( + **self._generation_inputs(batch), + **self.generation_kwargs, ) + texts = _decode_completion_texts(processor, batch, generated_ids) return [ build_ocr_result( text, @@ -357,6 +544,9 @@ def _ocr_batch_sync(self, pages: list[DocumentPage]) -> list[OCRResult]: def _load_runtime(self) -> _HFRuntime: return _load_hf_runtime() + def _generation_inputs(self, batch: Any) -> dict[str, object]: + return dict(cast("dict[str, object]", batch)) + def _resolve_model_source(self) -> str: return self.model_id @@ -426,6 +616,152 @@ def _log_prompt_payload( ) +@dataclass(slots=True) +class ChandraOCR2OCRBackend(HuggingFaceVisionOCRBackend): + """Preset OCR backend for ``datalab-to/chandra-ocr-2``.""" + + model_id: str = CHANDRA_OCR_2_MODEL_ID + template: OCRPromptTemplateLike = CHANDRA_OCR_2_OCR_TEMPLATE + model_name: str | None = "chandra-ocr-2" + + def _get_processor(self, runtime: _HFRuntime) -> Any: + processor = super()._get_processor(runtime) + tokenizer = getattr(processor, "tokenizer", None) + if tokenizer is not None and getattr(tokenizer, "padding_side", None) != "left": + tokenizer.padding_side = "left" + return processor + + def _get_model(self, runtime: _HFRuntime) -> Any: + if self._model is None: + with self._init_lock: + if self._model is None: + model_kwargs = { + **_default_chandra_ocr_2_model_kwargs(), + **self.model_kwargs, + } + self._model = runtime.model_cls.from_pretrained( + self._get_model_source(), + trust_remote_code=self.trust_remote_code, + **model_kwargs, + ) + eval_method = getattr(self._model, "eval", None) + if callable(eval_method): + eval_method() + return self._model + + def _build_chandra_batch(self, processor: Any, conversations: list[OCRConversation]) -> Any: + processor_apply = getattr(processor, "apply_chat_template", None) + if not callable(processor_apply): + raise ConfigurationError("Chandra OCR 2 requires `processor.apply_chat_template(...)` support.") + return processor_apply( + conversations, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + padding=True, + ) + + def _resolve_chandra_generation_kwargs(self, processor: Any, model: Any) -> dict[str, object]: + generation_kwargs = dict(self.generation_kwargs) + eos_token_ids: list[int] = [] + eos_token_id = getattr(getattr(model, "generation_config", None), "eos_token_id", None) + if isinstance(eos_token_id, int): + eos_token_ids.append(eos_token_id) + elif isinstance(eos_token_id, list): + eos_token_ids.extend(token_id for token_id in eos_token_id if isinstance(token_id, int)) + + tokenizer = getattr(processor, "tokenizer", None) + convert_tokens_to_ids = getattr(tokenizer, "convert_tokens_to_ids", None) + if callable(convert_tokens_to_ids): + im_end_id = convert_tokens_to_ids("<|im_end|>") + if isinstance(im_end_id, int) and im_end_id >= 0 and im_end_id not in eos_token_ids: + eos_token_ids.append(im_end_id) + + if eos_token_ids: + generation_kwargs["eos_token_id"] = eos_token_ids + return generation_kwargs + + def _ocr_sync(self, page: DocumentPage) -> OCRResult: + prepared_page = preprocess_backend_page( + page, + image_preprocessor=self.image_preprocessor, + ) + runtime = self._load_runtime() + processor = self._get_processor(runtime) + model = self._get_model(runtime) + + rendered, conversation = render_ocr_prompt( + processor, + self.template, + prepared_page, + add_generation_prompt=True, + ) + self._log_prompt_payload( + rendered_prompt=rendered, + conversation=conversation, + batch_size=1, + ) + batch = self._build_chandra_batch(processor, [conversation]) + batch = _move_batch_to_model(batch, model) + generated_ids = model.generate( + **batch, + **self._resolve_chandra_generation_kwargs(processor, model), + ) + text = _decode_completion_texts(processor, batch, generated_ids)[0] + return build_ocr_result( + text, + provider_name=self.provider_name, + model_name=self.model_name or self.model_id, + text_postprocessor=self.text_postprocessor, + ) + + def _ocr_batch_sync(self, pages: list[DocumentPage]) -> list[OCRResult]: + if not pages: + return [] + + runtime = self._load_runtime() + processor = self._get_processor(runtime) + model = self._get_model(runtime) + conversations: list[OCRConversation] = [] + + for page in pages: + prepared_page = preprocess_backend_page( + page, + image_preprocessor=self.image_preprocessor, + ) + rendered, conversation = render_ocr_prompt( + processor, + self.template, + prepared_page, + add_generation_prompt=True, + ) + conversations.append(conversation) + if not self._has_logged_prompt: + self._log_prompt_payload( + rendered_prompt=rendered, + conversation=conversation, + batch_size=len(pages), + ) + + batch = self._build_chandra_batch(processor, conversations) + batch = _move_batch_to_model(batch, model) + generated_ids = model.generate( + **batch, + **self._resolve_chandra_generation_kwargs(processor, model), + ) + texts = _decode_completion_texts(processor, batch, generated_ids) + return [ + build_ocr_result( + text, + provider_name=self.provider_name, + model_name=self.model_name or self.model_id, + text_postprocessor=self.text_postprocessor, + ) + for text in texts + ] + + @dataclass(slots=True) class Churro3BOCRBackend(HuggingFaceVisionOCRBackend): """Preset OCR backend for ``stanford-oval/churro-3B``.""" @@ -435,6 +771,132 @@ class Churro3BOCRBackend(HuggingFaceVisionOCRBackend): model_name: str | None = "churro-3B" +@dataclass(slots=True) +class DeepSeekOCR2OCRBackend(HuggingFaceVisionOCRBackend): + """Preset OCR backend for ``deepseek-ai/DeepSeek-OCR-2``.""" + + model_id: str = DEEPSEEK_OCR_2_MODEL_ID + template: OCRPromptTemplateLike = DEEPSEEK_OCR_2_OCR_TEMPLATE + model_name: str | None = "DeepSeek-OCR-2" + trust_remote_code: bool = True + model_kwargs: dict[str, object] = field(default_factory=lambda: {"use_safetensors": True}) + generation_kwargs: dict[str, object] = field(default_factory=lambda: {"max_new_tokens": 8_192}) + base_size: int = 1_024 + image_size: int = 768 + crop_mode: bool = True + + def _load_runtime(self) -> _HFRuntime: + return _load_hf_auto_model_runtime() + + def _get_model(self, runtime: _HFRuntime) -> Any: + torch = _ensure_deepseek_ocr_2_cuda_runtime() + if self._model is None: + with self._init_lock: + if self._model is None: + model_kwargs = dict(self.model_kwargs) + model = runtime.model_cls.from_pretrained( + self._get_model_source(), + trust_remote_code=self.trust_remote_code, + **model_kwargs, + ) + eval_method = getattr(model, "eval", None) + if callable(eval_method): + model = eval_method() + cuda_method = getattr(model, "cuda", None) + if callable(cuda_method) and "device_map" not in model_kwargs: + model = cuda_method() + to_method = getattr(model, "to", None) + if ( + callable(to_method) + and "torch_dtype" not in model_kwargs + and "dtype" not in model_kwargs + ): + model = to_method(torch.bfloat16) + self._model = model + return self._model + + def _infer_deepseek_page( + self, + *, + tokenizer: Any, + model: Any, + page: DocumentPage, + batch_size: int, + ) -> OCRResult: + conversation = build_ocr_conversation(self.template, page) + rendered_prompt = _deepseek_ocr_2_prompt_from_conversation(conversation) + self._log_prompt_payload( + rendered_prompt=rendered_prompt, + conversation=conversation, + batch_size=batch_size, + ) + + infer_method = getattr(model, "infer", None) + if not callable(infer_method): + raise ConfigurationError("DeepSeek-OCR-2 requires a model object with `infer(...)` support.") + + with TemporaryDirectory(prefix="churro-deepseek-ocr-2-") as output_dir: + image_path = Path(output_dir) / "page.png" + page.image.save(image_path) + text = infer_method( + tokenizer, + prompt=rendered_prompt, + image_file=str(image_path), + output_path=output_dir, + base_size=self.base_size, + image_size=self.image_size, + crop_mode=self.crop_mode, + save_results=False, + eval_mode=True, + ) + if not isinstance(text, str): + raise ProviderError("DeepSeek-OCR-2 returned no OCR text.") + return build_ocr_result( + text, + provider_name=self.provider_name, + model_name=self.model_name or self.model_id, + text_postprocessor=self.text_postprocessor, + ) + + def _ocr_sync(self, page: DocumentPage) -> OCRResult: + prepared_page = preprocess_backend_page( + page, + image_preprocessor=self.image_preprocessor, + ) + runtime = self._load_runtime() + tokenizer = self._get_processor(runtime) + model = self._get_model(runtime) + return self._infer_deepseek_page( + tokenizer=tokenizer, + model=model, + page=prepared_page, + batch_size=1, + ) + + def _ocr_batch_sync(self, pages: list[DocumentPage]) -> list[OCRResult]: + if not pages: + return [] + + runtime = self._load_runtime() + tokenizer = self._get_processor(runtime) + model = self._get_model(runtime) + results: list[OCRResult] = [] + for page in pages: + prepared_page = preprocess_backend_page( + page, + image_preprocessor=self.image_preprocessor, + ) + results.append( + self._infer_deepseek_page( + tokenizer=tokenizer, + model=model, + page=prepared_page, + batch_size=len(pages), + ) + ) + return results + + @dataclass(slots=True) class DotsOCR15OCRBackend(HuggingFaceVisionOCRBackend): """Preset OCR backend for ``kristaller486/dots.ocr-1.5``. @@ -452,6 +914,11 @@ class DotsOCR15OCRBackend(HuggingFaceVisionOCRBackend): def _load_runtime(self) -> _HFRuntime: return _load_hf_causal_runtime() + def _generation_inputs(self, batch: Any) -> dict[str, object]: + generation_inputs = super()._generation_inputs(batch) + generation_inputs.pop("mm_token_type_ids", None) + return generation_inputs + def _resolve_model_source(self) -> str: return _prepare_dots_ocr_model_dir(self.model_id) @@ -477,11 +944,280 @@ def _get_model(self, runtime: _HFRuntime) -> Any: trust_remote_code=self.trust_remote_code, **self.model_kwargs, ) + _patch_dots_ocr_prepare_inputs_for_generation(self._model) return self._model +@dataclass(slots=True) +class DotsMOCROCRBackend(DotsOCR15OCRBackend): + """Preset OCR backend for ``rednote-hilab/dots.mocr``.""" + + model_id: str = DOTS_MOCR_MODEL_ID + template: OCRPromptTemplateLike = DOTS_MOCR_OCR_TEMPLATE + model_name: str | None = "dots.mocr" + + +@dataclass(slots=True) +class PaddleOCRVL15OCRBackend(HuggingFaceVisionOCRBackend): + """Preset OCR backend for ``PaddlePaddle/PaddleOCR-VL-1.5``.""" + + model_id: str = PADDLEOCR_VL_1_5_MODEL_ID + template: OCRPromptTemplateLike = PADDLEOCR_VL_1_5_OCR_TEMPLATE + model_name: str | None = "PaddleOCR-VL-1.5" + generation_kwargs: dict[str, object] = field( + default_factory=lambda: {"max_new_tokens": 4_096, "do_sample": False} + ) + + def _get_processor(self, runtime: _HFRuntime) -> Any: + processor = super()._get_processor(runtime) + tokenizer = getattr(processor, "tokenizer", None) + if tokenizer is not None and getattr(tokenizer, "padding_side", None) != "left": + tokenizer.padding_side = "left" + return processor + + def _get_model(self, runtime: _HFRuntime) -> Any: + model = super()._get_model(runtime) + eval_method = getattr(model, "eval", None) + if callable(eval_method): + eval_method() + return model + + def _build_paddleocr_vl_batch( + self, + processor: Any, + conversations: OCRConversation | list[OCRConversation], + *, + padding: bool, + ) -> Any: + processor_apply = getattr(processor, "apply_chat_template", None) + if not callable(processor_apply): + raise ConfigurationError( + "PaddleOCR-VL-1.5 requires `processor.apply_chat_template(...)` support." + ) + return processor_apply( + conversations, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + processor_kwargs=_paddleocr_vl_processor_kwargs( + processor=processor, + padding=padding, + ), + ) + + def _ocr_sync(self, page: DocumentPage) -> OCRResult: + prepared_page = preprocess_backend_page( + page, + image_preprocessor=self.image_preprocessor, + ) + runtime = self._load_runtime() + processor = self._get_processor(runtime) + model = self._get_model(runtime) + + rendered, conversation = render_ocr_prompt( + processor, + self.template, + prepared_page, + add_generation_prompt=True, + ) + self._log_prompt_payload( + rendered_prompt=rendered, + conversation=conversation, + batch_size=1, + ) + batch = self._build_paddleocr_vl_batch(processor, conversation, padding=False) + batch = _move_batch_to_model(batch, model) + generated_ids = model.generate(**batch, **self.generation_kwargs) + text = _decode_completion_texts(processor, batch, generated_ids)[0] + return build_ocr_result( + text, + provider_name=self.provider_name, + model_name=self.model_name or self.model_id, + text_postprocessor=self.text_postprocessor, + ) + + def _ocr_batch_sync(self, pages: list[DocumentPage]) -> list[OCRResult]: + if not pages: + return [] + + runtime = self._load_runtime() + processor = self._get_processor(runtime) + model = self._get_model(runtime) + conversations: list[OCRConversation] = [] + + for page in pages: + prepared_page = preprocess_backend_page( + page, + image_preprocessor=self.image_preprocessor, + ) + rendered, conversation = render_ocr_prompt( + processor, + self.template, + prepared_page, + add_generation_prompt=True, + ) + conversations.append(conversation) + if not self._has_logged_prompt: + self._log_prompt_payload( + rendered_prompt=rendered, + conversation=conversation, + batch_size=len(pages), + ) + + batch = self._build_paddleocr_vl_batch(processor, conversations, padding=True) + batch = _move_batch_to_model(batch, model) + generated_ids = model.generate(**batch, **self.generation_kwargs) + texts = _decode_completion_texts(processor, batch, generated_ids) + return [ + build_ocr_result( + text, + provider_name=self.provider_name, + model_name=self.model_name or self.model_id, + text_postprocessor=self.text_postprocessor, + ) + for text in texts + ] + + +@dataclass(slots=True) +class LFM25VLOCRBackend(HuggingFaceVisionOCRBackend): + """Preset OCR backend for ``LiquidAI/LFM2.5-VL-1.6B``.""" + + model_id: str = LFM2_5_VL_1_6B_MODEL_ID + template: OCRPromptTemplateLike = LFM2_5_VL_1_6B_OCR_TEMPLATE + model_name: str | None = "LFM2.5-VL-1.6B" + _has_tied_lm_head: bool = field(default=False, init=False, repr=False) + + def _get_processor(self, runtime: _HFRuntime) -> Any: + processor = super()._get_processor(runtime) + tokenizer = getattr(processor, "tokenizer", None) + if tokenizer is not None and getattr(tokenizer, "padding_side", None) != "left": + tokenizer.padding_side = "left" + return processor + + def _get_model(self, runtime: _HFRuntime) -> Any: + model = super()._get_model(runtime) + if self._has_tied_lm_head: + return model + + with self._init_lock: + if self._has_tied_lm_head: + return model + lm_head = getattr(model, "lm_head", None) + get_input_embeddings = getattr(model, "get_input_embeddings", None) + if lm_head is None or not callable(get_input_embeddings): + self._has_tied_lm_head = True + return model + input_embeddings = get_input_embeddings() + weight = getattr(input_embeddings, "weight", None) + if weight is not None: + lm_head.weight = weight + self._has_tied_lm_head = True + return model + + def _build_lfm_batch( + self, + processor: Any, + conversations: OCRConversation | list[OCRConversation], + *, + padding: bool, + ) -> Any: + processor_apply = getattr(processor, "apply_chat_template", None) + if not callable(processor_apply): + raise ConfigurationError( + "Liquid LFM2.5-VL requires `processor.apply_chat_template(...)` support." + ) + return processor_apply( + conversations, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + padding=padding, + ) + + def _ocr_sync(self, page: DocumentPage) -> OCRResult: + prepared_page = preprocess_backend_page( + page, + image_preprocessor=self.image_preprocessor, + ) + runtime = self._load_runtime() + processor = self._get_processor(runtime) + model = self._get_model(runtime) + + rendered, conversation = render_ocr_prompt( + processor, + self.template, + prepared_page, + add_generation_prompt=True, + ) + self._log_prompt_payload( + rendered_prompt=rendered, + conversation=conversation, + batch_size=1, + ) + batch = self._build_lfm_batch(processor, conversation, padding=False) + batch = _move_batch_to_model(batch, model) + generated_ids = model.generate(**batch, **self.generation_kwargs) + text = _decode_completion_texts(processor, batch, generated_ids)[0] + return build_ocr_result( + text, + provider_name=self.provider_name, + model_name=self.model_name or self.model_id, + text_postprocessor=self.text_postprocessor, + ) + + def _ocr_batch_sync(self, pages: list[DocumentPage]) -> list[OCRResult]: + if not pages: + return [] + + runtime = self._load_runtime() + processor = self._get_processor(runtime) + model = self._get_model(runtime) + conversations: list[OCRConversation] = [] + + for page in pages: + prepared_page = preprocess_backend_page( + page, + image_preprocessor=self.image_preprocessor, + ) + rendered, conversation = render_ocr_prompt( + processor, + self.template, + prepared_page, + add_generation_prompt=True, + ) + conversations.append(conversation) + if not self._has_logged_prompt: + self._log_prompt_payload( + rendered_prompt=rendered, + conversation=conversation, + batch_size=len(pages), + ) + + batch = self._build_lfm_batch(processor, conversations, padding=True) + batch = _move_batch_to_model(batch, model) + generated_ids = model.generate(**batch, **self.generation_kwargs) + texts = _decode_completion_texts(processor, batch, generated_ids) + return [ + build_ocr_result( + text, + provider_name=self.provider_name, + model_name=self.model_name or self.model_id, + text_postprocessor=self.text_postprocessor, + ) + for text in texts + ] + + __all__ = [ + "ChandraOCR2OCRBackend", "Churro3BOCRBackend", + "DeepSeekOCR2OCRBackend", + "DotsMOCROCRBackend", "DotsOCR15OCRBackend", "HuggingFaceVisionOCRBackend", + "LFM25VLOCRBackend", + "PaddleOCRVL15OCRBackend", ] diff --git a/src/churro_ocr/providers/ocr.py b/src/churro_ocr/providers/ocr.py index 8f58f98..452b47d 100644 --- a/src/churro_ocr/providers/ocr.py +++ b/src/churro_ocr/providers/ocr.py @@ -10,8 +10,10 @@ from typing import Any from churro_ocr._internal.image import image_to_base64 +from churro_ocr._internal.install import install_command_hint from churro_ocr._internal.litellm import LiteLLMTransport from churro_ocr._internal.prompt_logging import log_prompt_payload_once +from churro_ocr._internal.retry import retry_api_call from churro_ocr.errors import ConfigurationError, ProviderError from churro_ocr.ocr import OCRBackend, OCRResult from churro_ocr.page_detection import DocumentPage @@ -23,6 +25,7 @@ TextPostprocessor, default_ocr_image_preprocessor, identity_text_postprocessor, + validate_mistral_ocr_model, ) from churro_ocr.templates import ( DEFAULT_OCR_TEMPLATE, @@ -30,6 +33,8 @@ build_ocr_conversation, ) +_MISTRAL_REQUEST_TIMEOUT_SECONDS = 60.0 + def _with_default_ocr_completion_kwargs(config: LiteLLMTransportConfig) -> LiteLLMTransportConfig: completion_kwargs: dict[str, object] = {"max_tokens": DEFAULT_OCR_MAX_TOKENS} @@ -100,6 +105,7 @@ async def ocr(self, page: DocumentPage) -> OCRResult: text = await self.transport.complete_text( model=self.model, messages=messages, + allow_empty=True, ) return build_ocr_result( text, @@ -183,7 +189,7 @@ async def _get_client(self) -> Any: from azure.core.credentials import AzureKeyCredential except ImportError as exc: # pragma: no cover - optional extra path raise ConfigurationError( - "Azure OCR requires the 'azure' extra. Install with `pip install \"churro-ocr[azure]\"`." + f"Azure OCR requires the `azure` runtime. {install_command_hint('azure')}" ) from exc client = DocumentIntelligenceClient( @@ -219,12 +225,20 @@ async def ocr(self, page: DocumentPage) -> OCRResult: set_logged=lambda: setattr(self, "_has_logged_prompt", True), ) client = await self._get_client() - poller = await client.begin_analyze_document( - model_id=self.model_id, - body=BytesIO(image_bytes), - content_type="application/octet-stream", + + async def _analyze_document() -> Any: + poller = await client.begin_analyze_document( + model_id=self.model_id, + body=BytesIO(image_bytes), + content_type="application/octet-stream", + ) + return await poller.result() + + result = await retry_api_call( + _analyze_document, + operation_name="Azure OCR request", + context=f"for model {self.model_id}", ) - result = await poller.result() if not isinstance(result.content, str): raise ProviderError("Azure Document Intelligence returned no OCR text.") return build_ocr_result( @@ -240,14 +254,14 @@ class MistralOCRBackend(OCRBackend): """Mistral OCR backend. :param api_key: Mistral API key used for OCR requests. - :param model: Mistral OCR model identifier. + :param model: Pinned Mistral OCR model identifier. :param model_name: Optional human-readable model name for result metadata. :param image_preprocessor: Image preprocessor applied before OCR. :param text_postprocessor: Text postprocessor applied after OCR. """ api_key: str - model: str = "mistral-ocr-latest" + model: str model_name: str | None = None image_preprocessor: ImagePreprocessor = default_ocr_image_preprocessor text_postprocessor: TextPostprocessor = identity_text_postprocessor @@ -256,6 +270,10 @@ class MistralOCRBackend(OCRBackend): _has_logged_prompt: bool = field(default=False, init=False, repr=False) _prompt_log_lock: Lock = field(default_factory=Lock, init=False, repr=False) + def __post_init__(self) -> None: + """Reject unsupported Mistral OCR aliases and unpinned model ids.""" + validate_mistral_ocr_model(self.model, context="Mistral OCR backend") + async def _get_client(self) -> Any: client = self._client if client is not None: @@ -269,8 +287,7 @@ async def _get_client(self) -> Any: from mistralai import Mistral except ImportError as exc: # pragma: no cover - optional extra path raise ConfigurationError( - "Mistral OCR requires the 'mistral' extra. " - 'Install with `pip install "churro-ocr[mistral]"`.' + f"Mistral OCR requires the `mistral` runtime. {install_command_hint('mistral')}" ) from exc client = Mistral(api_key=self.api_key) @@ -302,9 +319,21 @@ async def ocr(self, page: DocumentPage) -> OCRResult: set_logged=lambda: setattr(self, "_has_logged_prompt", True), ) client = await self._get_client() - response = await client.ocr.process_async( - model=self.model, - document={"type": "image_url", "image_url": image_url}, + document = {"type": "image_url", "image_url": image_url} + + async def _process_ocr() -> Any: + return await asyncio.wait_for( + client.ocr.process_async( + model=self.model, + document=document, + ), + timeout=_MISTRAL_REQUEST_TIMEOUT_SECONDS, + ) + + response = await retry_api_call( + _process_ocr, + operation_name="Mistral OCR request", + context=f"for model {self.model}", ) if not response.pages: raise ProviderError("Mistral OCR returned no pages.") diff --git a/src/churro_ocr/providers/page_detection.py b/src/churro_ocr/providers/page_detection.py index b967e4f..3ff6362 100644 --- a/src/churro_ocr/providers/page_detection.py +++ b/src/churro_ocr/providers/page_detection.py @@ -10,8 +10,10 @@ from PIL import Image, ImageDraw, ImageOps +from churro_ocr._internal.install import install_command_hint from churro_ocr._internal.litellm import LiteLLMTransport from churro_ocr._internal.logging import logger +from churro_ocr._internal.retry import retry_api_call from churro_ocr._internal.runtime import run_sync from churro_ocr.errors import ConfigurationError, ProviderError from churro_ocr.page_detection import PageCandidate, PageDetectionBackend @@ -1202,8 +1204,7 @@ async def detect(self, image: Image.Image) -> list[PageCandidate]: from azure.core.credentials import AzureKeyCredential except ImportError as exc: # pragma: no cover - optional extra path raise ConfigurationError( - "Azure page detection requires the 'azure' extra. " - 'Install with `pip install "churro-ocr[azure]"`.' + f"Azure page detection requires the `azure` runtime. {install_command_hint('azure')}" ) from exc buffer = BytesIO() @@ -1213,12 +1214,21 @@ async def detect(self, image: Image.Image) -> list[PageCandidate]: credential=AzureKeyCredential(self.api_key), ) try: - poller = await client.begin_analyze_document( - model_id=self.model_id, - body=BytesIO(buffer.getvalue()), - content_type="application/octet-stream", + image_bytes = buffer.getvalue() + + async def _analyze_document() -> Any: + poller = await client.begin_analyze_document( + model_id=self.model_id, + body=BytesIO(image_bytes), + content_type="application/octet-stream", + ) + return await poller.result() + + result = await retry_api_call( + _analyze_document, + operation_name="Azure page detection request", + context=f"for model {self.model_id}", ) - result = await poller.result() finally: await client.close() diff --git a/src/churro_ocr/providers/specs.py b/src/churro_ocr/providers/specs.py index 2ca709d..cc6fec8 100644 --- a/src/churro_ocr/providers/specs.py +++ b/src/churro_ocr/providers/specs.py @@ -2,21 +2,43 @@ from __future__ import annotations -from collections.abc import Callable +import re +from collections.abc import Callable, Sequence from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal, cast from PIL import Image -from churro_ocr._internal.image import prepare_ocr_image -from churro_ocr.prompts import DEFAULT_OCR_OUTPUT_TAG, strip_ocr_output_tag +from churro_ocr._internal.image import ensure_rgb, prepare_ocr_image, resize_image_to_fit +from churro_ocr.errors import ConfigurationError +from churro_ocr.prompts import ( + DEFAULT_OCR_OUTPUT_TAG, + parse_chandra_response, + parse_olmocr_response, + strip_ocr_output_tag, +) from churro_ocr.templates import ( + CHANDRA_OCR_2_MODEL_ID, + CHANDRA_OCR_2_OCR_TEMPLATE, CHURRO_3B_MODEL_ID, CHURRO_3B_XML_TEMPLATE, + DEEPSEEK_OCR_2_MODEL_ID, + DEEPSEEK_OCR_2_OCR_PROMPT, + DEEPSEEK_OCR_2_OCR_TEMPLATE, DEFAULT_OCR_TEMPLATE, + DOTS_MOCR_MODEL_ID, + DOTS_MOCR_OCR_TEMPLATE, DOTS_OCR_1_5_MODEL_ID, DOTS_OCR_1_5_OCR_TEMPLATE, + LFM2_5_VL_1_6B_MODEL_ID, + LFM2_5_VL_1_6B_OCR_TEMPLATE, + OLMOCR_2_7B_1025_FP8_MODEL_ID, + OLMOCR_2_7B_1025_MODEL_ID, + OLMOCR_2_7B_1025_OCR_TEMPLATE, + PADDLEOCR_VL_1_5_MODEL_ID, + PADDLEOCR_VL_1_5_OCR_PROMPT, + PADDLEOCR_VL_1_5_OCR_TEMPLATE, OCRConversation, OCRPromptTemplateLike, ) @@ -25,11 +47,39 @@ pass -OCRProvider = Literal["litellm", "openai-compatible", "azure", "mistral", "hf", "vllm"] +OCRProvider = Literal["litellm", "openai-compatible", "azure", "mistral", "hf"] +MistralOCRModel = Literal["mistral-ocr-2505", "mistral-ocr-2512"] ImagePreprocessor = Callable[[Image.Image], Image.Image] -TextPostprocessor = Callable[[str], str] +TextPostprocessorResult = str | tuple[str, dict[str, Any]] +TextPostprocessor = Callable[[str], TextPostprocessorResult] VisionInputBuilder = Callable[[OCRConversation], object] DEFAULT_OCR_MAX_TOKENS = 20_000 +CHANDRA_OCR_MAX_TOKENS = 12_384 +DEEPSEEK_OCR_2_MAX_TOKENS = 8_192 +OLMOCR_MAX_TOKENS = 8_000 +PADDLEOCR_VL_MAX_TOKENS = 4_096 +CHANDRA_MAX_IMAGE_SIZE = (3_072, 2_048) +CHANDRA_MIN_IMAGE_SIZE = (1_792, 28) +CHANDRA_IMAGE_GRID_SIZE = 28 +OLMOCR_TARGET_LONGEST_IMAGE_DIM = 1_288 +MISTRAL_OCR_MODEL_IDS: tuple[MistralOCRModel, ...] = ( + "mistral-ocr-2505", + "mistral-ocr-2512", +) + + +def validate_mistral_ocr_model( + model: str | None, + *, + context: str = "OCR provider 'mistral'", +) -> MistralOCRModel: + """Return a supported pinned Mistral OCR model id or raise a configuration error.""" + supported_models = ", ".join(MISTRAL_OCR_MODEL_IDS) + if model is None: + raise ConfigurationError(f"{context} requires `model` to be one of: {supported_models}.") + if model not in MISTRAL_OCR_MODEL_IDS: + raise ConfigurationError(f"{context} only supports `model` values {supported_models}; got {model!r}.") + return cast(MistralOCRModel, model) def identity_text_postprocessor(text: str) -> str: @@ -59,6 +109,159 @@ def default_ocr_text_postprocessor(text: str) -> str: return strip_ocr_output_tag(text, output_tag=DEFAULT_OCR_OUTPUT_TAG) +_CHAT_ROLE_PREFIXES = { + "assistant", + "assistant:", + "user", + "user:", + "system", + "system:", + "", + "", + "", + "<|assistant|>", + "<|assistant|>:", + "<|user|>", + "<|user|>:", + "<|system|>", + "<|system|>:", + "<｜assistant｜>", + "<｜assistant｜>:", + "<｜user｜>", + "<｜user｜>:", + "<｜system｜>", + "<｜system｜>:", +} + + +def _strip_leading_chat_scaffold(text: str, *, prompts: Sequence[str]) -> str: + """Remove echoed prompts and leading chat role markers from model output.""" + cleaned = text.strip() + if not cleaned: + return "" + + normalized_prompts = tuple(prompt.strip() for prompt in prompts if prompt and prompt.strip()) + for _ in range(8): + previous = cleaned + lowered = cleaned.casefold() + stripped_prompt = False + for prompt in normalized_prompts: + if lowered.startswith(prompt.casefold()): + cleaned = cleaned[len(prompt) :].lstrip() + stripped_prompt = True + break + if stripped_prompt: + continue + + lines = cleaned.splitlines() + if not lines: + return "" + first_line = lines[0].strip() + if first_line.casefold() in _CHAT_ROLE_PREFIXES: + cleaned = "\n".join(lines[1:]).lstrip() + continue + if re.fullmatch(r"<\|?(?:assistant|user|system)\|?>", first_line, flags=re.IGNORECASE): + cleaned = "\n".join(lines[1:]).lstrip() + continue + if cleaned == previous: + break + return cleaned.strip() + + +def olmocr_image_preprocessor(image: Image.Image) -> Image.Image: + """Resize an image to olmOCR's expected 1288px longest side and normalize to RGB.""" + return ensure_rgb( + resize_image_to_fit( + image, + OLMOCR_TARGET_LONGEST_IMAGE_DIM, + OLMOCR_TARGET_LONGEST_IMAGE_DIM, + ) + ) + + +def olmocr_text_postprocessor(text: str) -> TextPostprocessorResult: + """Extract plain text and metadata from olmOCR YAML/markdown output.""" + return parse_olmocr_response(text) + + +def lfm2_5_vl_text_postprocessor(text: str) -> str: + """Strip Liquid LFM2.5-VL chat scaffold and OCR wrapper tags.""" + prompt = getattr(LFM2_5_VL_1_6B_OCR_TEMPLATE, "user_prompt", None) + cleaned = _strip_leading_chat_scaffold(text, prompts=[prompt] if isinstance(prompt, str) else []) + return strip_ocr_output_tag(cleaned, output_tag=DEFAULT_OCR_OUTPUT_TAG) + + +def deepseek_ocr_2_text_postprocessor(text: str) -> str: + """Strip DeepSeek OCR 2 prompt echoes, chat scaffold, and trailing stop tokens.""" + cleaned = text.strip() + stop_token = "<｜end▁of▁sentence｜>" + while cleaned.endswith(stop_token): + cleaned = cleaned[: -len(stop_token)].rstrip() + cleaned = _strip_leading_chat_scaffold( + cleaned, + prompts=[ + f"\n{DEEPSEEK_OCR_2_OCR_PROMPT}", + DEEPSEEK_OCR_2_OCR_PROMPT, + ], + ) + return cleaned.strip() + + +def paddleocr_vl_text_postprocessor(text: str) -> str: + """Strip PaddleOCR-VL prompt echoes and leading chat scaffold from OCR output.""" + return _strip_leading_chat_scaffold(text, prompts=[PADDLEOCR_VL_1_5_OCR_PROMPT]) + + +def chandra_image_preprocessor(image: Image.Image) -> Image.Image: + """Resize an image using Chandra OCR 2's pixel-budget and 28px-grid scaling.""" + width, height = image.size + if width <= 0 or height <= 0: + return ensure_rgb(image) + + max_pixels = CHANDRA_MAX_IMAGE_SIZE[0] * CHANDRA_MAX_IMAGE_SIZE[1] + min_pixels = CHANDRA_MIN_IMAGE_SIZE[0] * CHANDRA_MIN_IMAGE_SIZE[1] + current_pixels = width * height + scale = 1.0 + if current_pixels > max_pixels: + scale = (max_pixels / current_pixels) ** 0.5 + elif current_pixels < min_pixels: + scale = (min_pixels / current_pixels) ** 0.5 + + original_aspect_ratio = width / height + width_blocks = max(1, round((width * scale) / CHANDRA_IMAGE_GRID_SIZE)) + height_blocks = max(1, round((height * scale) / CHANDRA_IMAGE_GRID_SIZE)) + + while (width_blocks * height_blocks * CHANDRA_IMAGE_GRID_SIZE**2) > max_pixels: + if width_blocks == 1 and height_blocks == 1: + break + if width_blocks == 1: + height_blocks -= 1 + continue + if height_blocks == 1: + width_blocks -= 1 + continue + + width_loss = abs(((width_blocks - 1) / height_blocks) - original_aspect_ratio) + height_loss = abs((width_blocks / (height_blocks - 1)) - original_aspect_ratio) + if width_loss < height_loss: + width_blocks -= 1 + else: + height_blocks -= 1 + + new_size = ( + width_blocks * CHANDRA_IMAGE_GRID_SIZE, + height_blocks * CHANDRA_IMAGE_GRID_SIZE, + ) + if new_size == (width, height): + return ensure_rgb(image) + return ensure_rgb(image.resize(new_size, resample=Image.Resampling.LANCZOS)) + + +def chandra_text_postprocessor(text: str) -> TextPostprocessorResult: + """Extract plain text and metadata from Chandra OCR 2 HTML-layout output.""" + return parse_chandra_response(text) + + @dataclass(slots=True, frozen=True) class LiteLLMTransportConfig: """Shared transport config for LiteLLM-based multimodal requests. @@ -109,24 +312,6 @@ class HuggingFaceOptions: backend_variant: str | None = None -@dataclass(slots=True, frozen=True) -class VLLMOptions: - """Provider options for local vLLM OCR backends. - - :param trust_remote_code: Whether to allow remote model code execution. - :param processor_kwargs: Extra kwargs passed to ``AutoProcessor.from_pretrained``. - :param llm_kwargs: Extra kwargs passed to the vLLM ``LLM`` constructor. - :param sampling_kwargs: Extra kwargs passed to vLLM sampling params. - :param limit_mm_per_prompt: Per-request multimodal limits passed to vLLM. - """ - - trust_remote_code: bool | None = None - processor_kwargs: dict[str, object] = field(default_factory=dict) - llm_kwargs: dict[str, object] = field(default_factory=dict) - sampling_kwargs: dict[str, object] = field(default_factory=dict) - limit_mm_per_prompt: dict[str, int] = field(default_factory=dict) - - @dataclass(slots=True, frozen=True) class AzureDocumentIntelligenceOptions: """Provider options for Azure Document Intelligence OCR. @@ -150,11 +335,7 @@ class MistralOptions: OCRProviderOptions = ( - OpenAICompatibleOptions - | HuggingFaceOptions - | VLLMOptions - | AzureDocumentIntelligenceOptions - | MistralOptions + OpenAICompatibleOptions | HuggingFaceOptions | AzureDocumentIntelligenceOptions | MistralOptions ) @@ -169,7 +350,6 @@ class OCRModelProfile: :param display_name: Optional human-readable model name. :param transport: Default LiteLLM transport settings for this profile. :param huggingface: Default Hugging Face backend options for this profile. - :param vllm: Default vLLM backend options for this profile. """ profile_name: str @@ -179,7 +359,6 @@ class OCRModelProfile: display_name: str | None = None transport: LiteLLMTransportConfig = field(default_factory=LiteLLMTransportConfig) huggingface: HuggingFaceOptions = field(default_factory=HuggingFaceOptions) - vllm: VLLMOptions = field(default_factory=VLLMOptions) @dataclass(slots=True, frozen=True) @@ -221,6 +400,56 @@ def churro_3b_profile() -> OCRModelProfile: ) +def chandra_ocr_2_profile() -> OCRModelProfile: + """Return the built-in ``datalab-to/chandra-ocr-2`` OCR profile.""" + return OCRModelProfile( + profile_name=CHANDRA_OCR_2_MODEL_ID, + template=CHANDRA_OCR_2_OCR_TEMPLATE, + image_preprocessor=chandra_image_preprocessor, + text_postprocessor=chandra_text_postprocessor, + display_name="chandra-ocr-2", + transport=LiteLLMTransportConfig( + completion_kwargs={ + "max_tokens": CHANDRA_OCR_MAX_TOKENS, + "temperature": 0.0, + "top_p": 0.1, + } + ), + huggingface=HuggingFaceOptions( + generation_kwargs={ + "max_new_tokens": CHANDRA_OCR_MAX_TOKENS, + }, + backend_variant="chandra-ocr-2", + ), + ) + + +def deepseek_ocr_2_profile() -> OCRModelProfile: + """Return the built-in ``deepseek-ai/DeepSeek-OCR-2`` OCR profile.""" + return OCRModelProfile( + profile_name=DEEPSEEK_OCR_2_MODEL_ID, + template=DEEPSEEK_OCR_2_OCR_TEMPLATE, + text_postprocessor=deepseek_ocr_2_text_postprocessor, + display_name="DeepSeek-OCR-2", + transport=LiteLLMTransportConfig( + completion_kwargs={ + "max_tokens": DEEPSEEK_OCR_2_MAX_TOKENS, + "temperature": 0.0, + } + ), + huggingface=HuggingFaceOptions( + model_kwargs={ + "use_safetensors": True, + }, + generation_kwargs={ + "max_new_tokens": DEEPSEEK_OCR_2_MAX_TOKENS, + }, + trust_remote_code=True, + backend_variant="deepseek-ocr-2", + ), + ) + + def dots_ocr_1_5_profile() -> OCRModelProfile: """Return the built-in ``kristaller486/dots.ocr-1.5`` OCR profile. @@ -231,24 +460,141 @@ def dots_ocr_1_5_profile() -> OCRModelProfile: template=DOTS_OCR_1_5_OCR_TEMPLATE, text_postprocessor=identity_text_postprocessor, display_name="dots.ocr-1.5", + transport=LiteLLMTransportConfig( + completion_kwargs={ + "max_tokens": 2_048, + "temperature": 0.0, + } + ), huggingface=HuggingFaceOptions( trust_remote_code=True, backend_variant="dots-ocr-1.5", ), - vllm=VLLMOptions( + ) + + +def dots_mocr_profile() -> OCRModelProfile: + """Return the built-in ``rednote-hilab/dots.mocr`` OCR profile.""" + return OCRModelProfile( + profile_name=DOTS_MOCR_MODEL_ID, + template=DOTS_MOCR_OCR_TEMPLATE, + text_postprocessor=identity_text_postprocessor, + display_name="dots.mocr", + transport=LiteLLMTransportConfig( + completion_kwargs={ + "max_tokens": DEFAULT_OCR_MAX_TOKENS, + "temperature": 0.0, + } + ), + huggingface=HuggingFaceOptions( trust_remote_code=True, + backend_variant="dots-mocr", + ), + ) + + +def paddleocr_vl_1_5_profile() -> OCRModelProfile: + """Return the built-in ``PaddlePaddle/PaddleOCR-VL-1.5`` OCR profile.""" + return OCRModelProfile( + profile_name=PADDLEOCR_VL_1_5_MODEL_ID, + template=PADDLEOCR_VL_1_5_OCR_TEMPLATE, + text_postprocessor=paddleocr_vl_text_postprocessor, + display_name="PaddleOCR-VL-1.5", + transport=LiteLLMTransportConfig( + completion_kwargs={ + "max_tokens": PADDLEOCR_VL_MAX_TOKENS, + "temperature": 0.0, + } + ), + huggingface=HuggingFaceOptions( + generation_kwargs={ + "max_new_tokens": PADDLEOCR_VL_MAX_TOKENS, + "do_sample": False, + }, + backend_variant="paddleocr-vl-1.5", ), ) +def _olmocr_profile(*, profile_name: str, display_name: str) -> OCRModelProfile: + return OCRModelProfile( + profile_name=profile_name, + template=OLMOCR_2_7B_1025_OCR_TEMPLATE, + image_preprocessor=olmocr_image_preprocessor, + text_postprocessor=olmocr_text_postprocessor, + display_name=display_name, + transport=LiteLLMTransportConfig( + completion_kwargs={ + "max_tokens": OLMOCR_MAX_TOKENS, + "temperature": 0.1, + } + ), + huggingface=HuggingFaceOptions( + generation_kwargs={ + "max_new_tokens": OLMOCR_MAX_TOKENS, + "temperature": 0.1, + "do_sample": True, + }, + ), + ) + + +def lfm2_5_vl_1_6b_profile() -> OCRModelProfile: + """Return the built-in ``LiquidAI/LFM2.5-VL-1.6B`` OCR profile.""" + return OCRModelProfile( + profile_name=LFM2_5_VL_1_6B_MODEL_ID, + template=LFM2_5_VL_1_6B_OCR_TEMPLATE, + text_postprocessor=lfm2_5_vl_text_postprocessor, + display_name="LFM2.5-VL-1.6B", + huggingface=HuggingFaceOptions( + generation_kwargs={ + "max_new_tokens": 512, + "do_sample": False, + "repetition_penalty": 1.05, + }, + backend_variant="lfm2.5-vl", + ), + ) + + +def olmocr_2_7b_1025_profile() -> OCRModelProfile: + """Return the built-in ``allenai/olmOCR-2-7B-1025`` OCR profile.""" + return _olmocr_profile( + profile_name=OLMOCR_2_7B_1025_MODEL_ID, + display_name="olmOCR-2-7B-1025", + ) + + +def olmocr_2_7b_1025_fp8_profile() -> OCRModelProfile: + """Return the built-in ``allenai/olmOCR-2-7B-1025-FP8`` OCR profile.""" + return _olmocr_profile( + profile_name=OLMOCR_2_7B_1025_FP8_MODEL_ID, + display_name="olmOCR-2-7B-1025-FP8", + ) + + def _profile_registry() -> dict[str, OCRModelProfile]: default_profile = default_ocr_profile() churro_profile = churro_3b_profile() + chandra_profile = chandra_ocr_2_profile() + deepseek_profile = deepseek_ocr_2_profile() + dots_mocr = dots_mocr_profile() dots_profile = dots_ocr_1_5_profile() + lfm2_5_vl_profile = lfm2_5_vl_1_6b_profile() + olmocr_profile = olmocr_2_7b_1025_profile() + olmocr_fp8_profile = olmocr_2_7b_1025_fp8_profile() + paddleocr_vl_profile = paddleocr_vl_1_5_profile() return { default_profile.profile_name: default_profile, churro_profile.profile_name: churro_profile, + chandra_profile.profile_name: chandra_profile, + deepseek_profile.profile_name: deepseek_profile, + dots_mocr.profile_name: dots_mocr, dots_profile.profile_name: dots_profile, + lfm2_5_vl_profile.profile_name: lfm2_5_vl_profile, + olmocr_profile.profile_name: olmocr_profile, + olmocr_fp8_profile.profile_name: olmocr_fp8_profile, + paddleocr_vl_profile.profile_name: paddleocr_vl_profile, } @@ -282,20 +628,33 @@ def resolve_ocr_profile( __all__ = [ "AzureDocumentIntelligenceOptions", "DEFAULT_OCR_MAX_TOKENS", + "chandra_image_preprocessor", + "chandra_ocr_2_profile", + "chandra_text_postprocessor", + "deepseek_ocr_2_profile", + "deepseek_ocr_2_text_postprocessor", "default_ocr_image_preprocessor", "default_ocr_profile", "default_ocr_text_postprocessor", "HuggingFaceOptions", "identity_text_postprocessor", + "lfm2_5_vl_text_postprocessor", + "lfm2_5_vl_1_6b_profile", "ImagePreprocessor", "LiteLLMTransportConfig", + "MistralOCRModel", + "MISTRAL_OCR_MODEL_IDS", "MistralOptions", + "olmocr_image_preprocessor", + "olmocr_text_postprocessor", + "paddleocr_vl_1_5_profile", + "paddleocr_vl_text_postprocessor", "OCRBackendSpec", "OCRModelProfile", "OCRProvider", "OpenAICompatibleOptions", "resolve_ocr_profile", "TextPostprocessor", + "validate_mistral_ocr_model", "VisionInputBuilder", - "VLLMOptions", ] diff --git a/src/churro_ocr/providers/vllm.py b/src/churro_ocr/providers/vllm.py deleted file mode 100644 index 1c05283..0000000 --- a/src/churro_ocr/providers/vllm.py +++ /dev/null @@ -1,209 +0,0 @@ -"""vLLM OCR backends.""" - -from __future__ import annotations - -import asyncio -import threading -from dataclasses import dataclass, field -from importlib import import_module -from typing import Any, cast - -from churro_ocr._internal.prompt_logging import log_prompt_payload_once -from churro_ocr.errors import ConfigurationError, ProviderError -from churro_ocr.ocr import OCRBackend, OCRResult -from churro_ocr.page_detection import DocumentPage -from churro_ocr.providers._shared import build_ocr_result, preprocess_backend_page, render_ocr_prompt -from churro_ocr.providers.specs import ( - DEFAULT_OCR_MAX_TOKENS, - ImagePreprocessor, - TextPostprocessor, - default_ocr_image_preprocessor, - identity_text_postprocessor, -) -from churro_ocr.templates import ( - DOTS_OCR_1_5_MODEL_ID, - DOTS_OCR_1_5_OCR_TEMPLATE, - OCRPromptTemplateLike, -) - - -def _load_vllm_processor_cls() -> Any: - try: - from transformers import AutoProcessor - except ImportError as exc: # pragma: no cover - optional extra path - raise ConfigurationError( - 'vLLM OCR requires transformers. Install with `pip install "churro-ocr[vllm]"`.' - ) from exc - - return AutoProcessor - - -def _load_vllm_runtime() -> tuple[Any, Any]: - try: - vllm = import_module("vllm") - except ImportError as exc: # pragma: no cover - optional extra path - raise ConfigurationError( - 'vLLM OCR requires the "vllm" extra. Install with `pip install "churro-ocr[vllm]"`.' - ) from exc - - vllm_any = cast(Any, vllm) - return vllm_any.LLM, vllm_any.SamplingParams - - -@dataclass(slots=True) -class VLLMVisionOCRBackend(OCRBackend): - """OCR backend for local multimodal models served by vLLM. - - :param model_id: Model identifier served by vLLM. - :param template: Prompt template used to render OCR input. - :param model_name: Optional human-readable model name for result metadata. - :param trust_remote_code: Whether to allow remote model code execution. - :param processor_kwargs: Extra kwargs passed to processor loading. - :param llm_kwargs: Extra kwargs passed to the vLLM ``LLM`` constructor. - :param sampling_kwargs: Extra sampling kwargs passed at inference time. - :param limit_mm_per_prompt: Multimodal limits passed to vLLM. - :param image_preprocessor: Image preprocessor applied before OCR. - :param text_postprocessor: Text postprocessor applied after OCR. - :param provider_name: Provider identifier written into OCR results. - """ - - model_id: str - template: OCRPromptTemplateLike - model_name: str | None = None - trust_remote_code: bool = False - processor_kwargs: dict[str, object] = field(default_factory=dict) - llm_kwargs: dict[str, object] = field(default_factory=dict) - sampling_kwargs: dict[str, object] = field(default_factory=dict) - limit_mm_per_prompt: dict[str, int] = field(default_factory=lambda: {"image": 1}) - image_preprocessor: ImagePreprocessor = default_ocr_image_preprocessor - text_postprocessor: TextPostprocessor = identity_text_postprocessor - provider_name: str = "vllm" - _processor: object | None = field(default=None, init=False, repr=False) - _llm: object | None = field(default=None, init=False, repr=False) - _init_lock: threading.RLock = field(default_factory=threading.RLock, init=False, repr=False) - _has_logged_prompt: bool = field(default=False, init=False, repr=False) - _prompt_log_lock: threading.Lock = field(default_factory=threading.Lock, init=False, repr=False) - - def __post_init__(self) -> None: - """Apply default sampling settings after dataclass initialization.""" - self.sampling_kwargs = { - "max_tokens": DEFAULT_OCR_MAX_TOKENS, - **self.sampling_kwargs, - } - - async def ocr(self, page: DocumentPage) -> OCRResult: - """Run OCR for one page. - - :param page: Page to transcribe. - :returns: Provider-agnostic OCR result. - """ - return (await self.ocr_batch([page]))[0] - - async def ocr_batch(self, pages: list[DocumentPage]) -> list[OCRResult]: - """Run OCR for multiple pages in one batch. - - :param pages: Pages to transcribe in batch order. - :returns: OCR results in the same order as ``pages``. - """ - return await asyncio.to_thread(self._ocr_batch_sync, pages) - - def _ocr_batch_sync(self, pages: list[DocumentPage]) -> list[OCRResult]: - if not pages: - return [] - - processor = self._get_processor() - llm = self._get_llm() - _, sampling_params_cls = _load_vllm_runtime() - prompts: list[dict[str, object]] = [] - - for page in pages: - prepared_page = preprocess_backend_page( - page, - image_preprocessor=self.image_preprocessor, - ) - rendered, _ = render_ocr_prompt( - processor, - self.template, - prepared_page, - add_generation_prompt=True, - ) - prompt_payload: dict[str, object] = { - "prompt": rendered, - "multi_modal_data": {"image": prepared_page.image}, - } - prompts.append(prompt_payload) - if not self._has_logged_prompt: - log_prompt_payload_once( - payload={ - "batch_size": len(pages), - "prompt": prompt_payload, - }, - provider_name=self.provider_name, - has_logged=lambda: self._has_logged_prompt, - lock=self._prompt_log_lock, - set_logged=lambda: setattr(self, "_has_logged_prompt", True), - ) - - request_outputs = llm.generate( - prompts, - sampling_params_cls(**self.sampling_kwargs), - use_tqdm=False, - ) - results: list[OCRResult] = [] - for request_output in request_outputs: - outputs = getattr(request_output, "outputs", None) - if not outputs: - raise ProviderError("vLLM OCR returned no outputs.") - text = getattr(outputs[0], "text", None) - if not isinstance(text, str): - raise ProviderError("vLLM OCR returned a non-text response.") - results.append( - build_ocr_result( - text, - provider_name=self.provider_name, - model_name=self.model_name or self.model_id, - text_postprocessor=self.text_postprocessor, - ) - ) - return results - - def _get_processor(self) -> Any: - if self._processor is None: - with self._init_lock: - if self._processor is None: - processor_cls = _load_vllm_processor_cls() - self._processor = processor_cls.from_pretrained( - self.model_id, - trust_remote_code=self.trust_remote_code, - **self.processor_kwargs, - ) - return self._processor - - def _get_llm(self) -> Any: - if self._llm is None: - with self._init_lock: - if self._llm is None: - llm_cls, _ = _load_vllm_runtime() - self._llm = llm_cls( - model=self.model_id, - trust_remote_code=self.trust_remote_code, - limit_mm_per_prompt=self.limit_mm_per_prompt, - **self.llm_kwargs, - ) - return self._llm - - -@dataclass(slots=True) -class DotsOCR15VLLMOCRBackend(VLLMVisionOCRBackend): - """Preset vLLM OCR backend for ``kristaller486/dots.ocr-1.5``.""" - - model_id: str = DOTS_OCR_1_5_MODEL_ID - template: OCRPromptTemplateLike = DOTS_OCR_1_5_OCR_TEMPLATE - model_name: str | None = "dots.ocr-1.5" - trust_remote_code: bool = True - - -__all__ = [ - "DotsOCR15VLLMOCRBackend", - "VLLMVisionOCRBackend", -] diff --git a/src/churro_ocr/templates/__init__.py b/src/churro_ocr/templates/__init__.py index 16f76b0..755da05 100644 --- a/src/churro_ocr/templates/__init__.py +++ b/src/churro_ocr/templates/__init__.py @@ -9,23 +9,55 @@ ) from churro_ocr.templates.hf import HFChatTemplate from churro_ocr.templates.presets import ( + CHANDRA_OCR_2_MODEL_ID, + CHANDRA_OCR_2_OCR_TEMPLATE, CHURRO_3B_MODEL_ID, CHURRO_3B_XML_TEMPLATE, + DEEPSEEK_OCR_2_MODEL_ID, + DEEPSEEK_OCR_2_OCR_PROMPT, + DEEPSEEK_OCR_2_OCR_TEMPLATE, DEFAULT_OCR_TEMPLATE, + DOTS_MOCR_MODEL_ID, + DOTS_MOCR_OCR_PROMPT, + DOTS_MOCR_OCR_TEMPLATE, DOTS_OCR_1_5_MODEL_ID, DOTS_OCR_1_5_OCR_PROMPT, DOTS_OCR_1_5_OCR_TEMPLATE, + LFM2_5_VL_1_6B_MODEL_ID, + LFM2_5_VL_1_6B_OCR_TEMPLATE, + OLMOCR_2_7B_1025_FP8_MODEL_ID, + OLMOCR_2_7B_1025_MODEL_ID, + OLMOCR_2_7B_1025_OCR_TEMPLATE, + PADDLEOCR_VL_1_5_MODEL_ID, + PADDLEOCR_VL_1_5_OCR_PROMPT, + PADDLEOCR_VL_1_5_OCR_TEMPLATE, ) __all__ = [ "build_ocr_conversation", "CHURRO_3B_MODEL_ID", "CHURRO_3B_XML_TEMPLATE", + "CHANDRA_OCR_2_MODEL_ID", + "CHANDRA_OCR_2_OCR_TEMPLATE", + "DEEPSEEK_OCR_2_MODEL_ID", + "DEEPSEEK_OCR_2_OCR_PROMPT", + "DEEPSEEK_OCR_2_OCR_TEMPLATE", "DEFAULT_OCR_TEMPLATE", + "DOTS_MOCR_MODEL_ID", + "DOTS_MOCR_OCR_PROMPT", + "DOTS_MOCR_OCR_TEMPLATE", "DOTS_OCR_1_5_MODEL_ID", "DOTS_OCR_1_5_OCR_PROMPT", "DOTS_OCR_1_5_OCR_TEMPLATE", + "LFM2_5_VL_1_6B_MODEL_ID", + "LFM2_5_VL_1_6B_OCR_TEMPLATE", + "PADDLEOCR_VL_1_5_MODEL_ID", + "PADDLEOCR_VL_1_5_OCR_PROMPT", + "PADDLEOCR_VL_1_5_OCR_TEMPLATE", "HFChatTemplate", + "OLMOCR_2_7B_1025_FP8_MODEL_ID", + "OLMOCR_2_7B_1025_MODEL_ID", + "OLMOCR_2_7B_1025_OCR_TEMPLATE", "OCRConversation", "OCRPromptTemplate", "OCRPromptTemplateCallable", diff --git a/src/churro_ocr/templates/hf.py b/src/churro_ocr/templates/hf.py index 16a1987..9032873 100644 --- a/src/churro_ocr/templates/hf.py +++ b/src/churro_ocr/templates/hf.py @@ -15,11 +15,13 @@ class HFChatTemplate: :param system_message: Optional system message prepended to the conversation. :param user_prompt: Optional user-side text prompt appended with the image. :param include_image: Whether to include the page image in the user message. + :param user_prompt_first: Whether to place the user prompt before the image. """ system_message: str | None = None user_prompt: str | None = None include_image: bool = True + user_prompt_first: bool = False def build_conversation(self, page: DocumentPage) -> OCRConversation: """Build a structured multimodal conversation for one OCR page. @@ -37,9 +39,11 @@ def build_conversation(self, page: DocumentPage) -> OCRConversation: ) user_content: list[dict[str, object]] = [] + if self.user_prompt and self.user_prompt_first: + user_content.append({"type": "text", "text": self.user_prompt}) if self.include_image: user_content.append({"type": "image", "image": page.image.copy()}) - if self.user_prompt: + if self.user_prompt and not self.user_prompt_first: user_content.append({"type": "text", "text": self.user_prompt}) conversation.append({"role": "user", "content": user_content}) diff --git a/src/churro_ocr/templates/presets.py b/src/churro_ocr/templates/presets.py index 9909db8..6e31226 100644 --- a/src/churro_ocr/templates/presets.py +++ b/src/churro_ocr/templates/presets.py @@ -2,11 +2,23 @@ from __future__ import annotations -from churro_ocr.prompts import DEFAULT_OCR_SYSTEM_PROMPT, DEFAULT_OCR_USER_PROMPT +from churro_ocr.prompts import ( + CHANDRA_OCR_LAYOUT_PROMPT, + DEFAULT_OCR_SYSTEM_PROMPT, + DEFAULT_OCR_USER_PROMPT, + OLMOCR_V4_YAML_PROMPT, +) from churro_ocr.templates.hf import HFChatTemplate CHURRO_3B_MODEL_ID = "stanford-oval/churro-3B" +CHANDRA_OCR_2_MODEL_ID = "datalab-to/chandra-ocr-2" +DEEPSEEK_OCR_2_MODEL_ID = "deepseek-ai/DeepSeek-OCR-2" DOTS_OCR_1_5_MODEL_ID = "kristaller486/dots.ocr-1.5" +DOTS_MOCR_MODEL_ID = "rednote-hilab/dots.mocr" +PADDLEOCR_VL_1_5_MODEL_ID = "PaddlePaddle/PaddleOCR-VL-1.5" +OLMOCR_2_7B_1025_MODEL_ID = "allenai/olmOCR-2-7B-1025" +OLMOCR_2_7B_1025_FP8_MODEL_ID = "allenai/olmOCR-2-7B-1025-FP8" +LFM2_5_VL_1_6B_MODEL_ID = "LiquidAI/LFM2.5-VL-1.6B" DEFAULT_OCR_TEMPLATE = HFChatTemplate( system_message=DEFAULT_OCR_SYSTEM_PROMPT, user_prompt=DEFAULT_OCR_USER_PROMPT, @@ -16,18 +28,62 @@ system_message="Transcribe the entirety of this historical document to XML format.", user_prompt=None, ) +CHANDRA_OCR_2_OCR_TEMPLATE = HFChatTemplate( + system_message=None, + user_prompt=CHANDRA_OCR_LAYOUT_PROMPT, +) +DEEPSEEK_OCR_2_OCR_PROMPT = "Free OCR." +DEEPSEEK_OCR_2_OCR_TEMPLATE = HFChatTemplate( + system_message=None, + user_prompt=DEEPSEEK_OCR_2_OCR_PROMPT, +) DOTS_OCR_1_5_OCR_PROMPT = "Extract the text content from this image." DOTS_OCR_1_5_OCR_TEMPLATE = HFChatTemplate( system_message=None, user_prompt=DOTS_OCR_1_5_OCR_PROMPT, ) +DOTS_MOCR_OCR_PROMPT = DOTS_OCR_1_5_OCR_PROMPT +DOTS_MOCR_OCR_TEMPLATE = HFChatTemplate( + system_message=None, + user_prompt=DOTS_MOCR_OCR_PROMPT, +) +PADDLEOCR_VL_1_5_OCR_PROMPT = "OCR:" +PADDLEOCR_VL_1_5_OCR_TEMPLATE = HFChatTemplate( + system_message=None, + user_prompt=PADDLEOCR_VL_1_5_OCR_PROMPT, +) +OLMOCR_2_7B_1025_OCR_TEMPLATE = HFChatTemplate( + system_message=None, + user_prompt=OLMOCR_V4_YAML_PROMPT, + user_prompt_first=True, +) +LFM2_5_VL_1_6B_OCR_TEMPLATE = HFChatTemplate( + system_message=None, + user_prompt="Transcribe all visible text from this historical document page in reading order.", +) __all__ = [ "CHURRO_3B_MODEL_ID", "CHURRO_3B_XML_TEMPLATE", + "CHANDRA_OCR_2_MODEL_ID", + "CHANDRA_OCR_2_OCR_TEMPLATE", + "DEEPSEEK_OCR_2_MODEL_ID", + "DEEPSEEK_OCR_2_OCR_PROMPT", + "DEEPSEEK_OCR_2_OCR_TEMPLATE", "DEFAULT_OCR_TEMPLATE", + "DOTS_MOCR_MODEL_ID", + "DOTS_MOCR_OCR_PROMPT", + "DOTS_MOCR_OCR_TEMPLATE", "DOTS_OCR_1_5_MODEL_ID", "DOTS_OCR_1_5_OCR_PROMPT", "DOTS_OCR_1_5_OCR_TEMPLATE", + "PADDLEOCR_VL_1_5_MODEL_ID", + "PADDLEOCR_VL_1_5_OCR_PROMPT", + "PADDLEOCR_VL_1_5_OCR_TEMPLATE", + "OLMOCR_2_7B_1025_FP8_MODEL_ID", + "OLMOCR_2_7B_1025_MODEL_ID", + "OLMOCR_2_7B_1025_OCR_TEMPLATE", + "LFM2_5_VL_1_6B_MODEL_ID", + "LFM2_5_VL_1_6B_OCR_TEMPLATE", ] diff --git a/tests/test_cli.py b/tests/test_cli.py index a6e6715..c564195 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4,16 +4,32 @@ import subprocess import sys from pathlib import Path +from types import SimpleNamespace import pytest +import typer from PIL import Image import churro_ocr.cli as cli_module from churro_ocr.cli import app +from churro_ocr.errors import ConfigurationError from churro_ocr.ocr import OCRResult from churro_ocr.page_detection import DocumentPage, PageDetectionResult from churro_ocr.prompts import DEFAULT_OCR_OUTPUT_TAG -from churro_ocr.templates import DEFAULT_OCR_TEMPLATE, DOTS_OCR_1_5_OCR_TEMPLATE +from churro_ocr.providers.specs import DEFAULT_OCR_MAX_TOKENS +from churro_ocr.templates import ( + CHANDRA_OCR_2_MODEL_ID, + CHANDRA_OCR_2_OCR_TEMPLATE, + DEEPSEEK_OCR_2_MODEL_ID, + DEEPSEEK_OCR_2_OCR_TEMPLATE, + DEFAULT_OCR_TEMPLATE, + DOTS_MOCR_OCR_TEMPLATE, + DOTS_OCR_1_5_OCR_TEMPLATE, + OLMOCR_2_7B_1025_MODEL_ID, + OLMOCR_2_7B_1025_OCR_TEMPLATE, + PADDLEOCR_VL_1_5_MODEL_ID, + PADDLEOCR_VL_1_5_OCR_TEMPLATE, +) def test_transcribe_cli_writes_output( @@ -111,26 +127,51 @@ def test_build_ocr_backend_aligns_templates_for_generic_models() -> None: base_url=None, api_version=None, ) - vllm_backend = cli_module._build_ocr_backend( - backend="vllm", + openai_backend = cli_module._build_ocr_backend( + backend="openai-compatible", model="example/model", endpoint=None, api_key=None, - base_url=None, + base_url="http://127.0.0.1:8000/v1", api_version=None, ) assert litellm_backend.template == DEFAULT_OCR_TEMPLATE - assert litellm_backend.template == hf_backend.template == vllm_backend.template + assert litellm_backend.template == hf_backend.template == openai_backend.template assert litellm_backend.model_name == "example/model" assert hf_backend.model_name == "example/model" - assert vllm_backend.model_name == "example/model" + assert openai_backend.model_name == "example/model" assert f"<{DEFAULT_OCR_OUTPUT_TAG}>" in litellm_backend.template.system_message assert f"" in litellm_backend.template.system_message assert f"<{DEFAULT_OCR_OUTPUT_TAG}>" in litellm_backend.template.user_prompt assert f"" in litellm_backend.template.user_prompt +def test_build_ocr_backend_requires_pinned_mistral_model() -> None: + with pytest.raises(typer.BadParameter, match="mistral-ocr-2505, mistral-ocr-2512"): + cli_module._build_ocr_backend( + backend="mistral", + model="mistral-ocr-latest", + endpoint=None, + api_key="secret", + base_url=None, + api_version=None, + ) + + +def test_build_ocr_backend_accepts_pinned_mistral_model() -> None: + mistral_backend = cli_module._build_ocr_backend( + backend="mistral", + model="mistral-ocr-2512", + endpoint=None, + api_key="secret", + base_url=None, + api_version=None, + ) + + assert mistral_backend.model == "mistral-ocr-2512" + + def test_build_ocr_backend_aligns_templates_for_dots() -> None: litellm_backend = cli_module._build_ocr_backend( backend="litellm", @@ -148,20 +189,151 @@ def test_build_ocr_backend_aligns_templates_for_dots() -> None: base_url=None, api_version=None, ) - vllm_backend = cli_module._build_ocr_backend( - backend="vllm", + openai_backend = cli_module._build_ocr_backend( + backend="openai-compatible", model="kristaller486/dots.ocr-1.5", endpoint=None, api_key=None, - base_url=None, + base_url="http://127.0.0.1:8000/v1", api_version=None, ) assert litellm_backend.template == DOTS_OCR_1_5_OCR_TEMPLATE - assert litellm_backend.template == hf_backend.template == vllm_backend.template + assert litellm_backend.template == hf_backend.template == openai_backend.template assert litellm_backend.model_name == "dots.ocr-1.5" assert hf_backend.model_name == "dots.ocr-1.5" - assert vllm_backend.model_name == "dots.ocr-1.5" + assert openai_backend.model_name == "dots.ocr-1.5" + assert litellm_backend.transport.config.completion_kwargs == { + "max_tokens": 2_048, + "temperature": 0.0, + } + assert openai_backend.transport.config.completion_kwargs == { + "max_tokens": 2_048, + "temperature": 0.0, + } + + +def test_build_ocr_backend_aligns_templates_for_dots_mocr() -> None: + litellm_backend = cli_module._build_ocr_backend( + backend="litellm", + model="rednote-hilab/dots.mocr", + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + hf_backend = cli_module._build_ocr_backend( + backend="hf", + model="rednote-hilab/dots.mocr", + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + openai_backend = cli_module._build_ocr_backend( + backend="openai-compatible", + model="rednote-hilab/dots.mocr", + endpoint=None, + api_key=None, + base_url="http://127.0.0.1:8000/v1", + api_version=None, + ) + + assert litellm_backend.template == DOTS_MOCR_OCR_TEMPLATE + assert litellm_backend.template == hf_backend.template == openai_backend.template + assert litellm_backend.model_name == "dots.mocr" + assert hf_backend.model_name == "dots.mocr" + assert openai_backend.model_name == "dots.mocr" + assert litellm_backend.transport.config.completion_kwargs == { + "max_tokens": DEFAULT_OCR_MAX_TOKENS, + "temperature": 0.0, + } + assert openai_backend.transport.config.completion_kwargs == { + "max_tokens": DEFAULT_OCR_MAX_TOKENS, + "temperature": 0.0, + } + + +def test_build_ocr_backend_aligns_templates_for_deepseek_ocr_2() -> None: + litellm_backend = cli_module._build_ocr_backend( + backend="litellm", + model=DEEPSEEK_OCR_2_MODEL_ID, + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + hf_backend = cli_module._build_ocr_backend( + backend="hf", + model=DEEPSEEK_OCR_2_MODEL_ID, + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + openai_backend = cli_module._build_ocr_backend( + backend="openai-compatible", + model=DEEPSEEK_OCR_2_MODEL_ID, + endpoint=None, + api_key=None, + base_url="http://127.0.0.1:8000/v1", + api_version=None, + ) + + assert litellm_backend.template == DEEPSEEK_OCR_2_OCR_TEMPLATE + assert litellm_backend.template == hf_backend.template == openai_backend.template + assert litellm_backend.model_name == "DeepSeek-OCR-2" + assert hf_backend.model_name == "DeepSeek-OCR-2" + assert openai_backend.model_name == "DeepSeek-OCR-2" + assert litellm_backend.transport.config.completion_kwargs == { + "max_tokens": 8_192, + "temperature": 0.0, + } + assert openai_backend.transport.config.completion_kwargs == { + "max_tokens": 8_192, + "temperature": 0.0, + } + + +def test_build_ocr_backend_aligns_templates_for_paddleocr_vl() -> None: + litellm_backend = cli_module._build_ocr_backend( + backend="litellm", + model=PADDLEOCR_VL_1_5_MODEL_ID, + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + hf_backend = cli_module._build_ocr_backend( + backend="hf", + model=PADDLEOCR_VL_1_5_MODEL_ID, + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + openai_backend = cli_module._build_ocr_backend( + backend="openai-compatible", + model=PADDLEOCR_VL_1_5_MODEL_ID, + endpoint=None, + api_key=None, + base_url="http://127.0.0.1:8000/v1", + api_version=None, + ) + + assert litellm_backend.template == PADDLEOCR_VL_1_5_OCR_TEMPLATE + assert litellm_backend.template == hf_backend.template == openai_backend.template + assert litellm_backend.model_name == "PaddleOCR-VL-1.5" + assert hf_backend.model_name == "PaddleOCR-VL-1.5" + assert openai_backend.model_name == "PaddleOCR-VL-1.5" + assert litellm_backend.transport.config.completion_kwargs == { + "max_tokens": 4_096, + "temperature": 0.0, + } + assert openai_backend.transport.config.completion_kwargs == { + "max_tokens": 4_096, + "temperature": 0.0, + } def test_build_ocr_backend_uses_generic_defaults_for_qwen_3_5_0_8b() -> None: @@ -181,21 +353,150 @@ def test_build_ocr_backend_uses_generic_defaults_for_qwen_3_5_0_8b() -> None: base_url=None, api_version=None, ) - vllm_backend = cli_module._build_ocr_backend( - backend="vllm", + openai_backend = cli_module._build_ocr_backend( + backend="openai-compatible", model="Qwen/Qwen3.5-0.8B", endpoint=None, api_key=None, - base_url=None, + base_url="http://127.0.0.1:8000/v1", api_version=None, ) assert litellm_backend.template == DEFAULT_OCR_TEMPLATE - assert litellm_backend.template == hf_backend.template == vllm_backend.template + assert litellm_backend.template == hf_backend.template == openai_backend.template assert litellm_backend.model_name == "Qwen/Qwen3.5-0.8B" assert hf_backend.model_name == "Qwen/Qwen3.5-0.8B" - assert vllm_backend.model_name == "Qwen/Qwen3.5-0.8B" - assert vllm_backend.llm_kwargs == {} + assert openai_backend.model_name == "Qwen/Qwen3.5-0.8B" + + +def test_build_ocr_backend_aligns_templates_for_olmocr() -> None: + litellm_backend = cli_module._build_ocr_backend( + backend="litellm", + model=OLMOCR_2_7B_1025_MODEL_ID, + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + hf_backend = cli_module._build_ocr_backend( + backend="hf", + model=OLMOCR_2_7B_1025_MODEL_ID, + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + openai_backend = cli_module._build_ocr_backend( + backend="openai-compatible", + model=OLMOCR_2_7B_1025_MODEL_ID, + endpoint=None, + api_key=None, + base_url="http://127.0.0.1:8000/v1", + api_version=None, + ) + + assert litellm_backend.template == OLMOCR_2_7B_1025_OCR_TEMPLATE + assert litellm_backend.template == hf_backend.template == openai_backend.template + assert litellm_backend.model_name == "olmOCR-2-7B-1025" + assert hf_backend.model_name == "olmOCR-2-7B-1025" + assert openai_backend.model_name == "olmOCR-2-7B-1025" + assert openai_backend.transport.config.completion_kwargs == { + "max_tokens": 8_000, + "temperature": 0.1, + } + + +def test_build_ocr_backend_aligns_templates_for_chandra() -> None: + litellm_backend = cli_module._build_ocr_backend( + backend="litellm", + model=CHANDRA_OCR_2_MODEL_ID, + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + hf_backend = cli_module._build_ocr_backend( + backend="hf", + model=CHANDRA_OCR_2_MODEL_ID, + endpoint=None, + api_key=None, + base_url=None, + api_version=None, + ) + openai_backend = cli_module._build_ocr_backend( + backend="openai-compatible", + model=CHANDRA_OCR_2_MODEL_ID, + endpoint=None, + api_key=None, + base_url="http://127.0.0.1:8000/v1", + api_version=None, + ) + + assert litellm_backend.template == CHANDRA_OCR_2_OCR_TEMPLATE + assert litellm_backend.template == hf_backend.template == openai_backend.template + assert litellm_backend.model_name == "chandra-ocr-2" + assert hf_backend.model_name == "chandra-ocr-2" + assert openai_backend.model_name == "chandra-ocr-2" + assert openai_backend.transport.config.completion_kwargs == { + "max_tokens": 12_384, + "temperature": 0.0, + "top_p": 0.1, + } + + +def test_install_command_invokes_runtime_installer( + monkeypatch: pytest.MonkeyPatch, + cli_runner, +) -> None: + captured: dict[str, object] = {} + + def _fake_install_runtime_dependencies(**kwargs: object) -> SimpleNamespace: + captured.update(kwargs) + return SimpleNamespace( + target="local", + notes=("runtime ready",), + ) + + monkeypatch.setattr( + "churro_ocr.cli.install_runtime_dependencies", + _fake_install_runtime_dependencies, + ) + + result = cli_runner.invoke( + app, + [ + "install", + "local", + "--torch-backend", + "cu126", + ], + ) + + assert result.exit_code == 0 + assert captured == { + "target": "local", + "torch_backend": "cu126", + } + assert "Installed runtime target: local" in result.output + assert "runtime ready" in result.output + + +def test_install_command_surfaces_configuration_errors( + monkeypatch: pytest.MonkeyPatch, + cli_runner, +) -> None: + def _raise_configuration_error(**_: object) -> SimpleNamespace: + raise ConfigurationError("missing uv") + + monkeypatch.setattr( + "churro_ocr.cli.install_runtime_dependencies", + _raise_configuration_error, + ) + + result = cli_runner.invoke(app, ["install", "hf"]) + + assert result.exit_code == 1 + assert "missing uv" in result.output def test_module_entrypoint_help() -> None: @@ -209,6 +510,7 @@ def test_module_entrypoint_help() -> None: assert result.returncode == 0 assert "transcribe" in result.stdout assert "extract-pages" in result.stdout + assert "install" in result.stdout def test_console_script_help() -> None: @@ -225,3 +527,4 @@ def test_console_script_help() -> None: assert result.returncode == 0 assert "transcribe" in result.stdout assert "extract-pages" in result.stdout + assert "install" in result.stdout diff --git a/tests/test_cli_contract.py b/tests/test_cli_contract.py index 0bffed9..d269117 100644 --- a/tests/test_cli_contract.py +++ b/tests/test_cli_contract.py @@ -23,12 +23,19 @@ def sample_image_path(write_image_file) -> Path: (["--backend", "litellm"], ("--model is required for backend=litellm",)), ( ["--backend", "openai-compatible", "--model", "local-model"], - ("required for", "backend=openai-compatible"), + ("--model and --base-url are required for", "backend=openai-compatible"), ), (["--backend", "azure"], ("--endpoint and --api-key are required for backend=azure",)), (["--backend", "mistral"], ("--api-key is required for backend=mistral",)), + ( + ["--backend", "mistral", "--api-key", "secret"], + ("--model is required for backend=mistral", "mistral-ocr-2505", "mistral-ocr-2512"), + ), + ( + ["--backend", "mistral", "--api-key", "secret", "--model", "mistral-ocr-latest"], + ("must be one of", "mistral-ocr-2505", "mistral-ocr-2512"), + ), (["--backend", "hf"], ("--model is required for backend=hf",)), - (["--backend", "vllm"], ("--model is required for backend=vllm",)), ], ) def test_transcribe_cli_validates_backend_requirements( @@ -48,7 +55,57 @@ def test_transcribe_cli_validates_backend_requirements( assert expected_part in output -def test_transcribe_cli_rejects_unsupported_backend(sample_image_path: Path, cli_runner) -> None: +def test_transcribe_cli_allows_openai_compatible_backend_without_api_key( + monkeypatch: pytest.MonkeyPatch, + sample_image_path: Path, + cli_runner, +) -> None: + captured: dict[str, object] = {} + + class _FakeBackend: + async def ocr(self, page: DocumentPage) -> OCRResult: + captured["size"] = (page.image.width, page.image.height) + return OCRResult(text="ok", provider_name="fake", model_name="fake-model") + + def _fake_build_ocr_backend(spec: cli_module.OCRBackendSpec) -> _FakeBackend: + captured["spec"] = spec + return _FakeBackend() + + monkeypatch.setattr( + "churro_ocr.cli.build_ocr_backend", + _fake_build_ocr_backend, + ) + + result = cli_runner.invoke( + app, + [ + "transcribe", + "--image", + str(sample_image_path), + "--backend", + "openai-compatible", + "--model", + "local-model", + "--base-url", + "http://127.0.0.1:8000/v1", + ], + ) + + assert result.exit_code == 0 + assert result.output.strip() == "ok" + spec = captured["spec"] + assert isinstance(spec, cli_module.OCRBackendSpec) + assert spec.transport is not None + assert spec.transport.api_base == "http://127.0.0.1:8000/v1" + assert spec.transport.api_key is None + + +@pytest.mark.parametrize("backend", ["unsupported"]) +def test_transcribe_cli_rejects_unsupported_backend( + sample_image_path: Path, + backend: str, + cli_runner, +) -> None: result = cli_runner.invoke( app, [ @@ -56,14 +113,14 @@ def test_transcribe_cli_rejects_unsupported_backend(sample_image_path: Path, cli "--image", str(sample_image_path), "--backend", - "unsupported", + backend, "--model", "example/model", ], ) assert result.exit_code != 0 - assert "Unsupported backend: unsupported" in result.output + assert f"Unsupported backend: {backend}" in result.output def test_transcribe_cli_echoes_text_without_output( diff --git a/tests/test_hf_ocr.py b/tests/test_hf_ocr.py index 8eabd3c..70bdb71 100644 --- a/tests/test_hf_ocr.py +++ b/tests/test_hf_ocr.py @@ -12,20 +12,51 @@ from churro_ocr.errors import ConfigurationError from churro_ocr.ocr import OCRClient from churro_ocr.page_detection import DocumentPage -from churro_ocr.prompts import DEFAULT_OCR_OUTPUT_TAG +from churro_ocr.prompts import ( + CHANDRA_OCR_LAYOUT_PROMPT, + DEFAULT_OCR_OUTPUT_TAG, + OLMOCR_V4_YAML_PROMPT, + parse_chandra_response, + parse_olmocr_response, +) from churro_ocr.providers import OCRBackendSpec, build_ocr_backend from churro_ocr.providers.hf import ( + ChandraOCR2OCRBackend, Churro3BOCRBackend, + DeepSeekOCR2OCRBackend, + DotsMOCROCRBackend, DotsOCR15OCRBackend, HuggingFaceVisionOCRBackend, + LFM25VLOCRBackend, + PaddleOCRVL15OCRBackend, +) +from churro_ocr.providers.specs import ( + DEFAULT_OCR_MAX_TOKENS, + deepseek_ocr_2_text_postprocessor, + lfm2_5_vl_text_postprocessor, ) -from churro_ocr.providers.specs import DEFAULT_OCR_MAX_TOKENS from churro_ocr.templates import ( + CHANDRA_OCR_2_MODEL_ID, + CHANDRA_OCR_2_OCR_TEMPLATE, CHURRO_3B_XML_TEMPLATE, + DEEPSEEK_OCR_2_MODEL_ID, + DEEPSEEK_OCR_2_OCR_PROMPT, + DEEPSEEK_OCR_2_OCR_TEMPLATE, + DOTS_MOCR_MODEL_ID, + DOTS_MOCR_OCR_PROMPT, + DOTS_MOCR_OCR_TEMPLATE, DOTS_OCR_1_5_MODEL_ID, DOTS_OCR_1_5_OCR_PROMPT, DOTS_OCR_1_5_OCR_TEMPLATE, + LFM2_5_VL_1_6B_MODEL_ID, + LFM2_5_VL_1_6B_OCR_TEMPLATE, + OLMOCR_2_7B_1025_MODEL_ID, + OLMOCR_2_7B_1025_OCR_TEMPLATE, + PADDLEOCR_VL_1_5_MODEL_ID, + PADDLEOCR_VL_1_5_OCR_PROMPT, + PADDLEOCR_VL_1_5_OCR_TEMPLATE, HFChatTemplate, + OCRConversation, ) @@ -34,15 +65,780 @@ def test_hf_chat_template_builds_expected_conversation() -> None: system_message="system text", user_prompt="user text", ) - page = DocumentPage.from_image(Image.new("RGB", (20, 20), color="white")) + page = DocumentPage.from_image(Image.new("RGB", (20, 20), color="white")) + + conversation = template.build_conversation(page) + + assert conversation[0]["role"] == "system" + assert conversation[0]["content"][0]["text"] == "system text" + assert conversation[1]["role"] == "user" + assert conversation[1]["content"][0]["type"] == "image" + assert conversation[1]["content"][1]["text"] == "user text" + + +def test_olmocr_template_builds_prompt_before_image() -> None: + page = DocumentPage.from_image(Image.new("RGB", (20, 20), color="white")) + + conversation = OLMOCR_2_7B_1025_OCR_TEMPLATE.build_conversation(page) + + assert conversation[0]["role"] == "user" + assert conversation[0]["content"][0]["text"] == OLMOCR_V4_YAML_PROMPT + assert conversation[0]["content"][1]["type"] == "image" + + +def test_chandra_template_builds_image_before_prompt() -> None: + page = DocumentPage.from_image(Image.new("RGB", (20, 20), color="white")) + + conversation = CHANDRA_OCR_2_OCR_TEMPLATE.build_conversation(page) + + assert conversation[0]["role"] == "user" + assert conversation[0]["content"][0]["type"] == "image" + assert conversation[0]["content"][1]["text"] == CHANDRA_OCR_LAYOUT_PROMPT + + +def test_deepseek_ocr_2_template_builds_image_before_prompt() -> None: + page = DocumentPage.from_image(Image.new("RGB", (20, 20), color="white")) + + conversation = DEEPSEEK_OCR_2_OCR_TEMPLATE.build_conversation(page) + + assert conversation[0]["role"] == "user" + assert conversation[0]["content"][0]["type"] == "image" + assert conversation[0]["content"][1]["text"] == DEEPSEEK_OCR_2_OCR_PROMPT + + +def test_parse_olmocr_response_extracts_plain_text_and_metadata() -> None: + text, metadata = parse_olmocr_response( + "---\n" + "primary_language: en\n" + "is_rotation_valid: true\n" + "rotation_correction: 0\n" + "is_table: true\n" + "is_diagram: false\n" + "---\n" + "# Heading\n\n" + "
YearValue
190042
\n\n" + "![Figure alt text](page_0_0_100_100.png)\n" + "Paragraph with [reference](https://example.test)." + ) + + assert text == "Heading\n\nYear | Value\n1900 | 42\n\nParagraph with reference." + assert metadata["front_matter"] == { + "primary_language": "en", + "is_rotation_valid": True, + "rotation_correction": 0, + "is_table": True, + "is_diagram": False, + } + assert "Heading" in cast("str", metadata["raw_markdown"]) + + +def test_parse_chandra_response_extracts_plain_text_and_metadata() -> None: + text, metadata = parse_chandra_response( + '

Title

\n' + '

Paragraph with reference.

\n' + '

' + "Checked item

\n" + '
' + "
YearValue
190042
" + ) + + assert text == "Title\n\nParagraph with reference.\n\n[x] Checked item\n\nYear | Value\n1900 | 42" + assert metadata == { + "raw_html": ( + '

Title

\n' + '

Paragraph with reference.

\n' + '

' + "Checked item

\n" + '
' + "
YearValue
190042
" + ), + } + + +def test_lfm25_text_postprocessor_strips_prompt_and_role_scaffold() -> None: + text = ( + "Transcribe all visible text from this historical document page in reading order.\n" + "assistant\n" + f"<{DEFAULT_OCR_OUTPUT_TAG}>\n" + "decoded text\n" + f"" + ) + + assert lfm2_5_vl_text_postprocessor(text) == "decoded text" + + +def test_lfm25_text_postprocessor_strips_role_only_prefix() -> None: + assert lfm2_5_vl_text_postprocessor("assistant:\nplain text") == "plain text" + + +def test_deepseek_ocr_2_text_postprocessor_strips_prompt_echo_and_stop_token() -> None: + assert ( + deepseek_ocr_2_text_postprocessor( + "\nFree OCR.\n<｜Assistant｜>\nplain text<｜end▁of▁sentence｜>" + ) + == "plain text" + ) + + +def test_build_ocr_backend_uses_chandra_profile_defaults_for_hf() -> None: + backend = cast( + "ChandraOCR2OCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=CHANDRA_OCR_2_MODEL_ID, + ) + ), + ) + + assert isinstance(backend, ChandraOCR2OCRBackend) + assert backend.template == CHANDRA_OCR_2_OCR_TEMPLATE + assert backend.model_name == "chandra-ocr-2" + assert backend.generation_kwargs == { + "max_new_tokens": 12_384, + } + assert backend.image_preprocessor(Image.new("RGB", (5_000, 3_000), color="white")).size == (3_248, 1_932) + + +def test_build_ocr_backend_uses_deepseek_ocr_2_profile_defaults_for_hf() -> None: + backend = cast( + "DeepSeekOCR2OCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=DEEPSEEK_OCR_2_MODEL_ID, + ) + ), + ) + + assert isinstance(backend, DeepSeekOCR2OCRBackend) + assert backend.template == DEEPSEEK_OCR_2_OCR_TEMPLATE + assert backend.model_name == "DeepSeek-OCR-2" + assert backend.generation_kwargs == {"max_new_tokens": 8_192} + assert backend.trust_remote_code is True + assert backend.processor_kwargs == {} + assert backend.model_kwargs == {"use_safetensors": True} + assert backend.base_size == 1_024 + assert backend.image_size == 768 + assert backend.crop_mode is True + + +def test_build_ocr_backend_uses_olmocr_profile_defaults_for_hf() -> None: + backend = cast( + "HuggingFaceVisionOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=OLMOCR_2_7B_1025_MODEL_ID, + ) + ), + ) + + assert backend.template == OLMOCR_2_7B_1025_OCR_TEMPLATE + assert backend.model_name == "olmOCR-2-7B-1025" + assert backend.generation_kwargs == { + "max_new_tokens": 8_000, + "temperature": 0.1, + "do_sample": True, + } + assert backend.image_preprocessor(Image.new("RGB", (5_000, 3_000), color="white")).size == (1_288, 772) + + +def test_build_ocr_backend_uses_lfm25_profile_defaults_for_hf() -> None: + backend = cast( + "LFM25VLOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=LFM2_5_VL_1_6B_MODEL_ID, + ) + ), + ) + + assert isinstance(backend, LFM25VLOCRBackend) + assert backend.template == LFM2_5_VL_1_6B_OCR_TEMPLATE + assert backend.model_name == "LFM2.5-VL-1.6B" + assert backend.generation_kwargs == { + "max_new_tokens": 512, + "do_sample": False, + "repetition_penalty": 1.05, + } + + +def test_build_ocr_backend_uses_paddleocr_vl_profile_defaults_for_hf() -> None: + backend = cast( + "PaddleOCRVL15OCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=PADDLEOCR_VL_1_5_MODEL_ID, + ) + ), + ) + + assert isinstance(backend, PaddleOCRVL15OCRBackend) + assert backend.template == PADDLEOCR_VL_1_5_OCR_TEMPLATE + assert backend.model_name == "PaddleOCR-VL-1.5" + assert backend.generation_kwargs == { + "max_new_tokens": 4_096, + "do_sample": False, + } + + +def test_build_ocr_backend_uses_dots_mocr_profile_defaults_for_hf() -> None: + backend = cast( + "DotsMOCROCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=DOTS_MOCR_MODEL_ID, + ) + ), + ) + + assert isinstance(backend, DotsMOCROCRBackend) + assert backend.template == DOTS_MOCR_OCR_TEMPLATE + assert backend.model_name == "dots.mocr" + assert backend.generation_kwargs == {"max_new_tokens": DEFAULT_OCR_MAX_TOKENS} + assert backend.trust_remote_code is True + assert backend.processor_kwargs == {} + assert backend.model_kwargs["dtype"] in {"auto", "float32"} + if backend.model_kwargs["dtype"] == "auto" and "device_map" in backend.model_kwargs: + assert backend.model_kwargs["device_map"] == "auto" + assert "max_memory" in backend.model_kwargs + + +@pytest.mark.asyncio +async def test_chandra_huggingface_backend_matches_upstream_chat_template_and_eos_behavior( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + class FakeAttentionMask: + def sum(self, dim: int) -> SimpleNamespace: + captured["attention_mask_sum_dim"] = dim + return SimpleNamespace(tolist=lambda: [3]) + + class FakeBatch(dict[str, object]): + def to(self, device: object) -> FakeBatch: + captured["device"] = device + return self + + class FakeTokenizer: + def __init__(self) -> None: + self.padding_side = "right" + + def convert_tokens_to_ids(self, token: str) -> int: + captured["stop_token_lookup"] = token + return 77 + + class FakeProcessor: + def __init__(self) -> None: + self.tokenizer = FakeTokenizer() + + def apply_chat_template( + self, + conversation: object, + *, + add_generation_prompt: bool, + tokenize: bool = True, + return_dict: bool = True, + return_tensors: str | None = None, + padding: bool | None = None, + ) -> object: + captured["add_generation_prompt"] = add_generation_prompt + captured["tokenize"] = tokenize + captured["return_dict"] = return_dict + if not tokenize: + captured["render_conversation"] = conversation + return "" + captured["tokenized_conversations"] = conversation + captured["return_tensors"] = return_tensors + captured["padding"] = padding + return FakeBatch( + { + "input_ids": SimpleNamespace(shape=(1, 3)), + "attention_mask": FakeAttentionMask(), + } + ) + + def batch_decode( + self, + generated_ids: object, + *, + skip_special_tokens: bool, + clean_up_tokenization_spaces: bool, + ) -> list[str]: + captured["generated_ids"] = generated_ids + captured["skip_special_tokens"] = skip_special_tokens + captured["clean_up_tokenization_spaces"] = clean_up_tokenization_spaces + return [ + '

Decoded ' + 'output.

' + ] + + class FakeProcessorCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeProcessor: + captured["processor_model_id"] = model_id + captured["processor_from_pretrained_kwargs"] = kwargs + return FakeProcessor() + + class FakeModel: + device = "fake-device" + dtype = None + generation_config = SimpleNamespace(eos_token_id=11) + + def eval(self) -> FakeModel: + captured["eval_called"] = True + return self + + def generate(self, **kwargs: object) -> list[list[int | str]]: + captured["generate_kwargs"] = kwargs + return [[0, 1, 2, "completion"]] + + class FakeModelCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeModel: + captured["model_model_id"] = model_id + captured["model_from_pretrained_kwargs"] = kwargs + return FakeModel() + + monkeypatch.setattr( + "churro_ocr.providers.hf._load_hf_runtime", + lambda: SimpleNamespace( + processor_cls=FakeProcessorCls, + model_cls=FakeModelCls, + process_vision_info=None, + ), + ) + + backend = cast( + "ChandraOCR2OCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=CHANDRA_OCR_2_MODEL_ID, + ) + ), + ) + result = await backend.ocr( + DocumentPage.from_image(Image.new("RGBA", (5_000, 3_000), color=(255, 255, 255, 255))) + ) + + assert result.text == "Decoded output." + assert result.metadata == { + "raw_html": ( + '

Decoded ' + 'output.

' + ), + } + assert captured["processor_model_id"] == CHANDRA_OCR_2_MODEL_ID + assert captured["model_model_id"] == CHANDRA_OCR_2_MODEL_ID + assert captured["eval_called"] is True + assert captured["render_conversation"][0]["role"] == "user" + render_content = cast("list[dict[str, object]]", captured["render_conversation"][0]["content"]) + assert render_content[0]["type"] == "image" + render_image = cast("Image.Image", render_content[0]["image"]) + assert render_image.size == (3_248, 1_932) + assert render_image.mode == "RGB" + assert render_content[1] == {"type": "text", "text": CHANDRA_OCR_LAYOUT_PROMPT} + tokenized_conversation = cast("list[OCRConversation]", captured["tokenized_conversations"]) + assert tokenized_conversation[0][0]["role"] == "user" + assert captured["tokenize"] is True + assert captured["return_dict"] is True + assert captured["return_tensors"] == "pt" + assert captured["padding"] is True + assert captured["device"] == "fake-device" + assert captured["attention_mask_sum_dim"] == 1 + assert captured["stop_token_lookup"] == "<|im_end|>" + assert captured["generate_kwargs"]["max_new_tokens"] == 12_384 + assert captured["generate_kwargs"]["eos_token_id"] == [11, 77] + assert captured["model_from_pretrained_kwargs"]["device_map"] == "auto" + assert "dtype" in cast("dict[str, object]", captured["model_from_pretrained_kwargs"]) + assert captured["skip_special_tokens"] is True + assert captured["clean_up_tokenization_spaces"] is False + + +@pytest.mark.asyncio +async def test_deepseek_ocr_2_huggingface_backend_uses_upstream_infer_contract( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + class FakeTokenizer: + pass + + class FakeTokenizerCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeTokenizer: + captured["tokenizer_model_id"] = model_id + captured["tokenizer_from_pretrained_kwargs"] = kwargs + return FakeTokenizer() + + class FakeModel: + def eval(self) -> FakeModel: + captured["eval_called"] = True + return self + + def cuda(self) -> FakeModel: + captured["cuda_called"] = True + return self + + def to(self, dtype: object) -> FakeModel: + captured["to_dtype"] = dtype + return self + + def infer(self, tokenizer: object, **kwargs: object) -> str: + captured["infer_tokenizer"] = tokenizer + captured["infer_kwargs"] = kwargs + image = Image.open(cast("str", kwargs["image_file"])) + captured["saved_image_mode"] = image.mode + captured["saved_image_size"] = image.size + captured["output_dir_exists"] = Path(cast("str", kwargs["output_path"])).exists() + return "\nFree OCR.\n<｜Assistant｜>\nDecoded text<｜end▁of▁sentence｜>" + + class FakeModelCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeModel: + captured["model_model_id"] = model_id + captured["model_from_pretrained_kwargs"] = kwargs + return FakeModel() + + monkeypatch.setattr( + "churro_ocr.providers.hf._load_hf_auto_model_runtime", + lambda: SimpleNamespace( + processor_cls=FakeTokenizerCls, + model_cls=FakeModelCls, + process_vision_info=None, + ), + ) + monkeypatch.setattr( + "churro_ocr.providers.hf._ensure_deepseek_ocr_2_cuda_runtime", + lambda: SimpleNamespace(bfloat16="fake-bfloat16"), + ) + + backend = cast( + "DeepSeekOCR2OCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=DEEPSEEK_OCR_2_MODEL_ID, + ) + ), + ) + result = await backend.ocr( + DocumentPage.from_image(Image.new("RGBA", (32, 16), color=(255, 255, 255, 255))) + ) + + assert result.text == "Decoded text" + assert result.metadata == {} + assert captured["tokenizer_model_id"] == DEEPSEEK_OCR_2_MODEL_ID + assert captured["model_model_id"] == DEEPSEEK_OCR_2_MODEL_ID + assert captured["tokenizer_from_pretrained_kwargs"] == {"trust_remote_code": True} + assert captured["model_from_pretrained_kwargs"] == { + "trust_remote_code": True, + "use_safetensors": True, + } + assert captured["eval_called"] is True + assert captured["cuda_called"] is True + assert captured["to_dtype"] == "fake-bfloat16" + assert captured["infer_tokenizer"].__class__ is FakeTokenizer + infer_kwargs = cast("dict[str, object]", captured["infer_kwargs"]) + assert infer_kwargs == { + "prompt": "\nFree OCR.", + "image_file": infer_kwargs["image_file"], + "output_path": infer_kwargs["output_path"], + "base_size": 1_024, + "image_size": 768, + "crop_mode": True, + "save_results": False, + "eval_mode": True, + } + assert captured["saved_image_mode"] == "RGB" + assert captured["saved_image_size"] == (32, 16) + assert captured["output_dir_exists"] is True + + +@pytest.mark.asyncio +async def test_lfm25_huggingface_backend_uses_tokenized_chat_template_and_ties_lm_head( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + class FakeAttentionMask: + def sum(self, dim: int) -> SimpleNamespace: + captured["attention_mask_sum_dim"] = dim + return SimpleNamespace(tolist=lambda: [4]) + + class FakeBatch(dict[str, object]): + def to(self, device: object) -> FakeBatch: + captured["device"] = device + return self + + class FakeProcessor: + def __init__(self) -> None: + self.tokenizer = SimpleNamespace(padding_side="right") + + def apply_chat_template( + self, + conversation: object, + *, + add_generation_prompt: bool, + tokenize: bool, + return_dict: bool | None = None, + return_tensors: str | None = None, + padding: bool | None = None, + ) -> object: + captured.setdefault("chat_calls", []).append( + { + "conversation": conversation, + "add_generation_prompt": add_generation_prompt, + "tokenize": tokenize, + "return_dict": return_dict, + "return_tensors": return_tensors, + "padding": padding, + } + ) + if not tokenize: + return "" + return FakeBatch( + { + "input_ids": SimpleNamespace(shape=(1, 4)), + "attention_mask": FakeAttentionMask(), + } + ) + + def __call__(self, **kwargs: object) -> object: + raise AssertionError("processor(...) should not be used for LFM2.5-VL") + + def batch_decode( + self, + generated_ids: object, + *, + skip_special_tokens: bool, + clean_up_tokenization_spaces: bool, + ) -> list[str]: + captured["generated_ids"] = generated_ids + captured["skip_special_tokens"] = skip_special_tokens + captured["clean_up_tokenization_spaces"] = clean_up_tokenization_spaces + return ["lfm transcription"] + + class FakeProcessorCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeProcessor: + captured["processor_model_id"] = model_id + captured["processor_from_pretrained_kwargs"] = kwargs + return FakeProcessor() + + class FakeLmHead: + weight = "original-weight" + + class FakeModel: + device = "fake-device" + dtype = None + + def __init__(self) -> None: + self.lm_head = FakeLmHead() + + def get_input_embeddings(self) -> SimpleNamespace: + captured["get_input_embeddings_called"] = True + return SimpleNamespace(weight="tied-weight") + + def generate(self, **kwargs: object) -> list[list[int | str]]: + captured["generate_kwargs"] = kwargs + return [[0, 1, 2, 3, "completion"]] + + class FakeModelCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeModel: + captured["model_model_id"] = model_id + captured["model_from_pretrained_kwargs"] = kwargs + model = FakeModel() + captured["model"] = model + return model + + monkeypatch.setattr( + "churro_ocr.providers.hf._load_hf_runtime", + lambda: SimpleNamespace( + processor_cls=FakeProcessorCls, + model_cls=FakeModelCls, + process_vision_info=None, + ), + ) + + backend = cast( + "LFM25VLOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=LFM2_5_VL_1_6B_MODEL_ID, + ) + ), + ) + result = await backend.ocr(DocumentPage.from_image(Image.new("RGB", (32, 32), color="white"))) + + assert result.text == "lfm transcription" + assert captured["processor_model_id"] == LFM2_5_VL_1_6B_MODEL_ID + assert captured["model_model_id"] == LFM2_5_VL_1_6B_MODEL_ID + assert captured["get_input_embeddings_called"] is True + assert cast("FakeModel", captured["model"]).lm_head.weight == "tied-weight" + assert cast("FakeProcessor", backend._processor).tokenizer.padding_side == "left" + assert len(cast("list[dict[str, object]]", captured["chat_calls"])) == 2 + assert cast("list[dict[str, object]]", captured["chat_calls"])[0]["tokenize"] is False + assert cast("list[dict[str, object]]", captured["chat_calls"])[1]["tokenize"] is True + assert cast("list[dict[str, object]]", captured["chat_calls"])[1]["return_dict"] is True + assert cast("list[dict[str, object]]", captured["chat_calls"])[1]["return_tensors"] == "pt" + assert cast("list[dict[str, object]]", captured["chat_calls"])[1]["padding"] is False + assert captured["device"] == "fake-device" + assert captured["attention_mask_sum_dim"] == 1 + assert captured["generate_kwargs"] == { + "input_ids": SimpleNamespace(shape=(1, 4)), + "attention_mask": cast("object", captured["generate_kwargs"]["attention_mask"]), + "max_new_tokens": 512, + "do_sample": False, + "repetition_penalty": 1.05, + } + assert captured["generated_ids"] == [["completion"]] + assert captured["skip_special_tokens"] is True + assert captured["clean_up_tokenization_spaces"] is False + + +@pytest.mark.asyncio +async def test_lfm25_huggingface_backend_batches_pages_with_tokenized_chat_template( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + class FakeAttentionMask: + def sum(self, dim: int) -> SimpleNamespace: + captured["attention_mask_sum_dim"] = dim + return SimpleNamespace(tolist=lambda: [4, 4]) + + class FakeBatch(dict[str, object]): + def to(self, device: object) -> FakeBatch: + captured["device"] = device + return self + + class FakeProcessor: + def __init__(self) -> None: + self.tokenizer = SimpleNamespace(padding_side="right") + + def apply_chat_template( + self, + conversation: object, + *, + add_generation_prompt: bool, + tokenize: bool, + return_dict: bool | None = None, + return_tensors: str | None = None, + padding: bool | None = None, + ) -> object: + captured.setdefault("chat_calls", []).append( + { + "conversation": conversation, + "add_generation_prompt": add_generation_prompt, + "tokenize": tokenize, + "return_dict": return_dict, + "return_tensors": return_tensors, + "padding": padding, + } + ) + if not tokenize: + return "" + return FakeBatch( + { + "input_ids": SimpleNamespace(shape=(2, 4)), + "attention_mask": FakeAttentionMask(), + } + ) + + def __call__(self, **kwargs: object) -> object: + raise AssertionError("processor(...) should not be used for LFM2.5-VL batches") + + def batch_decode( + self, + generated_ids: object, + *, + skip_special_tokens: bool, + clean_up_tokenization_spaces: bool, + ) -> list[str]: + captured["generated_ids"] = generated_ids + captured["skip_special_tokens"] = skip_special_tokens + captured["clean_up_tokenization_spaces"] = clean_up_tokenization_spaces + return ["first transcription", "second transcription"] + + class FakeProcessorCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeProcessor: + captured["processor_model_id"] = model_id + captured["processor_from_pretrained_kwargs"] = kwargs + return FakeProcessor() + + class FakeLmHead: + weight = "original-weight" + + class FakeModel: + device = "fake-device" + dtype = None + + def __init__(self) -> None: + self.lm_head = FakeLmHead() + + def get_input_embeddings(self) -> SimpleNamespace: + return SimpleNamespace(weight="tied-weight") + + def generate(self, **kwargs: object) -> list[list[int | str]]: + captured["generate_kwargs"] = kwargs + return [ + [0, 1, 2, 3, "first"], + [0, 1, 2, 3, "second"], + ] + + class FakeModelCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeModel: + captured["model_model_id"] = model_id + captured["model_from_pretrained_kwargs"] = kwargs + return FakeModel() + + monkeypatch.setattr( + "churro_ocr.providers.hf._load_hf_runtime", + lambda: SimpleNamespace( + processor_cls=FakeProcessorCls, + model_cls=FakeModelCls, + process_vision_info=None, + ), + ) - conversation = template.build_conversation(page) + backend = cast( + "LFM25VLOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=LFM2_5_VL_1_6B_MODEL_ID, + ) + ), + ) + results = await backend.ocr_batch( + [ + DocumentPage.from_image(Image.new("RGB", (32, 32), color="white")), + DocumentPage.from_image(Image.new("RGB", (32, 32), color="white")), + ] + ) - assert conversation[0]["role"] == "system" - assert conversation[0]["content"][0]["text"] == "system text" - assert conversation[1]["role"] == "user" - assert conversation[1]["content"][0]["type"] == "image" - assert conversation[1]["content"][1]["text"] == "user text" + assert [result.text for result in results] == ["first transcription", "second transcription"] + assert cast("FakeProcessor", backend._processor).tokenizer.padding_side == "left" + assert len(cast("list[dict[str, object]]", captured["chat_calls"])) == 3 + assert cast("list[dict[str, object]]", captured["chat_calls"])[2]["tokenize"] is True + assert cast("list[dict[str, object]]", captured["chat_calls"])[2]["padding"] is True + assert captured["device"] == "fake-device" + assert captured["attention_mask_sum_dim"] == 1 + assert captured["generated_ids"] == [["first"], ["second"]] + assert captured["skip_special_tokens"] is True + assert captured["clean_up_tokenization_spaces"] is False @pytest.mark.asyncio @@ -252,7 +1048,12 @@ def apply_chat_template( def __call__(self, **kwargs: object) -> FakeBatch: captured["processor_kwargs"] = kwargs - return FakeBatch({"input_ids": SimpleNamespace(shape=(1, 4))}) + return FakeBatch( + { + "input_ids": SimpleNamespace(shape=(1, 4)), + "mm_token_type_ids": "ignored-mm-token-type-ids", + } + ) def batch_decode( self, @@ -345,12 +1146,576 @@ def test_dots_ocr_15_backend_uses_expected_defaults() -> None: assert backend.trust_remote_code is True assert backend.processor_kwargs == {} assert backend.model_kwargs["dtype"] in {"auto", "float32"} - if backend.model_kwargs["dtype"] == "auto": + if backend.model_kwargs["dtype"] == "auto" and "device_map" in backend.model_kwargs: + assert backend.model_kwargs["device_map"] == "auto" + assert "max_memory" in backend.model_kwargs + assert backend.generation_kwargs == {"max_new_tokens": DEFAULT_OCR_MAX_TOKENS} + + +def test_deepseek_ocr_2_backend_uses_expected_defaults() -> None: + backend = DeepSeekOCR2OCRBackend() + + assert backend.model_id == DEEPSEEK_OCR_2_MODEL_ID + assert backend.template == DEEPSEEK_OCR_2_OCR_TEMPLATE + assert backend.model_name == "DeepSeek-OCR-2" + assert backend.trust_remote_code is True + assert backend.processor_kwargs == {} + assert backend.model_kwargs == {"use_safetensors": True} + assert backend.generation_kwargs == {"max_new_tokens": 8_192} + assert backend.base_size == 1_024 + assert backend.image_size == 768 + assert backend.crop_mode is True + + +def test_dots_mocr_backend_uses_expected_defaults() -> None: + backend = DotsMOCROCRBackend() + + assert backend.model_id == DOTS_MOCR_MODEL_ID + assert backend.template == DOTS_MOCR_OCR_TEMPLATE + assert backend.model_name == "dots.mocr" + assert backend.trust_remote_code is True + assert backend.processor_kwargs == {} + assert backend.model_kwargs["dtype"] in {"auto", "float32"} + if backend.model_kwargs["dtype"] == "auto" and "device_map" in backend.model_kwargs: assert backend.model_kwargs["device_map"] == "auto" assert "max_memory" in backend.model_kwargs assert backend.generation_kwargs == {"max_new_tokens": DEFAULT_OCR_MAX_TOKENS} +def test_dots_mocr_template_matches_upstream_prompt() -> None: + page = DocumentPage.from_image(Image.new("RGB", (20, 20), color="white")) + + conversation = DOTS_MOCR_OCR_TEMPLATE.build_conversation(page) + + assert conversation[0]["role"] == "user" + assert conversation[0]["content"][0]["type"] == "image" + assert conversation[0]["content"][1]["text"] == DOTS_MOCR_OCR_PROMPT + + +@pytest.mark.asyncio +async def test_paddleocr_vl_15_backend_uses_tokenized_chat_template_and_profile_defaults( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + class FakeAttentionMask: + def sum(self, dim: int) -> SimpleNamespace: + captured["attention_mask_sum_dim"] = dim + return SimpleNamespace(tolist=lambda: [4]) + + class FakeBatch(dict[str, object]): + def to(self, device: object) -> FakeBatch: + captured["device"] = device + return self + + class FakeTokenizer: + def __init__(self) -> None: + self.padding_side = "right" + + class FakeProcessor: + def __init__(self) -> None: + self.tokenizer = FakeTokenizer() + self.image_processor = SimpleNamespace(min_pixels=112_896, max_pixels=1_003_520) + + def apply_chat_template( + self, + conversation: object, + *, + add_generation_prompt: bool, + tokenize: bool, + return_dict: bool | None = None, + return_tensors: str | None = None, + processor_kwargs: dict[str, object] | None = None, + ) -> object: + captured.setdefault("chat_calls", []).append( + { + "conversation": conversation, + "add_generation_prompt": add_generation_prompt, + "tokenize": tokenize, + "return_dict": return_dict, + "return_tensors": return_tensors, + "processor_kwargs": processor_kwargs, + } + ) + if not tokenize: + return "" + return FakeBatch( + { + "input_ids": SimpleNamespace(shape=(1, 4)), + "attention_mask": FakeAttentionMask(), + } + ) + + def __call__(self, **kwargs: object) -> object: + raise AssertionError("processor(...) should not be used for PaddleOCR-VL") + + def batch_decode( + self, + generated_ids: object, + *, + skip_special_tokens: bool, + clean_up_tokenization_spaces: bool, + ) -> list[str]: + captured["generated_ids"] = generated_ids + captured["skip_special_tokens"] = skip_special_tokens + captured["clean_up_tokenization_spaces"] = clean_up_tokenization_spaces + return ["OCR:\nassistant\npaddle transcription"] + + class FakeProcessorCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeProcessor: + captured["processor_model_id"] = model_id + captured["processor_from_pretrained_kwargs"] = kwargs + return FakeProcessor() + + class FakeModel: + device = "fake-device" + dtype = None + + def eval(self) -> FakeModel: + captured["eval_called"] = True + return self + + def generate(self, **kwargs: object) -> list[list[int | str]]: + captured["generate_kwargs"] = kwargs + return [[0, 1, 2, 3, "completion"]] + + class FakeModelCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeModel: + captured["model_model_id"] = model_id + captured["model_from_pretrained_kwargs"] = kwargs + return FakeModel() + + monkeypatch.setattr( + "churro_ocr.providers.hf._load_hf_runtime", + lambda: SimpleNamespace( + processor_cls=FakeProcessorCls, + model_cls=FakeModelCls, + process_vision_info=None, + ), + ) + + backend = cast( + "PaddleOCRVL15OCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="hf", + model=PADDLEOCR_VL_1_5_MODEL_ID, + ) + ), + ) + result = await backend.ocr(DocumentPage.from_image(Image.new("RGB", (32, 32), color="white"))) + + assert result.text == "paddle transcription" + assert captured["processor_model_id"] == PADDLEOCR_VL_1_5_MODEL_ID + assert captured["model_model_id"] == PADDLEOCR_VL_1_5_MODEL_ID + assert captured["eval_called"] is True + assert cast("FakeProcessor", backend._processor).tokenizer.padding_side == "left" + assert len(cast("list[dict[str, object]]", captured["chat_calls"])) == 2 + assert cast("list[dict[str, object]]", captured["chat_calls"])[0]["tokenize"] is False + assert cast("list[dict[str, object]]", captured["chat_calls"])[1]["tokenize"] is True + assert cast("list[dict[str, object]]", captured["chat_calls"])[1]["return_dict"] is True + assert cast("list[dict[str, object]]", captured["chat_calls"])[1]["return_tensors"] == "pt" + assert cast("list[dict[str, object]]", captured["chat_calls"])[1]["processor_kwargs"] == { + "text_kwargs": { + "padding": False, + "return_mm_token_type_ids": True, + }, + "images_kwargs": { + "min_pixels": 112_896, + "max_pixels": 1_003_520, + }, + } + render_conversation = cast("list[dict[str, object]]", captured["chat_calls"])[0]["conversation"] + assert cast("list[dict[str, object]]", render_conversation)[0]["role"] == "user" + render_content = cast( + "list[dict[str, object]]", + cast("list[dict[str, object]]", render_conversation)[0]["content"], + ) + assert render_content[0]["type"] == "image" + assert render_content[1] == {"type": "text", "text": PADDLEOCR_VL_1_5_OCR_PROMPT} + assert captured["device"] == "fake-device" + assert captured["attention_mask_sum_dim"] == 1 + assert captured["generate_kwargs"] == { + "input_ids": SimpleNamespace(shape=(1, 4)), + "attention_mask": cast("object", captured["generate_kwargs"]["attention_mask"]), + "max_new_tokens": 4_096, + "do_sample": False, + } + assert captured["generated_ids"] == [["completion"]] + assert captured["skip_special_tokens"] is True + assert captured["clean_up_tokenization_spaces"] is False + + +def test_paddleocr_vl_15_backend_uses_expected_defaults() -> None: + backend = PaddleOCRVL15OCRBackend() + + assert backend.model_id == PADDLEOCR_VL_1_5_MODEL_ID + assert backend.template == PADDLEOCR_VL_1_5_OCR_TEMPLATE + assert backend.trust_remote_code is False + assert backend.processor_kwargs == {} + assert backend.model_kwargs == {} + assert backend.generation_kwargs == { + "max_new_tokens": 4_096, + "do_sample": False, + } + + +def test_patch_dots_ocr_prepare_inputs_for_generation_handles_missing_cache_position() -> None: + captured: dict[str, object] = {} + + class BaseModel: + def prepare_inputs_for_generation( + self, + input_ids: object, + *, + past_key_values: object = None, + inputs_embeds: object = None, + pixel_values: object = None, + attention_mask: object = None, + cache_position: object = None, + num_logits_to_keep: object = None, + **kwargs: object, + ) -> dict[str, object]: + del pixel_values + captured["base_call"] = { + "input_ids": input_ids, + "past_key_values": past_key_values, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "cache_position": cache_position, + "num_logits_to_keep": num_logits_to_keep, + "kwargs": kwargs, + } + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + class FakeDotsModel(BaseModel): + pass + + model = FakeDotsModel() + + hf_module._patch_dots_ocr_prepare_inputs_for_generation(model) + first_inputs = model.prepare_inputs_for_generation( + "tokens", + pixel_values="pixels", + attention_mask="mask", + cache_position=None, + ) + later_inputs = model.prepare_inputs_for_generation( + "tokens", + pixel_values="pixels", + attention_mask="mask", + cache_position=[3], + ) + + assert first_inputs == { + "input_ids": "tokens", + "attention_mask": "mask", + "pixel_values": "pixels", + } + assert later_inputs == { + "input_ids": "tokens", + "attention_mask": "mask", + } + assert captured["base_call"] == { + "input_ids": "tokens", + "past_key_values": None, + "inputs_embeds": None, + "attention_mask": "mask", + "cache_position": [3], + "num_logits_to_keep": None, + "kwargs": {}, + } + + hf_module._patch_dots_ocr_prepare_inputs_for_generation(model) + assert cast("Any", model)._churro_dots_prepare_inputs_patched is True + + +def test_patch_dots_ocr_prepare_inputs_for_generation_skips_wrapped_dots_method() -> None: + captured: dict[str, object] = {} + + class BaseModel: + def prepare_inputs_for_generation( + self, + input_ids: object, + *, + past_key_values: object = None, + inputs_embeds: object = None, + pixel_values: object = None, + attention_mask: object = None, + cache_position: object = None, + num_logits_to_keep: object = None, + **kwargs: object, + ) -> dict[str, object]: + del pixel_values + captured["base_call"] = { + "input_ids": input_ids, + "past_key_values": past_key_values, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "cache_position": cache_position, + "num_logits_to_keep": num_logits_to_keep, + "kwargs": kwargs, + } + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + class FakeDotsModel(BaseModel): + def prepare_inputs_for_generation( + self, + input_ids: object, + *, + past_key_values: object = None, + inputs_embeds: object = None, + pixel_values: object = None, + attention_mask: object = None, + cache_position: object = None, + num_logits_to_keep: object = None, + **kwargs: object, + ) -> dict[str, object]: + if cast("Any", cache_position)[0] == 0: + return {"pixel_values": pixel_values} + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + class WrappedFakeDotsModel(FakeDotsModel): + pass + + model = WrappedFakeDotsModel() + + hf_module._patch_dots_ocr_prepare_inputs_for_generation(model) + inputs = model.prepare_inputs_for_generation( + "tokens", + pixel_values="pixels", + attention_mask="mask", + cache_position=None, + ) + + assert inputs == { + "input_ids": "tokens", + "attention_mask": "mask", + "pixel_values": "pixels", + } + assert captured["base_call"] == { + "input_ids": "tokens", + "past_key_values": None, + "inputs_embeds": None, + "attention_mask": "mask", + "cache_position": None, + "num_logits_to_keep": None, + "kwargs": {}, + } + + +def test_patch_dots_ocr_prepare_inputs_for_generation_skips_duplicate_wrapped_dots_methods() -> None: + captured: dict[str, object] = {} + + class BaseModel: + def prepare_inputs_for_generation( + self, + input_ids: object, + *, + past_key_values: object = None, + inputs_embeds: object = None, + pixel_values: object = None, + attention_mask: object = None, + cache_position: object = None, + num_logits_to_keep: object = None, + **kwargs: object, + ) -> dict[str, object]: + del pixel_values + captured["base_call"] = { + "input_ids": input_ids, + "past_key_values": past_key_values, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "cache_position": cache_position, + "num_logits_to_keep": num_logits_to_keep, + "kwargs": kwargs, + } + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + } + + def shared_prepare_inputs_for_generation( + self: object, + input_ids: object, + *, + past_key_values: object = None, + inputs_embeds: object = None, + pixel_values: object = None, + attention_mask: object = None, + cache_position: object = None, + num_logits_to_keep: object = None, + **kwargs: object, + ) -> dict[str, object]: + del self, input_ids, past_key_values, inputs_embeds, pixel_values, attention_mask + del num_logits_to_keep, kwargs + if cast("Any", cache_position)[0] == 0: + return {"unexpected": True} + return {"unexpected": False} + + class FakeDotsOwner(BaseModel): + pass + + class WrappedFakeDotsModel(FakeDotsOwner): + pass + + cast("Any", FakeDotsOwner).prepare_inputs_for_generation = shared_prepare_inputs_for_generation + cast("Any", WrappedFakeDotsModel).prepare_inputs_for_generation = ( + shared_prepare_inputs_for_generation + ) + + model = WrappedFakeDotsModel() + + hf_module._patch_dots_ocr_prepare_inputs_for_generation(model) + inputs = model.prepare_inputs_for_generation( + "tokens", + pixel_values="pixels", + attention_mask="mask", + cache_position=None, + ) + + assert inputs == { + "input_ids": "tokens", + "attention_mask": "mask", + "pixel_values": "pixels", + } + assert captured["base_call"] == { + "input_ids": "tokens", + "past_key_values": None, + "inputs_embeds": None, + "attention_mask": "mask", + "cache_position": None, + "num_logits_to_keep": None, + "kwargs": {}, + } + + +def test_dots_ocr_15_backend_batch_strips_unused_mm_token_type_ids(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + class FakeAttentionMask: + def sum(self, dim: int) -> SimpleNamespace: + captured["attention_mask_sum_dim"] = dim + return SimpleNamespace(tolist=lambda: [4, 4]) + + class FakeBatch(dict[str, object]): + def to(self, device: object) -> FakeBatch: + captured["device"] = device + return self + + class FakeProcessor: + tokenizer = object() + + def apply_chat_template( + self, + conversation: list[dict[str, object]], + *, + add_generation_prompt: bool, + tokenize: bool, + ) -> str: + del add_generation_prompt, tokenize + image = cast(Image.Image, cast(list[dict[str, object]], conversation[0]["content"])[0]["image"]) + return f"" + + def __call__(self, **kwargs: object) -> FakeBatch: + captured["processor_kwargs"] = kwargs + return FakeBatch( + { + "input_ids": SimpleNamespace(shape=(2, 4)), + "attention_mask": FakeAttentionMask(), + "mm_token_type_ids": "ignored-mm-token-type-ids", + } + ) + + def batch_decode( + self, + generated_ids: object, + *, + skip_special_tokens: bool, + clean_up_tokenization_spaces: bool, + ) -> list[str]: + captured["generated_ids"] = generated_ids + captured["decode_kwargs"] = (skip_special_tokens, clean_up_tokenization_spaces) + return ["dots batch 1", "dots batch 2"] + + class FakeProcessorCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeProcessor: + captured["processor_model_id"] = model_id + captured["processor_from_pretrained_kwargs"] = kwargs + return FakeProcessor() + + class FakeModel: + device = "fake-device" + + def generate(self, **kwargs: object) -> list[list[int]]: + captured["generate_kwargs"] = kwargs + return [ + [100, 101, 102, 103, 104, 105], + [200, 201, 202, 203, 204, 205], + ] + + class FakeModelCls: + @staticmethod + def from_pretrained(model_id: str, **kwargs: object) -> FakeModel: + captured["model_model_id"] = model_id + captured["model_from_pretrained_kwargs"] = kwargs + return FakeModel() + + def fake_process_vision_info( + conversation: list[dict[str, object]], + *, + return_video_kwargs: bool, + return_video_metadata: bool, + ) -> tuple[object, None, None]: + del conversation, return_video_kwargs, return_video_metadata + return "fake-image-inputs", None, None + + monkeypatch.setattr( + "churro_ocr.providers.hf._load_hf_causal_runtime", + lambda: SimpleNamespace( + processor_cls=FakeProcessorCls, + model_cls=FakeModelCls, + process_vision_info=fake_process_vision_info, + ), + ) + monkeypatch.setattr( + "churro_ocr.providers.hf._prepare_dots_ocr_model_dir", + lambda model_id: model_id, + ) + + backend = DotsOCR15OCRBackend() + results = backend._ocr_batch_sync( + [ + DocumentPage.from_image(Image.new("RGB", (32, 32), color="white")), + DocumentPage.from_image(Image.new("RGB", (64, 64), color="white")), + ] + ) + + assert [result.text for result in results] == ["dots batch 1", "dots batch 2"] + assert captured["processor_kwargs"]["text"] == ["", ""] + assert captured["processor_kwargs"]["images"] == [["fake-image-inputs"], ["fake-image-inputs"]] + assert captured["generate_kwargs"] == { + "input_ids": SimpleNamespace(shape=(2, 4)), + "attention_mask": cast("object", captured["generate_kwargs"]["attention_mask"]), + "max_new_tokens": DEFAULT_OCR_MAX_TOKENS, + } + assert captured["attention_mask_sum_dim"] == 1 + assert captured["generated_ids"] == [[104, 105], [204, 205]] + assert captured["decode_kwargs"] == (True, False) + + @pytest.mark.asyncio async def test_huggingface_vision_ocr_backend_strips_default_output_tags( monkeypatch: pytest.MonkeyPatch, @@ -459,6 +1824,7 @@ def test_hf_runtime_loaders_use_installed_modules( cast(Any, transformers_module).AutoModelForImageTextToText = image_text_model_cls cast(Any, transformers_module).AutoModelForCausalLM = causal_model_cls + monkeypatch.setitem(sys.modules, "torch", ModuleType("torch")) monkeypatch.setitem(sys.modules, "qwen_vl_utils", qwen_module) monkeypatch.setitem(sys.modules, "transformers", transformers_module) diff --git a/tests/test_install.py b/tests/test_install.py new file mode 100644 index 0000000..7748202 --- /dev/null +++ b/tests/test_install.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest + +import churro_ocr._internal.install as install_module +from churro_ocr.errors import ConfigurationError + + +class _FakeDistribution: + def __init__(self, *, requires: list[str] | None) -> None: + self.requires = requires + + +def test_install_runtime_dependencies_installs_hf_and_torch_with_uv( + monkeypatch: pytest.MonkeyPatch, +) -> None: + commands: list[list[str]] = [] + + monkeypatch.setattr( + install_module.metadata, + "distribution", + lambda _: _FakeDistribution( + requires=[ + 'qwen-vl-utils; extra == "hf"', + 'transformers>=5,<6; extra == "hf"', + ] + ), + ) + monkeypatch.setattr(install_module.shutil, "which", lambda name: "/usr/bin/uv" if name == "uv" else None) + monkeypatch.setattr( + install_module.subprocess, + "run", + lambda command, check=True: commands.append(list(command)) or SimpleNamespace(returncode=0), + ) + + result = install_module.install_runtime_dependencies( + target="hf", + torch_backend="cu126", + ) + + assert result.target == "hf" + assert commands == [ + [ + "/usr/bin/uv", + "pip", + "install", + "--python", + sys.executable, + "--upgrade", + "qwen-vl-utils", + "transformers>=5,<6", + ], + [ + "/usr/bin/uv", + "pip", + "install", + "--python", + sys.executable, + "--upgrade", + "--torch-backend=cu126", + "torch", + "torchvision", + ], + ] + + +def test_install_runtime_dependencies_installs_local_client_with_uv( + monkeypatch: pytest.MonkeyPatch, +) -> None: + commands: list[list[str]] = [] + + monkeypatch.setattr( + install_module.metadata, + "distribution", + lambda _: _FakeDistribution(requires=['litellm[caching]==1.82.3; extra == "local"']), + ) + monkeypatch.setattr(install_module.shutil, "which", lambda name: "/usr/bin/uv" if name == "uv" else None) + monkeypatch.setattr( + install_module.subprocess, + "run", + lambda command, check=True: commands.append(list(command)) or SimpleNamespace(returncode=0), + ) + + result = install_module.install_runtime_dependencies(target="local") + + assert result.target == "local" + assert commands == [ + [ + "/usr/bin/uv", + "pip", + "install", + "--python", + sys.executable, + "--upgrade", + "litellm[caching]==1.82.3", + ] + ] + assert result.notes == () + + +def test_install_runtime_dependencies_requires_uv(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(install_module.shutil, "which", lambda _: None) + + with pytest.raises(ConfigurationError, match="requires `uv` on PATH"): + install_module.install_runtime_dependencies(target="hf") diff --git a/tests/test_internal_helpers.py b/tests/test_internal_helpers.py index 7f33cfe..427bdb4 100644 --- a/tests/test_internal_helpers.py +++ b/tests/test_internal_helpers.py @@ -1,5 +1,6 @@ from __future__ import annotations +import re import sys from base64 import b64encode from threading import Lock @@ -11,13 +12,16 @@ import churro_ocr._internal.litellm as litellm_module import churro_ocr._internal.prompt_logging as prompt_logging_module +import churro_ocr._internal.retry as retry_module from churro_ocr._internal import logging as logging_module from churro_ocr._internal.image import image_to_base64, load_image from churro_ocr._internal.litellm import LiteLLMTransport, complete_text from churro_ocr.errors import ConfigurationError, ProviderError +from churro_ocr.page_detection import DocumentPage +from churro_ocr.providers._shared import render_ocr_prompt from churro_ocr.providers.hf import _load_hf_causal_runtime, _load_hf_runtime from churro_ocr.providers.specs import LiteLLMTransportConfig -from churro_ocr.providers.vllm import _load_vllm_processor_cls, _load_vllm_runtime +from churro_ocr.templates import HFChatTemplate def _make_fake_litellm_module(*, acompletion: object, completion_cost: object | None = None) -> ModuleType: @@ -93,6 +97,42 @@ def test_prepare_messages_from_conversation_converts_images_and_preserves_unknow ] +def test_render_ocr_prompt_supports_transformers_v5_chat_template_contract() -> None: + captured: dict[str, object] = {} + + class FakeProcessor: + def apply_chat_template( + self, + conversation: list[dict[str, object]], + *, + add_generation_prompt: bool, + tokenize: bool = True, + return_dict: bool = True, + ) -> object: + captured["conversation"] = conversation + captured["add_generation_prompt"] = add_generation_prompt + captured["tokenize"] = tokenize + captured["return_dict"] = return_dict + if not tokenize: + return "" + return {"input_ids": [1, 2, 3]} if return_dict else [1, 2, 3] + + page = DocumentPage.from_image(Image.new("RGB", (16, 16), color="white")) + + rendered, conversation = render_ocr_prompt( + FakeProcessor(), + HFChatTemplate(user_prompt="prompt"), + page, + add_generation_prompt=True, + ) + + assert rendered == "" + assert conversation == captured["conversation"] + assert captured["add_generation_prompt"] is True + assert captured["tokenize"] is False + assert captured["return_dict"] is True + + @pytest.mark.asyncio async def test_complete_text_wrapper_passes_transport_config(monkeypatch: pytest.MonkeyPatch) -> None: captured: dict[str, object] = {} @@ -104,12 +144,14 @@ async def _fake_complete_text( messages: list[dict[str, object]], timeout_seconds: int = 600, output_json: bool = False, + allow_empty: bool = False, ) -> str: captured["config"] = self.config captured["model"] = model captured["messages"] = messages captured["timeout_seconds"] = timeout_seconds captured["output_json"] = output_json + captured["allow_empty"] = allow_empty return "ok" monkeypatch.setattr( @@ -137,6 +179,7 @@ async def _fake_complete_text( assert config.completion_kwargs == {"temperature": 0} assert captured["timeout_seconds"] == 42 assert captured["output_json"] is True + assert captured["allow_empty"] is False def test_extract_response_cost_uses_completion_cost_fallback(monkeypatch: pytest.MonkeyPatch) -> None: @@ -223,7 +266,10 @@ def test_configure_disk_cache_enables_and_updates_cache(monkeypatch: pytest.Monk async def test_transport_complete_text_raises_provider_error_on_failure( monkeypatch: pytest.MonkeyPatch, ) -> None: + calls = {"acompletion": 0} + async def _failing_acompletion(**_: object) -> object: + calls["acompletion"] += 1 raise RuntimeError("boom") fake_module = _make_fake_litellm_module(acompletion=_failing_acompletion) @@ -236,6 +282,49 @@ async def _failing_acompletion(**_: object) -> object: model="example/model", messages=[{"role": "user", "content": [{"type": "text", "text": "hello"}]}], ) + assert calls == {"acompletion": 1} + + +@pytest.mark.asyncio +async def test_transport_complete_text_retries_transient_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"acompletion": 0} + sleep_calls: list[float] = [] + + class FakeLiteLLMError(Exception): + def __init__(self, status_code: int, headers: dict[str, str] | None = None) -> None: + self.status_code = status_code + self.headers = headers or {} + self.raw_response = SimpleNamespace(headers=self.headers) + super().__init__(f"Status {status_code}") + + async def _fake_sleep(delay: float) -> None: + sleep_calls.append(delay) + + async def _flaky_acompletion(**_: object) -> object: + calls["acompletion"] += 1 + if calls["acompletion"] == 1: + raise FakeLiteLLMError(429, headers={"retry-after": "3"}) + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))], + _hidden_params={}, + ) + + fake_module = _make_fake_litellm_module(acompletion=_flaky_acompletion) + monkeypatch.setitem(sys.modules, "litellm", fake_module) + monkeypatch.setattr(litellm_module, "_INITIALIZED", False) + monkeypatch.setattr(retry_module, "retry_sleep", _fake_sleep) + + transport = LiteLLMTransport() + result = await transport.complete_text( + model="example/model", + messages=[{"role": "user", "content": [{"type": "text", "text": "hello"}]}], + ) + + assert result == "ok" + assert calls == {"acompletion": 2} + assert sleep_calls == [3.0] @pytest.mark.asyncio @@ -258,6 +347,28 @@ async def _empty_acompletion(**_: object) -> object: ) +@pytest.mark.asyncio +async def test_transport_complete_text_allows_empty_output(monkeypatch: pytest.MonkeyPatch) -> None: + async def _empty_acompletion(**_: object) -> object: + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=" "))], + _hidden_params={}, + ) + + fake_module = _make_fake_litellm_module(acompletion=_empty_acompletion) + monkeypatch.setitem(sys.modules, "litellm", fake_module) + monkeypatch.setattr(litellm_module, "_INITIALIZED", False) + + transport = LiteLLMTransport() + result = await transport.complete_text( + model="example/model", + messages=[{"role": "user", "content": [{"type": "text", "text": "hello"}]}], + allow_empty=True, + ) + + assert result == "" + + def test_logger_adapter_supports_success_fallback_and_other_levels() -> None: records: list[tuple[str, str]] = [] @@ -351,20 +462,34 @@ def debug(self, message: str, *args: object) -> None: @pytest.mark.parametrize( - ("loader", "dependency_name"), + ("loader", "dependency_name", "message"), [ - (_load_vllm_processor_cls, "transformers"), - (_load_vllm_runtime, "vllm"), - (_load_hf_runtime, "qwen_vl_utils"), - (_load_hf_causal_runtime, "qwen_vl_utils"), + (_load_hf_runtime, "torch", "PyTorch runtime"), + (_load_hf_runtime, "qwen_vl_utils", "install hf"), + (_load_hf_causal_runtime, "torch", "PyTorch runtime"), + (_load_hf_causal_runtime, "qwen_vl_utils", "install hf"), ], ) def test_optional_dependency_loaders_raise_configuration_error( loader: Any, dependency_name: str, + message: str, patch_import_failure, + monkeypatch: pytest.MonkeyPatch, ) -> None: - patch_import_failure(failing_name=dependency_name) + if dependency_name == "torch": + + def _fake_import_module(name: str) -> object: + if name == "torch": + raise ImportError("missing torch") + return __import__(name) + + monkeypatch.setattr("churro_ocr.providers.hf.import_module", _fake_import_module) + elif dependency_name == "qwen_vl_utils": + monkeypatch.setitem(sys.modules, "torch", ModuleType("torch")) + patch_import_failure(failing_name=dependency_name) + else: + patch_import_failure(failing_name=dependency_name) - with pytest.raises(ConfigurationError): + with pytest.raises(ConfigurationError, match=re.escape(message)): loader() diff --git a/tests/test_package_check.py b/tests/test_package_check.py index 4ee8953..a4dfcfa 100644 --- a/tests/test_package_check.py +++ b/tests/test_package_check.py @@ -32,7 +32,7 @@ def _always_missing(_: str): monkeypatch.setattr(package_check.metadata, "distribution", _always_missing) - package_check._audit_dependency_licenses(_metadata_message('vllm<1,>=0.18; extra == "vllm"')) + package_check._audit_dependency_licenses(_metadata_message('mistralai<2,>=1.6.0; extra == "mistral"')) def test_license_audit_fails_for_missing_base_dependency(monkeypatch: pytest.MonkeyPatch) -> None: @@ -43,3 +43,13 @@ def _always_missing(_: str): with pytest.raises(RuntimeError, match="pillow \\(not installed in the Pixi audit environment\\)"): package_check._audit_dependency_licenses(_metadata_message("Pillow<12,>=10.4.0")) + + +def test_local_runtime_packaging_policy_rejects_direct_torch_runtime_pin() -> None: + metadata_message = _metadata_message( + 'transformers[torch]>=5,<6; extra == "hf"', + 'torchvision; extra == "all"', + ) + + with pytest.raises(RuntimeError, match="must not pin local PyTorch"): + package_check._assert_local_runtime_packaging_policy(metadata_message) diff --git a/tests/test_provider_api_contracts.py b/tests/test_provider_api_contracts.py index bde918f..71aa64c 100644 --- a/tests/test_provider_api_contracts.py +++ b/tests/test_provider_api_contracts.py @@ -14,7 +14,6 @@ MistralOptions, OCRBackendSpec, OpenAICompatibleOptions, - VLLMOptions, build_ocr_backend, resolve_ocr_profile, ) @@ -59,7 +58,7 @@ def test_provider_dir_lists_lazy_exports() -> None: (OCRBackendSpec(provider="litellm"), "OCR provider 'litellm' requires `model`."), ( OCRBackendSpec(provider="openai-compatible", model="local-model"), - "OCR provider 'openai-compatible' requires `transport.api_base` and `transport.api_key`.", + "OCR provider 'openai-compatible' requires `transport.api_base`.", ), ( OCRBackendSpec(provider="azure"), @@ -69,21 +68,42 @@ def test_provider_dir_lists_lazy_exports() -> None: OCRBackendSpec(provider="mistral"), "OCR provider 'mistral' requires MistralOptions(api_key=...).", ), + ( + OCRBackendSpec( + provider="mistral", + options=MistralOptions(api_key="secret"), + ), + "OCR provider 'mistral' requires `model` to be one of: mistral-ocr-2505, mistral-ocr-2512.", + ), + ( + OCRBackendSpec( + provider="mistral", + model="mistral-ocr-latest", + options=MistralOptions(api_key="secret"), + ), + ( + "OCR provider 'mistral' only supports `model` values mistral-ocr-2505, " + "mistral-ocr-2512; got 'mistral-ocr-latest'." + ), + ), ( OCRBackendSpec( provider="hf", model="example/model", - options=cast("Any", VLLMOptions()), + options=cast("Any", OpenAICompatibleOptions()), ), - "OCR provider 'hf' requires options of type HuggingFaceOptions, got VLLMOptions.", + "OCR provider 'hf' requires options of type HuggingFaceOptions, got OpenAICompatibleOptions.", ), ( OCRBackendSpec( - provider="vllm", + provider="openai-compatible", model="example/model", options=cast("Any", HuggingFaceOptions()), ), - "OCR provider 'vllm' requires options of type VLLMOptions, got HuggingFaceOptions.", + ( + "OCR provider 'openai-compatible' requires options of type " + "OpenAICompatibleOptions, got HuggingFaceOptions." + ), ), ], ) @@ -99,7 +119,6 @@ def test_build_ocr_backend_supports_custom_openai_model_prefix() -> None: model="local-model", transport=LiteLLMTransportConfig( api_base="http://127.0.0.1:8000/v1", - api_key="dummy", ), options=OpenAICompatibleOptions(model_prefix="custom"), ) @@ -134,6 +153,7 @@ def test_build_ocr_backend_accepts_provider_specific_options() -> None: mistral_backend = build_ocr_backend( OCRBackendSpec( provider="mistral", + model="mistral-ocr-2512", options=MistralOptions(api_key="secret"), ) ) @@ -142,4 +162,4 @@ def test_build_ocr_backend_accepts_provider_specific_options() -> None: assert isinstance(mistral_backend, MistralOCRBackend) assert azure_backend.model_id == "layout-model" assert azure_backend.model_name == "layout-model" - assert mistral_backend.model == "mistral-ocr-latest" + assert mistral_backend.model == "mistral-ocr-2512" diff --git a/tests/test_providers.py b/tests/test_providers.py index 065d08b..c7c9828 100644 --- a/tests/test_providers.py +++ b/tests/test_providers.py @@ -9,10 +9,16 @@ import pytest from PIL import Image +import churro_ocr._internal.retry as retry_module from churro_ocr._internal.litellm import LiteLLMTransport from churro_ocr.errors import ProviderError from churro_ocr.page_detection import DocumentPage -from churro_ocr.prompts import DEFAULT_BOUNDARY_DETECTION_PROMPT, DEFAULT_OCR_OUTPUT_TAG +from churro_ocr.prompts import ( + CHANDRA_OCR_LAYOUT_PROMPT, + DEFAULT_BOUNDARY_DETECTION_PROMPT, + DEFAULT_OCR_OUTPUT_TAG, + OLMOCR_V4_YAML_PROMPT, +) from churro_ocr.providers import ( AzureDocumentIntelligenceOptions, AzurePageDetector, @@ -22,7 +28,6 @@ MistralOptions, OCRBackendSpec, OpenAICompatibleOptions, - VLLMOptions, build_ocr_backend, locate_text_block_bbox_with_llm, resolve_ocr_profile, @@ -32,6 +37,7 @@ AzureDocumentIntelligenceOCRBackend, LiteLLMVisionOCRBackend, MistralOCRBackend, + OpenAICompatibleOCRBackend, ) from churro_ocr.providers.page_detection import ( _normalize_azure_page_polygon, @@ -39,8 +45,21 @@ locate_text_block_bbox_with_llm_sync, ) from churro_ocr.providers.specs import DEFAULT_OCR_MAX_TOKENS -from churro_ocr.providers.vllm import VLLMVisionOCRBackend -from churro_ocr.templates import DEFAULT_OCR_TEMPLATE +from churro_ocr.templates import ( + CHANDRA_OCR_2_MODEL_ID, + CHANDRA_OCR_2_OCR_TEMPLATE, + DEEPSEEK_OCR_2_MODEL_ID, + DEEPSEEK_OCR_2_OCR_PROMPT, + DEEPSEEK_OCR_2_OCR_TEMPLATE, + DEFAULT_OCR_TEMPLATE, + DOTS_MOCR_OCR_TEMPLATE, + OLMOCR_2_7B_1025_FP8_MODEL_ID, + OLMOCR_2_7B_1025_MODEL_ID, + OLMOCR_2_7B_1025_OCR_TEMPLATE, + PADDLEOCR_VL_1_5_MODEL_ID, + PADDLEOCR_VL_1_5_OCR_PROMPT, + PADDLEOCR_VL_1_5_OCR_TEMPLATE, +) def _extract_user_text_parts(messages: list[dict[str, Any]]) -> list[str]: @@ -59,6 +78,7 @@ async def test_litellm_ocr_backend_uses_transport(monkeypatch: pytest.MonkeyPatc image = Image.new("RGB", (10, 10), color="white") page = DocumentPage(page_index=0, image=image, source_index=0) prompt_logs: list[str] = [] + captured: dict[str, object] = {} class FakeLogger: def debug(self, message: str, *args: object) -> None: @@ -74,7 +94,21 @@ def _fake_prepare_messages( assert image_detail == "high" return [{"role": "user", "content": [{"type": "text", "text": "prompt"}]}] - async def _fake_complete_text(self, **_: object) -> str: # noqa: ANN001 + async def _fake_complete_text( + self: LiteLLMTransport, + *, + model: str, + messages: list[dict[str, object]], + timeout_seconds: int = 600, + output_json: bool = False, + allow_empty: bool = False, + ) -> str: + captured["transport"] = self + captured["model"] = model + captured["messages"] = messages + captured["timeout_seconds"] = timeout_seconds + captured["output_json"] = output_json + captured["allow_empty"] = allow_empty return "transcribed text" monkeypatch.setattr( @@ -95,6 +129,12 @@ async def _fake_complete_text(self, **_: object) -> str: # noqa: ANN001 assert result.text == "transcribed text" assert result.model_name == "gpt-4.1-mini" assert backend.transport.config.completion_kwargs == {"max_tokens": DEFAULT_OCR_MAX_TOKENS} + assert captured["transport"] is backend.transport + assert captured["model"] == "gpt-4.1-mini" + assert captured["messages"] == [{"role": "user", "content": [{"type": "text", "text": "prompt"}]}] + assert captured["timeout_seconds"] == 600 + assert captured["output_json"] is False + assert captured["allow_empty"] is True assert len(prompt_logs) == 1 assert "First OCR prompt payload for litellm" in prompt_logs[0] @@ -190,6 +230,31 @@ async def _fake_complete_text(self, **_: object) -> str: # noqa: ANN001 assert result.text == "transcribed text" +@pytest.mark.asyncio +async def test_litellm_ocr_backend_accepts_empty_transport_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + async def _fake_complete_text(self, **_: object) -> str: # noqa: ANN001 + return "" + + monkeypatch.setattr( + "churro_ocr._internal.litellm._prepare_messages_from_conversation", + lambda *_args, **_kwargs: [], + ) + monkeypatch.setattr("churro_ocr._internal.litellm.LiteLLMTransport.complete_text", _fake_complete_text) + + backend = build_ocr_backend( + OCRBackendSpec( + provider="litellm", + model="gpt-4.1-mini", + profile=resolve_ocr_profile(model_id="gpt-4.1-mini"), + ) + ) + result = await backend.ocr(DocumentPage.from_image(Image.new("RGB", (10, 10), color="white"))) + + assert result.text == "" + + @pytest.mark.asyncio async def test_openai_compatible_backend_reports_display_model( monkeypatch: pytest.MonkeyPatch, @@ -293,6 +358,93 @@ def __init__(self, key: str) -> None: assert calls == {"client_inits": 1, "requests": 2} +@pytest.mark.asyncio +async def test_azure_ocr_backend_retries_transient_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"client_inits": 0, "requests": 0} + sleep_calls: list[float] = [] + image = Image.new("RGB", (10, 10), color="white") + encoded = base64.b64encode(b"image-bytes").decode("ascii") + + class FakeAzureError(Exception): + def __init__(self, status_code: int, headers: dict[str, str] | None = None) -> None: + self.status_code = status_code + self.headers = headers or {} + self.response = SimpleNamespace(headers=self.headers) + super().__init__(f"Status {status_code}") + + class FakePoller: + async def result(self) -> SimpleNamespace: + return SimpleNamespace(content="azure text") + + class FakeClient: + def __init__(self, *, endpoint: str, credential: Any) -> None: + calls["client_inits"] += 1 + assert endpoint == "https://example.test" + assert credential.key == "secret" + + async def begin_analyze_document( + self, + *, + model_id: str, + body: Any, + content_type: str, + ) -> FakePoller: + calls["requests"] += 1 + assert model_id == "prebuilt-layout" + assert body.read() == b"image-bytes" + assert content_type == "application/octet-stream" + if calls["requests"] < 3: + raise FakeAzureError(503) + return FakePoller() + + class FakeAzureKeyCredential: + def __init__(self, key: str) -> None: + self.key = key + + async def _fake_sleep(delay: float) -> None: + sleep_calls.append(delay) + + azure_document_module = ModuleType("azure.ai.documentintelligence.aio") + cast("Any", azure_document_module).DocumentIntelligenceClient = FakeClient + azure_credentials_module = ModuleType("azure.core.credentials") + cast("Any", azure_credentials_module).AzureKeyCredential = FakeAzureKeyCredential + monkeypatch.setitem(sys.modules, "azure.ai.documentintelligence.aio", azure_document_module) + monkeypatch.setitem(sys.modules, "azure.core.credentials", azure_credentials_module) + monkeypatch.setattr(retry_module, "retry_sleep", _fake_sleep) + monkeypatch.setattr( + "churro_ocr.providers.ocr.image_to_base64", + lambda actual_image, format_name: ( + ( + "image/jpeg", + encoded, + ) + if actual_image.size == image.size and actual_image.mode == "RGB" and format_name == "JPEG" + else ("", "") + ), + ) + + backend = cast( + "AzureDocumentIntelligenceOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="azure", + options=AzureDocumentIntelligenceOptions( + endpoint="https://example.test", + api_key="secret", + ), + ) + ), + ) + + result = await backend.ocr(DocumentPage(page_index=0, image=image, source_index=0)) + + assert result.text == "azure text" + assert calls == {"client_inits": 1, "requests": 3} + assert sleep_calls == [1.0, 2.0] + + @pytest.mark.asyncio async def test_mistral_ocr_backend_reuses_client(monkeypatch: pytest.MonkeyPatch) -> None: calls = {"client_inits": 0, "requests": 0} @@ -301,7 +453,7 @@ async def test_mistral_ocr_backend_reuses_client(monkeypatch: pytest.MonkeyPatch class FakeOCRNamespace: async def process_async(self, *, model: str, document: dict[str, str]) -> SimpleNamespace: calls["requests"] += 1 - assert model == "mistral-ocr-latest" + assert model == "mistral-ocr-2512" assert document == { "type": "image_url", "image_url": "data:image/jpeg;base64,encoded-image", @@ -334,6 +486,7 @@ def __init__(self, *, api_key: str) -> None: build_ocr_backend( OCRBackendSpec( provider="mistral", + model="mistral-ocr-2512", options=MistralOptions(api_key="secret"), ) ), @@ -349,182 +502,661 @@ def __init__(self, *, api_key: str) -> None: @pytest.mark.asyncio -async def test_vllm_ocr_backend_reuses_engine_and_batches(monkeypatch: pytest.MonkeyPatch) -> None: - calls = {"processor_inits": 0, "engine_inits": 0} +async def test_mistral_ocr_backend_retries_transient_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"client_inits": 0, "requests": 0} + image = Image.new("RGB", (10, 10), color="white") + sleep_calls: list[float] = [] + + class FakeMistralError(Exception): + def __init__(self, status_code: int, headers: dict[str, str] | None = None) -> None: + self.status_code = status_code + self.headers = headers or {} + self.raw_response = SimpleNamespace(headers=self.headers) + super().__init__(f"Status {status_code}") + + class FakeOCRNamespace: + async def process_async(self, *, model: str, document: dict[str, str]) -> SimpleNamespace: + calls["requests"] += 1 + assert model == "mistral-ocr-2512" + assert document == { + "type": "image_url", + "image_url": "data:image/jpeg;base64,encoded-image", + } + if calls["requests"] < 3: + raise FakeMistralError(520) + return SimpleNamespace(pages=[SimpleNamespace(markdown="mistral text")]) + + class FakeMistralClient: + def __init__(self, *, api_key: str) -> None: + calls["client_inits"] += 1 + assert api_key == "secret" + self.ocr = FakeOCRNamespace() + + async def _fake_sleep(delay: float) -> None: + sleep_calls.append(delay) + + mistral_module = ModuleType("mistralai") + cast("Any", mistral_module).Mistral = FakeMistralClient + monkeypatch.setitem(sys.modules, "mistralai", mistral_module) + monkeypatch.setattr( + "churro_ocr.providers.ocr.image_to_base64", + lambda actual_image, format_name: ( + ( + "image/jpeg", + "encoded-image", + ) + if actual_image.size == image.size and actual_image.mode == "RGB" and format_name == "JPEG" + else ("", "") + ), + ) + monkeypatch.setattr(retry_module, "retry_sleep", _fake_sleep) + + backend = cast( + "MistralOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="mistral", + model="mistral-ocr-2512", + options=MistralOptions(api_key="secret"), + ) + ), + ) + page = DocumentPage(page_index=0, image=image, source_index=0) + + result = await backend.ocr(page) + + assert result.text == "mistral text" + assert calls == {"client_inits": 1, "requests": 3} + assert sleep_calls == [1.0, 2.0] + + +@pytest.mark.asyncio +async def test_mistral_ocr_backend_retries_request_timeouts( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"client_inits": 0, "requests": 0, "wait_for": 0} + image = Image.new("RGB", (10, 10), color="white") + sleep_calls: list[float] = [] + + class FakeOCRNamespace: + async def process_async(self, *, model: str, document: dict[str, str]) -> SimpleNamespace: + calls["requests"] += 1 + assert model == "mistral-ocr-2512" + assert document == { + "type": "image_url", + "image_url": "data:image/jpeg;base64,encoded-image", + } + return SimpleNamespace(pages=[SimpleNamespace(markdown="mistral text")]) + + class FakeMistralClient: + def __init__(self, *, api_key: str) -> None: + calls["client_inits"] += 1 + assert api_key == "secret" + self.ocr = FakeOCRNamespace() + + async def _fake_sleep(delay: float) -> None: + sleep_calls.append(delay) + + async def _fake_wait_for(awaitable: Any, **kwargs: float) -> Any: + calls["wait_for"] += 1 + timeout = kwargs["timeout"] + assert timeout == 60.0 + if calls["wait_for"] == 1: + close = getattr(awaitable, "close", None) + if callable(close): + close() + raise TimeoutError + return await awaitable + + mistral_module = ModuleType("mistralai") + cast("Any", mistral_module).Mistral = FakeMistralClient + monkeypatch.setitem(sys.modules, "mistralai", mistral_module) + monkeypatch.setattr( + "churro_ocr.providers.ocr.image_to_base64", + lambda actual_image, format_name: ( + ( + "image/jpeg", + "encoded-image", + ) + if actual_image.size == image.size and actual_image.mode == "RGB" and format_name == "JPEG" + else ("", "") + ), + ) + monkeypatch.setattr(retry_module, "retry_sleep", _fake_sleep) + monkeypatch.setattr("churro_ocr.providers.ocr.asyncio.wait_for", _fake_wait_for) + + backend = cast( + "MistralOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="mistral", + model="mistral-ocr-2512", + options=MistralOptions(api_key="secret"), + ) + ), + ) + page = DocumentPage(page_index=0, image=image, source_index=0) + + result = await backend.ocr(page) + + assert result.text == "mistral text" + assert calls == {"client_inits": 1, "requests": 1, "wait_for": 2} + assert sleep_calls == [1.0] + + +@pytest.mark.asyncio +async def test_mistral_ocr_backend_does_not_retry_non_transient_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"client_inits": 0, "requests": 0} + image = Image.new("RGB", (10, 10), color="white") + sleep_calls: list[float] = [] + + class FakeMistralError(Exception): + def __init__(self, status_code: int) -> None: + self.status_code = status_code + self.headers: dict[str, str] = {} + self.raw_response = SimpleNamespace(headers=self.headers) + super().__init__(f"Status {status_code}") + + class FakeOCRNamespace: + async def process_async(self, *, model: str, document: dict[str, str]) -> SimpleNamespace: + calls["requests"] += 1 + assert model == "mistral-ocr-2505" + assert document == { + "type": "image_url", + "image_url": "data:image/jpeg;base64,encoded-image", + } + raise FakeMistralError(400) + + class FakeMistralClient: + def __init__(self, *, api_key: str) -> None: + calls["client_inits"] += 1 + assert api_key == "secret" + self.ocr = FakeOCRNamespace() + + async def _fake_sleep(delay: float) -> None: + sleep_calls.append(delay) + + mistral_module = ModuleType("mistralai") + cast("Any", mistral_module).Mistral = FakeMistralClient + monkeypatch.setitem(sys.modules, "mistralai", mistral_module) + monkeypatch.setattr( + "churro_ocr.providers.ocr.image_to_base64", + lambda actual_image, format_name: ( + ( + "image/jpeg", + "encoded-image", + ) + if actual_image.size == image.size and actual_image.mode == "RGB" and format_name == "JPEG" + else ("", "") + ), + ) + monkeypatch.setattr(retry_module, "retry_sleep", _fake_sleep) + + backend = cast( + "MistralOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="mistral", + model="mistral-ocr-2505", + options=MistralOptions(api_key="secret"), + ) + ), + ) + page = DocumentPage(page_index=0, image=image, source_index=0) + + with pytest.raises(FakeMistralError, match="Status 400"): + await backend.ocr(page) + + assert calls == {"client_inits": 1, "requests": 1} + assert sleep_calls == [] + + +def test_build_ocr_backend_uses_olmocr_profile_defaults_for_openai_compatible() -> None: + backend = cast( + "OpenAICompatibleOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="openai-compatible", + model=OLMOCR_2_7B_1025_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), + ) + ), + ) + + assert backend.template == OLMOCR_2_7B_1025_OCR_TEMPLATE + assert backend.model_name == "olmOCR-2-7B-1025" + assert backend.transport.config.completion_kwargs == { + "max_tokens": 8_000, + "temperature": 0.1, + } + assert backend.image_preprocessor(Image.new("RGB", (5_000, 3_000), color="white")).size == (1_288, 772) + + +def test_build_ocr_backend_uses_chandra_profile_defaults_for_openai_compatible() -> None: + backend = cast( + "OpenAICompatibleOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="openai-compatible", + model=CHANDRA_OCR_2_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), + ) + ), + ) + + assert backend.template == CHANDRA_OCR_2_OCR_TEMPLATE + assert backend.model_name == "chandra-ocr-2" + assert backend.transport.config.completion_kwargs == { + "max_tokens": 12_384, + "temperature": 0.0, + "top_p": 0.1, + } + assert backend.image_preprocessor(Image.new("RGB", (5_000, 3_000), color="white")).size == (3_248, 1_932) + + +def test_build_ocr_backend_resolves_olmocr_fp8_profile_defaults_for_openai_compatible() -> None: + backend = cast( + "OpenAICompatibleOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="openai-compatible", + model=OLMOCR_2_7B_1025_FP8_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), + ) + ), + ) + + assert backend.template == OLMOCR_2_7B_1025_OCR_TEMPLATE + assert backend.model_name == "olmOCR-2-7B-1025-FP8" + assert backend.transport.config.completion_kwargs == { + "max_tokens": 8_000, + "temperature": 0.1, + } + + +def test_build_ocr_backend_uses_paddleocr_vl_profile_defaults_for_openai_compatible() -> None: + backend = cast( + "OpenAICompatibleOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="openai-compatible", + model=PADDLEOCR_VL_1_5_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), + ) + ), + ) + + assert backend.template == PADDLEOCR_VL_1_5_OCR_TEMPLATE + assert backend.model_name == "PaddleOCR-VL-1.5" + assert backend.transport.config.completion_kwargs == { + "max_tokens": 4_096, + "temperature": 0.0, + } + + +def test_build_ocr_backend_uses_dots_mocr_profile_defaults_for_openai_compatible() -> None: + backend = cast( + "OpenAICompatibleOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="openai-compatible", + model="rednote-hilab/dots.mocr", + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), + ) + ), + ) + + assert backend.template == DOTS_MOCR_OCR_TEMPLATE + assert backend.model_name == "dots.mocr" + assert backend.transport.config.completion_kwargs == { + "max_tokens": DEFAULT_OCR_MAX_TOKENS, + "temperature": 0.0, + } + + +def test_build_ocr_backend_uses_deepseek_ocr_2_profile_defaults_for_openai_compatible() -> None: + backend = cast( + "OpenAICompatibleOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="openai-compatible", + model=DEEPSEEK_OCR_2_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), + ) + ), + ) + + assert backend.template == DEEPSEEK_OCR_2_OCR_TEMPLATE + assert backend.model_name == "DeepSeek-OCR-2" + assert backend.transport.config.completion_kwargs == { + "max_tokens": 8_192, + "temperature": 0.0, + } + + +@pytest.mark.asyncio +async def test_openai_compatible_backend_uses_deepseek_ocr_2_prompt_and_postprocessing( + monkeypatch: pytest.MonkeyPatch, +) -> None: captured: dict[str, object] = {} - class FakeProcessor: - tokenizer = None + def _fake_prepare_messages_from_conversation( + self: LiteLLMTransport, + conversation: list[dict[str, object]], + ) -> list[dict[str, object]]: + captured["conversation"] = conversation + captured["completion_kwargs"] = dict(self.config.completion_kwargs) + return conversation - def apply_chat_template( - self, - conversation: list[dict[str, object]], - *, - add_generation_prompt: bool, - tokenize: bool, - ) -> str: - assert add_generation_prompt is True - assert tokenize is False - content = cast("list[dict[str, object]]", conversation[0]["content"]) - return f"prompt:{content[1]['text']}" - - class FakeProcessorCls: - @staticmethod - def from_pretrained(model_id: str, **kwargs: object) -> FakeProcessor: - calls["processor_inits"] += 1 - captured["processor_model_id"] = model_id - captured["processor_kwargs"] = kwargs - return FakeProcessor() - - class FakeSamplingParams: - def __init__(self, **kwargs: object) -> None: - captured["sampling_kwargs"] = kwargs - - class FakeLLM: - def __init__(self, **kwargs: object) -> None: - calls["engine_inits"] += 1 - captured["llm_kwargs"] = kwargs - - def generate( - self, - prompts: list[dict[str, object]], - sampling_params: FakeSamplingParams, - *, - use_tqdm: bool, - ) -> list[SimpleNamespace]: - del sampling_params - captured["prompts"] = prompts - captured["use_tqdm"] = use_tqdm - return [ - SimpleNamespace(outputs=[SimpleNamespace(text=f"text:{index}")]) - for index, _prompt in enumerate(prompts) - ] - - vllm_module = ModuleType("vllm") - cast("Any", vllm_module).LLM = FakeLLM - cast("Any", vllm_module).SamplingParams = FakeSamplingParams - monkeypatch.setitem(sys.modules, "vllm", vllm_module) - monkeypatch.setattr("churro_ocr.providers.vllm._load_vllm_processor_cls", lambda: FakeProcessorCls) + async def _fake_complete_text( + self: LiteLLMTransport, + *, + model: str, + messages: list[dict[str, object]], + timeout_seconds: int = 600, + output_json: bool = False, + allow_empty: bool = False, + ) -> str: + captured["model"] = model + captured["messages"] = messages + captured["timeout_seconds"] = timeout_seconds + captured["output_json"] = output_json + captured["allow_empty"] = allow_empty + return "\nFree OCR.\n<｜Assistant｜>\nDecoded text<｜end▁of▁sentence｜>" + + monkeypatch.setattr( + "churro_ocr._internal.litellm.LiteLLMTransport.prepare_messages_from_conversation", + _fake_prepare_messages_from_conversation, + ) + monkeypatch.setattr("churro_ocr._internal.litellm.LiteLLMTransport.complete_text", _fake_complete_text) backend = cast( - "VLLMVisionOCRBackend", + "OpenAICompatibleOCRBackend", build_ocr_backend( OCRBackendSpec( - provider="vllm", - model="kristaller486/dots.ocr-1.5", - options=VLLMOptions(), + provider="openai-compatible", + model=DEEPSEEK_OCR_2_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), ) ), ) - pages = [ - DocumentPage.from_image(Image.new("RGBA", (5_000, 3_000), color=(255, 255, 255, 255))), - DocumentPage.from_image(Image.new("RGB", (12, 12), color="white")), - ] + result = await backend.ocr(DocumentPage.from_image(Image.new("RGB", (10, 10), color="white"))) - first = await backend.ocr(pages[0]) - second_batch = await backend.ocr_batch(pages) - - assert first.text == "text:0" - assert [result.text for result in second_batch] == ["text:0", "text:1"] - assert calls == {"processor_inits": 1, "engine_inits": 1} - assert captured["processor_model_id"] == "kristaller486/dots.ocr-1.5" - assert captured["processor_kwargs"] == {"trust_remote_code": True} - assert captured["llm_kwargs"] == { - "model": "kristaller486/dots.ocr-1.5", - "trust_remote_code": True, - "limit_mm_per_prompt": {"image": 1}, + assert result.text == "Decoded text" + assert result.model_name == "DeepSeek-OCR-2" + assert captured["model"] == f"openai/{DEEPSEEK_OCR_2_MODEL_ID}" + assert captured["completion_kwargs"] == { + "max_tokens": 8_192, + "temperature": 0.0, } - assert captured["sampling_kwargs"] == {"max_tokens": DEFAULT_OCR_MAX_TOKENS} - assert captured["use_tqdm"] is False - prompt_batch = cast("list[dict[str, object]]", captured["prompts"]) - assert isinstance(prompt_batch, list) - assert len(prompt_batch) == 2 - assert prompt_batch[0]["prompt"] == "prompt:Extract the text content from this image." - prompt_media = cast("dict[str, object]", prompt_batch[0]["multi_modal_data"]) - prompt_image = cast("Image.Image", prompt_media["image"]) - assert isinstance(prompt_image, Image.Image) - assert prompt_image.size == (2_500, 1_500) + conversation = cast("list[dict[str, object]]", captured["conversation"]) + content = cast("list[dict[str, object]]", conversation[0]["content"]) + assert content[0]["type"] == "image" + assert content[1] == {"type": "text", "text": DEEPSEEK_OCR_2_OCR_PROMPT} + + +@pytest.mark.asyncio +async def test_openai_compatible_backend_uses_olmocr_prompt_and_plain_text_postprocessing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + def _fake_prepare_messages_from_conversation( + self: LiteLLMTransport, + conversation: list[dict[str, object]], + ) -> list[dict[str, object]]: + captured["conversation"] = conversation + captured["completion_kwargs"] = dict(self.config.completion_kwargs) + return [{"role": "user", "content": [{"type": "text", "text": "prompt"}]}] + + async def _fake_complete_text( + self: LiteLLMTransport, + *, + model: str, + messages: list[dict[str, object]], + timeout_seconds: int = 600, + output_json: bool = False, + allow_empty: bool = False, + ) -> str: + captured["model"] = model + captured["messages"] = messages + captured["timeout_seconds"] = timeout_seconds + captured["output_json"] = output_json + captured["allow_empty"] = allow_empty + captured["completion_kwargs"] = dict(self.config.completion_kwargs) + return ( + "---\n" + "primary_language: en\n" + "is_rotation_valid: true\n" + "rotation_correction: 0\n" + "is_table: true\n" + "is_diagram: false\n" + "---\n" + "# Ledger\n\n" + "" + "
YearValue
190042
\n\n" + "Paragraph with [note](https://example.test).\n" + "![Figure alt text](page_0_0_100_100.png)\n" + ) + + monkeypatch.setattr( + LiteLLMTransport, + "prepare_messages_from_conversation", + _fake_prepare_messages_from_conversation, + ) + monkeypatch.setattr(LiteLLMTransport, "complete_text", _fake_complete_text) + + backend = cast( + "OpenAICompatibleOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="openai-compatible", + model=OLMOCR_2_7B_1025_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), + ) + ), + ) + result = await backend.ocr( + DocumentPage.from_image(Image.new("RGBA", (5_000, 3_000), color=(255, 255, 255, 255))) + ) + + assert result.text == "Ledger\n\nYear | Value\n1900 | 42\n\nParagraph with note." + assert result.metadata == { + "front_matter": { + "primary_language": "en", + "is_rotation_valid": True, + "rotation_correction": 0, + "is_table": True, + "is_diagram": False, + }, + "raw_markdown": ( + "# Ledger\n\n" + "
YearValue
190042
\n\n" + "Paragraph with [note](https://example.test).\n" + "![Figure alt text](page_0_0_100_100.png)" + ), + } + assert captured["model"] == f"openai/{OLMOCR_2_7B_1025_MODEL_ID}" + assert captured["messages"] == [{"role": "user", "content": [{"type": "text", "text": "prompt"}]}] + assert captured["timeout_seconds"] == 600 + assert captured["output_json"] is False + assert captured["allow_empty"] is True + assert captured["completion_kwargs"] == { + "max_tokens": 8_000, + "temperature": 0.1, + } + conversation = cast("list[dict[str, object]]", captured["conversation"]) + assert conversation[0]["role"] == "user" + user_content = cast("list[dict[str, object]]", conversation[0]["content"]) + assert user_content[0] == {"type": "text", "text": OLMOCR_V4_YAML_PROMPT} + assert user_content[1]["type"] == "image" + prompt_image = cast("Image.Image", user_content[1]["image"]) + assert prompt_image.size == (1_288, 772) assert prompt_image.mode == "RGB" -def test_vllm_backend_defaults_are_public() -> None: +@pytest.mark.asyncio +async def test_openai_compatible_backend_uses_chandra_prompt_and_plain_text_postprocessing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} + + def _fake_prepare_messages_from_conversation( + self: LiteLLMTransport, + conversation: list[dict[str, object]], + ) -> list[dict[str, object]]: + captured["conversation"] = conversation + captured["completion_kwargs"] = dict(self.config.completion_kwargs) + return [{"role": "user", "content": [{"type": "text", "text": "prompt"}]}] + + async def _fake_complete_text( + self: LiteLLMTransport, + *, + model: str, + messages: list[dict[str, object]], + timeout_seconds: int = 600, + output_json: bool = False, + allow_empty: bool = False, + ) -> str: + captured["model"] = model + captured["messages"] = messages + captured["timeout_seconds"] = timeout_seconds + captured["output_json"] = output_json + captured["allow_empty"] = allow_empty + captured["completion_kwargs"] = dict(self.config.completion_kwargs) + return ( + '

Ledger

\n' + '

Paragraph with note.

\n' + '
' + "
YearValue
190042
" + ) + + monkeypatch.setattr( + LiteLLMTransport, + "prepare_messages_from_conversation", + _fake_prepare_messages_from_conversation, + ) + monkeypatch.setattr(LiteLLMTransport, "complete_text", _fake_complete_text) + backend = cast( - "VLLMVisionOCRBackend", + "OpenAICompatibleOCRBackend", build_ocr_backend( OCRBackendSpec( - provider="vllm", - model="kristaller486/dots.ocr-1.5", - options=VLLMOptions(), + provider="openai-compatible", + model=CHANDRA_OCR_2_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), ) ), ) + result = await backend.ocr( + DocumentPage.from_image(Image.new("RGBA", (5_000, 3_000), color=(255, 255, 255, 255))) + ) - assert backend.model_id == "kristaller486/dots.ocr-1.5" - assert backend.model_name == "dots.ocr-1.5" - assert backend.provider_name == "vllm" - assert backend.processor_kwargs == {} + assert result.text == "Ledger\n\nParagraph with note.\n\nYear | Value\n1900 | 42" + assert result.metadata == { + "raw_html": ( + '

Ledger

\n' + '

Paragraph with note.

\n' + '
' + "
YearValue
190042
" + ), + } + assert captured["model"] == f"openai/{CHANDRA_OCR_2_MODEL_ID}" + assert captured["messages"] == [{"role": "user", "content": [{"type": "text", "text": "prompt"}]}] + assert captured["timeout_seconds"] == 600 + assert captured["output_json"] is False + assert captured["allow_empty"] is True + assert captured["completion_kwargs"] == { + "max_tokens": 12_384, + "temperature": 0.0, + "top_p": 0.1, + } + conversation = cast("list[dict[str, object]]", captured["conversation"]) + assert conversation[0]["role"] == "user" + user_content = cast("list[dict[str, object]]", conversation[0]["content"]) + assert user_content[0]["type"] == "image" + prompt_image = cast("Image.Image", user_content[0]["image"]) + assert prompt_image.size == (3_248, 1_932) + assert prompt_image.mode == "RGB" + assert user_content[1] == {"type": "text", "text": CHANDRA_OCR_LAYOUT_PROMPT} @pytest.mark.asyncio -async def test_vllm_ocr_backend_strips_default_output_tags(monkeypatch: pytest.MonkeyPatch) -> None: - class FakeProcessor: - tokenizer = None +async def test_openai_compatible_backend_uses_paddleocr_vl_prompt_and_strips_prompt_echo( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, object] = {} - def apply_chat_template( - self, - conversation: list[dict[str, object]], - *, - add_generation_prompt: bool, - tokenize: bool, - ) -> str: - del conversation, add_generation_prompt, tokenize - return "prompt" - - class FakeProcessorCls: - @staticmethod - def from_pretrained(model_id: str, **kwargs: object) -> FakeProcessor: - del model_id, kwargs - return FakeProcessor() - - class FakeSamplingParams: - def __init__(self, **kwargs: object) -> None: - del kwargs - - class FakeLLM: - def __init__(self, **kwargs: object) -> None: - del kwargs - - def generate( - self, - prompts: list[dict[str, object]], - sampling_params: FakeSamplingParams, - *, - use_tqdm: bool, - ) -> list[SimpleNamespace]: - del prompts, sampling_params, use_tqdm - return [ - SimpleNamespace( - outputs=[ - SimpleNamespace( - text=(f"<{DEFAULT_OCR_OUTPUT_TAG}>\npage text\n") - ) - ] - ) - ] + def _fake_prepare_messages_from_conversation( + self: LiteLLMTransport, + conversation: list[dict[str, object]], + ) -> list[dict[str, object]]: + captured["conversation"] = conversation + captured["completion_kwargs"] = dict(self.config.completion_kwargs) + return [{"role": "user", "content": [{"type": "text", "text": "prompt"}]}] + + async def _fake_complete_text( + self: LiteLLMTransport, + *, + model: str, + messages: list[dict[str, object]], + timeout_seconds: int = 600, + output_json: bool = False, + allow_empty: bool = False, + ) -> str: + captured["model"] = model + captured["messages"] = messages + captured["timeout_seconds"] = timeout_seconds + captured["output_json"] = output_json + captured["allow_empty"] = allow_empty + captured["completion_kwargs"] = dict(self.config.completion_kwargs) + return "OCR:\nassistant\npaddle transcription" - vllm_module = ModuleType("vllm") - cast("Any", vllm_module).LLM = FakeLLM - cast("Any", vllm_module).SamplingParams = FakeSamplingParams - monkeypatch.setitem(sys.modules, "vllm", vllm_module) - monkeypatch.setattr("churro_ocr.providers.vllm._load_vllm_processor_cls", lambda: FakeProcessorCls) + monkeypatch.setattr( + LiteLLMTransport, + "prepare_messages_from_conversation", + _fake_prepare_messages_from_conversation, + ) + monkeypatch.setattr(LiteLLMTransport, "complete_text", _fake_complete_text) backend = cast( - "VLLMVisionOCRBackend", - build_ocr_backend(OCRBackendSpec(provider="vllm", model="example/model")), + "OpenAICompatibleOCRBackend", + build_ocr_backend( + OCRBackendSpec( + provider="openai-compatible", + model=PADDLEOCR_VL_1_5_MODEL_ID, + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), + ) + ), + ) + result = await backend.ocr( + DocumentPage.from_image(Image.new("RGBA", (5_000, 3_000), color=(255, 255, 255, 255))) ) - result = await backend.ocr(DocumentPage.from_image(Image.new("RGB", (10, 10), color="white"))) - assert result.text == "page text" + assert result.text == "paddle transcription" + assert captured["model"] == f"openai/{PADDLEOCR_VL_1_5_MODEL_ID}" + assert captured["messages"] == [{"role": "user", "content": [{"type": "text", "text": "prompt"}]}] + assert captured["timeout_seconds"] == 600 + assert captured["output_json"] is False + assert captured["allow_empty"] is True + assert captured["completion_kwargs"] == { + "max_tokens": 4_096, + "temperature": 0.0, + } + conversation = cast("list[dict[str, object]]", captured["conversation"]) + assert conversation[0]["role"] == "user" + user_content = cast("list[dict[str, object]]", conversation[0]["content"]) + assert user_content[0]["type"] == "image" + prompt_image = cast("Image.Image", user_content[0]["image"]) + assert prompt_image.size == (2_500, 1_500) + assert prompt_image.mode == "RGB" + assert user_content[1] == {"type": "text", "text": PADDLEOCR_VL_1_5_OCR_PROMPT} @pytest.mark.asyncio @@ -1010,6 +1642,83 @@ def __init__(self, key: str) -> None: assert calls["closes"] == 1 +@pytest.mark.asyncio +async def test_azure_page_detector_retries_transient_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = {"client_inits": 0, "requests": 0, "closes": 0} + sleep_calls: list[float] = [] + + class FakeAzureError(Exception): + def __init__(self, status_code: int, headers: dict[str, str] | None = None) -> None: + self.status_code = status_code + self.headers = headers or {} + self.response = SimpleNamespace(headers=self.headers) + super().__init__(f"Status {status_code}") + + class FakePoller: + async def result(self) -> SimpleNamespace: + return SimpleNamespace( + pages=[ + SimpleNamespace( + polygon=[0, 0, 100, 0, 100, 200, 0, 200, 0, 0], + width=100, + height=200, + page_number=1, + unit="pixel", + angle=None, + ) + ] + ) + + class FakeClient: + def __init__(self, *, endpoint: str, credential: Any) -> None: + calls["client_inits"] += 1 + assert endpoint == "https://example.test" + assert credential.key == "secret" + + async def begin_analyze_document( + self, + *, + model_id: str, + body: Any, + content_type: str, + ) -> FakePoller: + calls["requests"] += 1 + assert model_id == "prebuilt-layout" + assert body.read() + assert content_type == "application/octet-stream" + if calls["requests"] == 1: + raise FakeAzureError(429, headers={"retry-after": "4"}) + return FakePoller() + + async def close(self) -> None: + calls["closes"] += 1 + + class FakeAzureKeyCredential: + def __init__(self, key: str) -> None: + self.key = key + + async def _fake_sleep(delay: float) -> None: + sleep_calls.append(delay) + + azure_document_module = ModuleType("azure.ai.documentintelligence.aio") + cast(Any, azure_document_module).DocumentIntelligenceClient = FakeClient + azure_credentials_module = ModuleType("azure.core.credentials") + cast(Any, azure_credentials_module).AzureKeyCredential = FakeAzureKeyCredential + monkeypatch.setitem(sys.modules, "azure.ai.documentintelligence.aio", azure_document_module) + monkeypatch.setitem(sys.modules, "azure.core.credentials", azure_credentials_module) + monkeypatch.setattr(retry_module, "retry_sleep", _fake_sleep) + + detector = AzurePageDetector(endpoint="https://example.test", api_key="secret") + candidates = await detector.detect(Image.new("RGB", (100, 200), color="white")) + + assert len(candidates) == 1 + assert candidates[0].bbox == (0.0, 0.0, 100.0, 200.0) + assert calls == {"client_inits": 1, "requests": 2, "closes": 1} + assert sleep_calls == [4.0] + + def test_azure_page_detector_type_is_public() -> None: detector = AzurePageDetector(endpoint="https://example.test", api_key="secret") assert detector.model_id == "prebuilt-layout" @@ -1017,11 +1726,12 @@ def test_azure_page_detector_type_is_public() -> None: def test_build_ocr_backend_resolves_profile_defaults() -> None: backend = cast( - "VLLMVisionOCRBackend", + "OpenAICompatibleOCRBackend", build_ocr_backend( OCRBackendSpec( - provider="vllm", + provider="openai-compatible", model="stanford-oval/churro-3B", + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), ) ), ) @@ -1032,19 +1742,19 @@ def test_build_ocr_backend_resolves_profile_defaults() -> None: def test_build_ocr_backend_uses_generic_defaults_for_qwen_model() -> None: backend = cast( - "VLLMVisionOCRBackend", + "OpenAICompatibleOCRBackend", build_ocr_backend( OCRBackendSpec( - provider="vllm", + provider="openai-compatible", model="Qwen/Qwen3.5-0.8B", + transport=LiteLLMTransportConfig(api_base="http://127.0.0.1:8000/v1"), ) ), ) assert backend.model_name == "Qwen/Qwen3.5-0.8B" assert backend.template == DEFAULT_OCR_TEMPLATE - assert backend.llm_kwargs == {} - assert backend.sampling_kwargs == {"max_tokens": DEFAULT_OCR_MAX_TOKENS} + assert backend.transport.config.completion_kwargs == {"max_tokens": DEFAULT_OCR_MAX_TOKENS} def test_build_ocr_backend_merges_hf_overrides_with_profile_defaults() -> None: diff --git a/tests/test_tooling_benchmark.py b/tests/test_tooling_benchmark.py index e672903..98bd750 100644 --- a/tests/test_tooling_benchmark.py +++ b/tests/test_tooling_benchmark.py @@ -8,11 +8,17 @@ from datasets import Dataset from PIL import Image +from churro_ocr.ocr import OCRResult from churro_ocr.providers.hf import HuggingFaceVisionOCRBackend from churro_ocr.providers.ocr import LiteLLMVisionOCRBackend from churro_ocr.providers.specs import DEFAULT_OCR_MAX_TOKENS -from churro_ocr.providers.vllm import VLLMVisionOCRBackend -from churro_ocr.templates import CHURRO_3B_XML_TEMPLATE +from churro_ocr.templates import ( + CHURRO_3B_XML_TEMPLATE, + DEEPSEEK_OCR_2_OCR_TEMPLATE, + DOTS_MOCR_OCR_TEMPLATE, + DOTS_OCR_1_5_OCR_TEMPLATE, + PADDLEOCR_VL_1_5_OCR_TEMPLATE, +) from tooling.benchmarking import benchmark from tooling.evaluation.types import BenchmarkDatasetExample @@ -104,23 +110,39 @@ def test_validate_options_requires_model_for_hf() -> None: assert benchmark._validate_options(options) == 1 -def test_validate_options_requires_model_for_vllm() -> None: - options = benchmark.BenchmarkOptions( - backend="vllm", +def test_validate_options_requires_pinned_model_for_mistral() -> None: + missing_model = benchmark.BenchmarkOptions( + backend="mistral", model=None, dataset_split="dev", + api_key="secret", ) - assert benchmark._validate_options(options) == 1 + alias_model = benchmark.BenchmarkOptions( + backend="mistral", + model="mistral-ocr-latest", + dataset_split="dev", + api_key="secret", + ) + pinned_model = benchmark.BenchmarkOptions( + backend="mistral", + model="mistral-ocr-2512", + dataset_split="dev", + api_key="secret", + ) + + assert benchmark._validate_options(missing_model) == 1 + assert benchmark._validate_options(alias_model) == 1 + assert benchmark._validate_options(pinned_model) == 0 -def test_validate_options_rejects_invalid_vllm_gpu_memory_utilization() -> None: +def test_validate_options_allows_openai_compatible_without_api_key() -> None: options = benchmark.BenchmarkOptions( - backend="vllm", - model="Qwen/Qwen3.5-0.8B", + backend="openai-compatible", dataset_split="dev", - vllm_gpu_memory_utilization=1.5, + model="local-model", + base_url="http://127.0.0.1:8000/v1", ) - assert benchmark._validate_options(options) == 1 + assert benchmark._validate_options(options) == 0 def test_validate_options_rejects_invalid_split() -> None: @@ -155,24 +177,18 @@ def test_parse_args_accepts_subset_filters() -> None: assert options.document_type == "print" -def test_parse_args_accepts_vllm_resource_overrides() -> None: - options = benchmark.parse_args( - [ - "--backend", - "vllm", - "--dataset-split", - "dev", - "--model", - "Qwen/Qwen3.5-0.8B", - "--vllm-gpu-memory-utilization", - "0.25", - "--vllm-cpu-offload-gb", - "8", - ] - ) - - assert options.vllm_gpu_memory_utilization == pytest.approx(0.25) - assert options.vllm_cpu_offload_gb == pytest.approx(8.0) +def test_parse_args_rejects_unsupported_backend() -> None: + with pytest.raises(SystemExit): + benchmark.parse_args( + [ + "--backend", + "unsupported", + "--dataset-split", + "dev", + "--model", + "Qwen/Qwen3.5-0.8B", + ] + ) def test_build_ocr_backend_enables_disk_cache_for_litellm( @@ -198,6 +214,24 @@ def test_build_ocr_backend_enables_disk_cache_for_litellm( assert backend.transport.config.completion_kwargs == {"max_tokens": DEFAULT_OCR_MAX_TOKENS} +def test_build_ocr_backend_allows_openai_compatible_without_api_key() -> None: + backend = cast( + "LiteLLMVisionOCRBackend", + benchmark._build_ocr_backend( + benchmark.BenchmarkOptions( + backend="openai-compatible", + dataset_split="dev", + model="local-model", + base_url="http://127.0.0.1:8000/v1", + ) + ), + ) + + assert backend.provider_name == "openai-compatible" + assert backend.transport.config.api_base == "http://127.0.0.1:8000/v1" + assert backend.transport.config.api_key is None + + def test_build_ocr_backend_uses_dots_preset_for_hf() -> None: backend = cast( "HuggingFaceVisionOCRBackend", @@ -214,38 +248,178 @@ def test_build_ocr_backend_uses_dots_preset_for_hf() -> None: assert backend.processor_kwargs == {} assert backend.trust_remote_code is True assert backend.model_kwargs["dtype"] in {"auto", "float32"} - if backend.model_kwargs["dtype"] == "auto": + if backend.model_kwargs["dtype"] == "auto" and "device_map" in backend.model_kwargs: assert backend.model_kwargs["device_map"] == "auto" - assert "max_memory" in backend.model_kwargs + if "max_memory" in backend.model_kwargs: + assert backend.model_kwargs["max_memory"] assert backend.generation_kwargs == {"max_new_tokens": DEFAULT_OCR_MAX_TOKENS} -def test_build_ocr_backend_uses_dots_preset_for_vllm() -> None: +def test_build_ocr_backend_uses_dots_preset_for_openai_compatible() -> None: backend = cast( - "VLLMVisionOCRBackend", + "LiteLLMVisionOCRBackend", benchmark._build_ocr_backend( benchmark.BenchmarkOptions( - backend="vllm", + backend="openai-compatible", dataset_split="dev", model="kristaller486/dots.ocr-1.5", + base_url="http://127.0.0.1:8000/v1", ) ), ) + assert backend.provider_name == "openai-compatible" assert backend.model_name == "dots.ocr-1.5" + assert backend.template == DOTS_OCR_1_5_OCR_TEMPLATE + assert backend.transport.config.completion_kwargs == { + "max_tokens": 2_048, + "temperature": 0.0, + } + + +def test_build_ocr_backend_uses_dots_mocr_preset_for_hf() -> None: + backend = cast( + "HuggingFaceVisionOCRBackend", + benchmark._build_ocr_backend( + benchmark.BenchmarkOptions( + backend="hf", + dataset_split="dev", + model="rednote-hilab/dots.mocr", + ) + ), + ) + + assert backend.model_name == "dots.mocr" assert backend.processor_kwargs == {} assert backend.trust_remote_code is True - assert backend.sampling_kwargs == {"max_tokens": DEFAULT_OCR_MAX_TOKENS} + assert backend.model_kwargs["dtype"] in {"auto", "float32"} + if backend.model_kwargs["dtype"] == "auto" and "device_map" in backend.model_kwargs: + assert backend.model_kwargs["device_map"] == "auto" + if "max_memory" in backend.model_kwargs: + assert backend.model_kwargs["max_memory"] + assert backend.generation_kwargs == {"max_new_tokens": DEFAULT_OCR_MAX_TOKENS} -def test_build_ocr_backend_uses_churro_preset_template_for_vllm() -> None: +def test_build_ocr_backend_uses_dots_mocr_preset_for_openai_compatible() -> None: backend = cast( - "VLLMVisionOCRBackend", + "LiteLLMVisionOCRBackend", benchmark._build_ocr_backend( benchmark.BenchmarkOptions( - backend="vllm", + backend="openai-compatible", + dataset_split="dev", + model="rednote-hilab/dots.mocr", + base_url="http://127.0.0.1:8000/v1", + ) + ), + ) + + assert backend.provider_name == "openai-compatible" + assert backend.model_name == "dots.mocr" + assert backend.template == DOTS_MOCR_OCR_TEMPLATE + assert backend.transport.config.completion_kwargs == { + "max_tokens": DEFAULT_OCR_MAX_TOKENS, + "temperature": 0.0, + } + + +def test_build_ocr_backend_uses_deepseek_ocr_2_preset_for_hf() -> None: + backend = cast( + "HuggingFaceVisionOCRBackend", + benchmark._build_ocr_backend( + benchmark.BenchmarkOptions( + backend="hf", + dataset_split="dev", + model="deepseek-ai/DeepSeek-OCR-2", + ) + ), + ) + + assert backend.model_name == "DeepSeek-OCR-2" + assert backend.processor_kwargs == {} + assert backend.trust_remote_code is True + assert backend.model_kwargs == { + "device_map": "auto", + "torch_dtype": "auto", + "use_safetensors": True, + } + assert backend.generation_kwargs == {"max_new_tokens": 8_192} + + +def test_build_ocr_backend_uses_deepseek_ocr_2_preset_for_openai_compatible() -> None: + backend = cast( + "LiteLLMVisionOCRBackend", + benchmark._build_ocr_backend( + benchmark.BenchmarkOptions( + backend="openai-compatible", + dataset_split="dev", + model="deepseek-ai/DeepSeek-OCR-2", + base_url="http://127.0.0.1:8000/v1", + ) + ), + ) + + assert backend.provider_name == "openai-compatible" + assert backend.model_name == "DeepSeek-OCR-2" + assert backend.template == DEEPSEEK_OCR_2_OCR_TEMPLATE + assert backend.transport.config.completion_kwargs == { + "max_tokens": 8_192, + "temperature": 0.0, + } + + +def test_build_ocr_backend_uses_paddleocr_vl_preset_for_hf() -> None: + backend = cast( + "HuggingFaceVisionOCRBackend", + benchmark._build_ocr_backend( + benchmark.BenchmarkOptions( + backend="hf", + dataset_split="dev", + model="PaddlePaddle/PaddleOCR-VL-1.5", + ) + ), + ) + + assert backend.model_name == "PaddleOCR-VL-1.5" + assert backend.processor_kwargs == {} + assert backend.trust_remote_code is False + assert backend.model_kwargs == {"device_map": "auto", "torch_dtype": "auto"} + assert backend.generation_kwargs == { + "max_new_tokens": 4_096, + "do_sample": False, + } + + +def test_build_ocr_backend_uses_paddleocr_vl_preset_for_openai_compatible() -> None: + backend = cast( + "LiteLLMVisionOCRBackend", + benchmark._build_ocr_backend( + benchmark.BenchmarkOptions( + backend="openai-compatible", + dataset_split="dev", + model="PaddlePaddle/PaddleOCR-VL-1.5", + base_url="http://127.0.0.1:8000/v1", + ) + ), + ) + + assert backend.provider_name == "openai-compatible" + assert backend.model_name == "PaddleOCR-VL-1.5" + assert backend.template == PADDLEOCR_VL_1_5_OCR_TEMPLATE + assert backend.transport.config.completion_kwargs == { + "max_tokens": 4_096, + "temperature": 0.0, + } + + +def test_build_ocr_backend_uses_churro_preset_template_for_openai_compatible() -> None: + backend = cast( + "LiteLLMVisionOCRBackend", + benchmark._build_ocr_backend( + benchmark.BenchmarkOptions( + backend="openai-compatible", dataset_split="dev", model="stanford-oval/churro-3B", + base_url="http://127.0.0.1:8000/v1", ) ), ) @@ -254,29 +428,24 @@ def test_build_ocr_backend_uses_churro_preset_template_for_vllm() -> None: assert backend.model_name == "churro-3B" -def test_build_ocr_backend_uses_generic_qwen_model_name_for_vllm() -> None: +def test_build_ocr_backend_uses_generic_qwen_model_name_for_openai_compatible() -> None: backend = cast( - "VLLMVisionOCRBackend", + "LiteLLMVisionOCRBackend", benchmark._build_ocr_backend( benchmark.BenchmarkOptions( - backend="vllm", + backend="openai-compatible", dataset_split="dev", model="Qwen/Qwen3.5-0.8B", - vllm_gpu_memory_utilization=0.25, - vllm_cpu_offload_gb=8.0, + base_url="http://127.0.0.1:8000/v1", ) ), ) assert backend.model_name == "Qwen/Qwen3.5-0.8B" - assert backend.llm_kwargs == { - "gpu_memory_utilization": 0.25, - "cpu_offload_gb": 8.0, - } - assert backend.sampling_kwargs == {"max_tokens": DEFAULT_OCR_MAX_TOKENS} + assert backend.transport.config.completion_kwargs == {"max_tokens": DEFAULT_OCR_MAX_TOKENS} -def test_build_ocr_backend_aligns_hf_and_vllm_templates_for_generic_models() -> None: +def test_build_ocr_backend_aligns_hf_and_openai_compatible_templates_for_generic_models() -> None: hf_backend = cast( "HuggingFaceVisionOCRBackend", benchmark._build_ocr_backend( @@ -287,18 +456,19 @@ def test_build_ocr_backend_aligns_hf_and_vllm_templates_for_generic_models() -> ) ), ) - vllm_backend = cast( - "VLLMVisionOCRBackend", + openai_backend = cast( + "LiteLLMVisionOCRBackend", benchmark._build_ocr_backend( benchmark.BenchmarkOptions( - backend="vllm", + backend="openai-compatible", dataset_split="dev", model="example/model", + base_url="http://127.0.0.1:8000/v1", ) ), ) - assert hf_backend.template == vllm_backend.template + assert hf_backend.template == openai_backend.template @pytest.mark.asyncio @@ -334,7 +504,9 @@ async def fake_predict(ds, options, *, total_pages): # noqa: ANN001 assert selected[0]["example_id"] == "1" assert options.max_concurrency == 2 assert total_pages is None - return [benchmark._build_evaluation_example(selected[0])], ["prediction"] + return [ + benchmark._build_evaluation_example(selected[0]) + ], [{"text": "prediction", "metadata": {"raw_html": "

prediction

"}}] monkeypatch.setattr(benchmark, "_predict_texts", fake_predict) @@ -366,7 +538,7 @@ def fake_compute_metrics(ds, predictions, output_prefix, elapsed_time): # noqa: assert result == 0 assert captured["dataset"] == [benchmark._build_evaluation_example(dataset[1])] - assert captured["predictions"] == ["prediction"] + assert captured["predictions"] == [{"text": "prediction", "metadata": {"raw_html": "

prediction

"}}] assert captured["output_prefix"] == str(tmp_path / "outputs") assert captured["elapsed_time"] == pytest.approx(3.5) @@ -507,10 +679,20 @@ async def test_predict_texts_updates_progress_and_preserves_order(monkeypatch) - ] class FakeProgressBar: - def __init__(self, *, total: int | None, desc: str, unit: str) -> None: + def __init__( + self, + *, + total: int | None, + desc: str, + unit: str, + mininterval: float, + smoothing: float, + ) -> None: self.total = total self.desc = desc self.unit = unit + self.mininterval = mininterval + self.smoothing = smoothing self.updates: list[int] = [] self.postfixes: list[dict[str, int]] = [] self.refresh_count = 0 @@ -538,19 +720,33 @@ def refresh(self) -> None: progress_bars: list[FakeProgressBar] = [] - def fake_tqdm(*, total: int | None, desc: str, unit: str) -> FakeProgressBar: - progress_bar = FakeProgressBar(total=total, desc=desc, unit=unit) + def fake_tqdm( + *, + total: int | None, + desc: str, + unit: str, + mininterval: float, + smoothing: float, + ) -> FakeProgressBar: + progress_bar = FakeProgressBar( + total=total, + desc=desc, + unit=unit, + mininterval=mininterval, + smoothing=smoothing, + ) progress_bars.append(progress_bar) return progress_bar - class FakeOCRResult: - def __init__(self, text: str) -> None: - self.text = text - class FakeOCRBackend: async def ocr(self, page): # noqa: ANN001 await asyncio.sleep(page.width / 1000) - return FakeOCRResult(text=f"page-{page.width}") + return OCRResult( + text=f"page-{page.width}", + provider_name="fake", + model_name="fake-model", + metadata={"page_width": page.width}, + ) monkeypatch.setattr(benchmark, "tqdm", fake_tqdm) monkeypatch.setattr(benchmark, "_build_ocr_backend", lambda _: FakeOCRBackend()) @@ -569,12 +765,18 @@ async def ocr(self, page): # noqa: ANN001 total_pages=3, ) - assert predictions == ["page-3", "page-1", "page-2"] + assert predictions == [ + {"text": "page-3", "metadata": {"page_width": 3}}, + {"text": "page-1", "metadata": {"page_width": 1}}, + {"text": "page-2", "metadata": {"page_width": 2}}, + ] assert evaluation_examples == [benchmark._build_evaluation_example(example) for example in dataset] assert len(progress_bars) == 1 assert progress_bars[0].total == 3 assert progress_bars[0].desc == "OCR" assert progress_bars[0].unit == "page" + assert progress_bars[0].mininterval == benchmark.PROGRESS_BAR_MININTERVAL_SECONDS + assert progress_bars[0].smoothing == benchmark.PROGRESS_BAR_SMOOTHING assert progress_bars[0].updates == [1, 1, 1] assert progress_bars[0].postfixes[-1] == { "submitted": 3, @@ -607,14 +809,18 @@ async def test_predict_texts_uses_batch_backend_with_max_concurrency_as_batch_si ] captured_batch_sizes: list[int] = [] - class FakeOCRResult: - def __init__(self, text: str) -> None: - self.text = text - class FakeBatchBackend: async def ocr_batch(self, pages): # noqa: ANN001 captured_batch_sizes.append(len(pages)) - return [FakeOCRResult(text=f"page-{page.width}") for page in pages] + return [ + OCRResult( + text=f"page-{page.width}", + provider_name="fake", + model_name="fake-model", + metadata={"page_width": page.width}, + ) + for page in pages + ] monkeypatch.setattr(benchmark, "_build_ocr_backend", lambda _: FakeBatchBackend()) @@ -632,7 +838,11 @@ async def ocr_batch(self, pages): # noqa: ANN001 ) assert captured_batch_sizes == [2, 1] - assert predictions == ["page-3", "page-1", "page-2"] + assert predictions == [ + {"text": "page-3", "metadata": {"page_width": 3}}, + {"text": "page-1", "metadata": {"page_width": 1}}, + {"text": "page-2", "metadata": {"page_width": 2}}, + ] assert evaluation_examples == [benchmark._build_evaluation_example(example) for example in dataset] @@ -648,13 +858,17 @@ class FakeLogger: def info(self, message: str, *args: object) -> None: logged_messages.append(message % args if args else message) - class FakeOCRResult: - def __init__(self, text: str) -> None: - self.text = text - class FakeBatchBackend: async def ocr_batch(self, pages): # noqa: ANN001 - return [FakeOCRResult(text=f"page-{page.width}") for page in pages] + return [ + OCRResult( + text=f"page-{page.width}", + provider_name="fake", + model_name="fake-model", + metadata={"page_width": page.width}, + ) + for page in pages + ] monkeypatch.setattr(benchmark, "logger", FakeLogger()) monkeypatch.setattr(benchmark, "_build_ocr_backend", lambda _: FakeBatchBackend()) @@ -672,49 +886,15 @@ async def ocr_batch(self, pages): # noqa: ANN001 total_pages=2, ) - assert predictions == ["page-3", "page-1"] + assert predictions == [ + {"text": "page-3", "metadata": {"page_width": 3}}, + {"text": "page-1", "metadata": {"page_width": 1}}, + ] assert logged_messages == [ "First benchmark OCR output for backend=hf model=kristaller486/dots.ocr-1.5:\npage-3" ] -@pytest.mark.asyncio -async def test_predict_texts_uses_max_concurrency_for_vllm_batch_backend(monkeypatch) -> None: - dataset = [ - _benchmark_example(str(index), size=(index + 1, index + 1), transcription=f"text-{index}") - for index in range(10) - ] - captured_batch_sizes: list[int] = [] - - class FakeOCRResult: - def __init__(self, text: str) -> None: - self.text = text - - class FakeBatchBackend: - async def ocr_batch(self, pages): # noqa: ANN001 - captured_batch_sizes.append(len(pages)) - return [FakeOCRResult(text=f"page-{page.width}") for page in pages] - - monkeypatch.setattr(benchmark, "_build_ocr_backend", lambda _: FakeBatchBackend()) - - options = benchmark.BenchmarkOptions( - backend="vllm", - dataset_split="dev", - model="Qwen/Qwen3.5-0.8B", - max_concurrency=2, - ) - - evaluation_examples, predictions = await benchmark._predict_texts( - dataset, - options, - total_pages=10, - ) - - assert captured_batch_sizes == [2, 2, 2, 2, 2] - assert predictions == [f"page-{index + 1}" for index in range(10)] - assert evaluation_examples == [benchmark._build_evaluation_example(example) for example in dataset] - - @pytest.mark.asyncio async def test_predict_texts_logs_first_submitted_output_once_for_non_batch_backend(monkeypatch) -> None: dataset: list[BenchmarkDatasetExample] = [ @@ -728,14 +908,15 @@ class FakeLogger: def info(self, message: str, *args: object) -> None: logged_messages.append(message % args if args else message) - class FakeOCRResult: - def __init__(self, text: str) -> None: - self.text = text - class FakeOCRBackend: async def ocr(self, page): # noqa: ANN001 await asyncio.sleep(page.width / 1000) - return FakeOCRResult(text=f"page-{page.width}") + return OCRResult( + text=f"page-{page.width}", + provider_name="fake", + model_name="fake-model", + metadata={"page_width": page.width}, + ) monkeypatch.setattr(benchmark, "logger", FakeLogger()) monkeypatch.setattr(benchmark, "_build_ocr_backend", lambda _: FakeOCRBackend()) @@ -754,5 +935,9 @@ async def ocr(self, page): # noqa: ANN001 total_pages=3, ) - assert predictions == ["page-3", "page-1", "page-2"] + assert predictions == [ + {"text": "page-3", "metadata": {"page_width": 3}}, + {"text": "page-1", "metadata": {"page_width": 1}}, + {"text": "page-2", "metadata": {"page_width": 2}}, + ] assert logged_messages == ["First benchmark OCR output for backend=azure model=:\npage-3"] diff --git a/tests/test_tooling_evaluate_page.py b/tests/test_tooling_evaluate_page.py index 1963af4..d0c0293 100644 --- a/tests/test_tooling_evaluate_page.py +++ b/tests/test_tooling_evaluate_page.py @@ -97,3 +97,21 @@ def test_calculate_metrics_strips_output_tags_before_normalization( result = evaluate_page_module.calculate_metrics((example, predicted_text)) assert result["normalized_predicted_text"] == expected + + +def test_calculate_metrics_from_text_lazily_initializes_bleu_metric(monkeypatch) -> None: + init_calls = 0 + fake_metric = SimpleNamespace(compute=lambda *_args, **_kwargs: {"bleu": 0.25}) + + def fake_initialize_metrics() -> None: + nonlocal init_calls + init_calls += 1 + monkeypatch.setattr(evaluate_page_module, "bleu_metric", fake_metric) + + monkeypatch.setattr(evaluate_page_module, "bleu_metric", None) + monkeypatch.setattr(evaluate_page_module, "initialize_metrics", fake_initialize_metrics) + + result = evaluate_page_module.calculate_metrics_from_text("pred", "gold", "English", "Latin") + + assert init_calls == 1 + assert result["bleu"] == 0.25 diff --git a/tests/test_tooling_metrics.py b/tests/test_tooling_metrics.py index f328c13..0493f52 100644 --- a/tests/test_tooling_metrics.py +++ b/tests/test_tooling_metrics.py @@ -5,7 +5,7 @@ from typing import cast from tooling.evaluation import metrics -from tooling.evaluation.types import EvaluationExample, PageEvaluationResult +from tooling.evaluation.types import BenchmarkPrediction, EvaluationExample, PageEvaluationResult def test_calculate_language_and_type_metrics_handles_missing_categories() -> None: @@ -55,7 +55,7 @@ def test_compute_metrics_writes_expected_outputs( "document_type": "print", }, ] - predicted_texts = [""] + predictions: list[BenchmarkPrediction] = [{"text": "", "metadata": {"raw_html": "

"}}] def fake_batch_evaluate(ds, preds): # noqa: ANN001 assert ds == dataset @@ -81,7 +81,7 @@ def fake_batch_evaluate(ds, preds): # noqa: ANN001 output_prefix = tmp_path / "results" combined = metrics.compute_metrics( dataset=dataset, - predicted_texts=predicted_texts, + predictions=predictions, output_prefix=output_prefix, elapsed_time=4.567, ) @@ -90,6 +90,7 @@ def fake_batch_evaluate(ds, preds): # noqa: ANN001 assert outputs == [ { "example_id": "file1", + "metadata": {"raw_html": "

"}, "main_language": "english", "document_type": "print", "normalized_levenshtein_similarity": 0.9, diff --git a/tests/test_tooling_support.py b/tests/test_tooling_support.py index af2dabe..32b3ad6 100644 --- a/tests/test_tooling_support.py +++ b/tests/test_tooling_support.py @@ -270,6 +270,85 @@ def test_evaluate_page_metric_helpers_cover_initialization_and_single_batch_path assert rows[0]["example_id"] == "row-1" +def test_batch_evaluate_initializes_worker_metrics_for_multi_example( + monkeypatch: pytest.MonkeyPatch, +) -> None: + init_calls = 0 + captured_initializer = None + + def fake_initialize_metrics() -> None: + nonlocal init_calls + init_calls += 1 + + class _FakePool: + def __init__(self, *, processes: int, initializer) -> None: # noqa: ANN001 + nonlocal captured_initializer + assert processes == 2 + captured_initializer = initializer + + def __enter__(self) -> _FakePool: + assert captured_initializer is not None + captured_initializer() + return self + + def __exit__(self, exc_type, exc, tb) -> bool: # noqa: ANN001 + return False + + def imap(self, func, iterable): # noqa: ANN001 + return map(func, iterable) + + monkeypatch.setattr(evaluate_page_module, "initialize_metrics", fake_initialize_metrics) + monkeypatch.setattr(evaluate_page_module.multiprocessing, "cpu_count", lambda: 2) + monkeypatch.setattr(evaluate_page_module.multiprocessing, "Pool", _FakePool) + monkeypatch.setattr( + evaluate_page_module, + "evaluate_page", + lambda inputs: cast( + "PageEvaluationResult", + { + "example_id": inputs[0]["example_id"], + "normalized_levenshtein_similarity": 1.0, + "is_empty": 0.0, + }, + ), + ) + + aggregate, rows = evaluate_page_module.batch_evaluate( + dataset=[ + cast( + "BenchmarkDatasetExample", + { + "image": "image", + "cleaned_transcription": "", + "dataset_id": "dataset-1", + "document_type": "print", + "example_id": "row-1", + "main_language": "English", + "main_script": "Latin", + }, + ), + cast( + "BenchmarkDatasetExample", + { + "image": "image", + "cleaned_transcription": "", + "dataset_id": "dataset-2", + "document_type": "print", + "example_id": "row-2", + "main_language": "English", + "main_script": "Latin", + }, + ), + ], + predicted_texts=["predicted-1", "predicted-2"], + ) + + assert captured_initializer is fake_initialize_metrics + assert init_calls == 2 + assert aggregate == {"normalized_levenshtein_similarity": 1.0, "is_empty": 0.0} + assert [row["example_id"] for row in rows] == ["row-1", "row-2"] + + def test_calculate_metrics_from_text_and_internal_error_fallback(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(evaluate_page_module, "initialize_metrics", lambda: None) monkeypatch.setattr( diff --git a/tooling/benchmarking/benchmark.py b/tooling/benchmarking/benchmark.py index 2ca1622..88ae9e5 100644 --- a/tooling/benchmarking/benchmark.py +++ b/tooling/benchmarking/benchmark.py @@ -22,6 +22,7 @@ sys.path.insert(0, _REPO_SRC_PATH_STR) from churro_ocr._internal.logging import logger +from churro_ocr.errors import ConfigurationError from churro_ocr.ocr import BatchOCRBackend, OCRBackend, OCRBackendLike from churro_ocr.page_detection import DocumentPage from churro_ocr.providers import ( @@ -32,19 +33,26 @@ MistralOptions, OCRBackendSpec, OpenAICompatibleOptions, - VLLMOptions, ) +from churro_ocr.providers.specs import MISTRAL_OCR_MODEL_IDS, validate_mistral_ocr_model from tooling.benchmarking.dataset import ( DatasetSelection, DatasetSubset, load_dataset_split, ) from tooling.evaluation.metrics import compute_metrics -from tooling.evaluation.types import BenchmarkDatasetExample, EvaluationExample, to_evaluation_example +from tooling.evaluation.types import ( + BenchmarkDatasetExample, + BenchmarkPrediction, + EvaluationExample, + to_evaluation_example, +) CHURRO_DATASET_ID = "stanford-oval/churro-dataset" VALID_DATASET_SPLITS = {"dev", "test"} -VALID_OCR_BACKENDS = {"litellm", "openai-compatible", "azure", "mistral", "hf", "vllm"} +VALID_OCR_BACKENDS = {"litellm", "openai-compatible", "azure", "mistral", "hf"} +PROGRESS_BAR_SMOOTHING = 0.05 +PROGRESS_BAR_MININTERVAL_SECONDS = 1.0 BENCHMARK_DATASET_COLUMNS = ( "image", "cleaned_transcription", @@ -73,8 +81,6 @@ class BenchmarkOptions: api_key: str | None = None base_url: str | None = None api_version: str | None = None - vllm_gpu_memory_utilization: float | None = None - vllm_cpu_offload_gb: float | None = None def dataset_subset(self) -> DatasetSubset: """Return the normalized subset filters for this benchmark run.""" @@ -103,13 +109,11 @@ def build_parser(*, add_help: bool = True) -> argparse.ArgumentParser: parser.add_argument("--input-size", type=int, default=0) parser.add_argument("--offset", type=int, default=0) parser.add_argument("--output-dir", type=Path, default=None) - parser.add_argument("--max-concurrency", type=int, default=32) + parser.add_argument("--max-concurrency", type=int, default=16) parser.add_argument("--endpoint", default=None) parser.add_argument("--api-key", default=None) parser.add_argument("--base-url", default=None) parser.add_argument("--api-version", default=None) - parser.add_argument("--vllm-gpu-memory-utilization", type=float, default=None) - parser.add_argument("--vllm-cpu-offload-gb", type=float, default=None) return parser @@ -130,8 +134,6 @@ def parse_args(argv: list[str] | None = None) -> BenchmarkOptions: api_key=namespace.api_key, base_url=namespace.base_url, api_version=namespace.api_version, - vllm_gpu_memory_utilization=namespace.vllm_gpu_memory_utilization, - vllm_cpu_offload_gb=namespace.vllm_cpu_offload_gb, ) @@ -148,31 +150,27 @@ def _validate_options(options: BenchmarkOptions) -> int: if options.backend == "litellm" and not options.model: logger.error("--model is required for backend=litellm.") return 1 - if options.backend == "openai-compatible" and ( - not options.model or not options.base_url or not options.api_key - ): - logger.error("--model, --base-url, and --api-key are required for backend=openai-compatible.") + if options.backend == "openai-compatible" and (not options.model or not options.base_url): + logger.error("--model and --base-url are required for backend=openai-compatible.") return 1 if options.backend == "hf" and not options.model: logger.error("--model is required for backend=hf.") return 1 - if options.backend == "vllm" and not options.model: - logger.error("--model is required for backend=vllm.") - return 1 - if options.vllm_gpu_memory_utilization is not None and not ( - 0.0 < options.vllm_gpu_memory_utilization <= 1.0 - ): - logger.error("--vllm-gpu-memory-utilization must be in the range (0, 1].") - return 1 - if options.vllm_cpu_offload_gb is not None and options.vllm_cpu_offload_gb < 0.0: - logger.error("--vllm-cpu-offload-gb must be non-negative.") - return 1 if options.backend == "azure" and (not options.endpoint or not options.api_key): logger.error("--endpoint and --api-key are required for backend=azure.") return 1 - if options.backend == "mistral" and not options.api_key: - logger.error("--api-key is required for backend=mistral.") - return 1 + if options.backend == "mistral": + if not options.api_key: + logger.error("--api-key is required for backend=mistral.") + return 1 + try: + validate_mistral_ocr_model(options.model) + except ConfigurationError: + logger.error( + "--model is required for backend=mistral and must be one of: %s.", + ", ".join(MISTRAL_OCR_MODEL_IDS), + ) + return 1 return 0 @@ -202,10 +200,22 @@ def _default_litellm_cache_dir() -> Path: return Path(__file__).resolve().parents[2] / "workdir" / "cache" / "litellm" +def _create_progress_bar(*, total: int | None, desc: str, unit: str) -> tqdm[object]: + """Return a tqdm progress bar tuned for steadier ETA updates.""" + return tqdm( + total=total, + desc=desc, + unit=unit, + mininterval=PROGRESS_BAR_MININTERVAL_SECONDS, + smoothing=PROGRESS_BAR_SMOOTHING, + ) + + def _build_evaluation_example(example: BenchmarkDatasetExample) -> EvaluationExample: """Keep only the fields needed for evaluation after OCR completes.""" return to_evaluation_example(example) + def _selected_dataset_examples( dataset_stream: Iterable[BenchmarkDatasetExample], options: BenchmarkOptions, @@ -232,7 +242,6 @@ def _build_ocr_backend(options: BenchmarkOptions) -> OCRBackendLike: if options.backend == "openai-compatible": assert options.model is not None assert options.base_url is not None - assert options.api_key is not None return build_ocr_backend( OCRBackendSpec( provider="openai-compatible", @@ -267,25 +276,13 @@ def _build_ocr_backend(options: BenchmarkOptions) -> OCRBackendLike: options=HuggingFaceOptions(model_kwargs={"device_map": "auto", "torch_dtype": "auto"}), ) ) - if options.backend == "vllm": - assert options.model is not None - llm_kwargs: dict[str, object] = {} - if options.vllm_gpu_memory_utilization is not None: - llm_kwargs["gpu_memory_utilization"] = options.vllm_gpu_memory_utilization - if options.vllm_cpu_offload_gb is not None: - llm_kwargs["cpu_offload_gb"] = options.vllm_cpu_offload_gb - return build_ocr_backend( - OCRBackendSpec( - provider="vllm", - model=options.model, - options=VLLMOptions(llm_kwargs=llm_kwargs), - ) - ) assert options.api_key is not None + assert options.model is not None + mistral_model = validate_mistral_ocr_model(options.model) return build_ocr_backend( OCRBackendSpec( provider="mistral", - model=options.model or "mistral-ocr-latest", + model=mistral_model, options=MistralOptions(api_key=options.api_key), ) ) @@ -306,17 +303,17 @@ async def _predict_texts( options: BenchmarkOptions, *, total_pages: int | None = None, -) -> tuple[list[EvaluationExample], list[str]]: +) -> tuple[list[EvaluationExample], list[BenchmarkPrediction]]: ocr_backend = _build_ocr_backend(options) max_in_flight = max(1, options.max_concurrency) has_logged_first_output = False if isinstance(ocr_backend, BatchOCRBackend): dataset_iterator = iter(dataset) evaluation_examples: list[EvaluationExample] = [] - predicted_texts: list[str] = [] + predictions: list[BenchmarkPrediction] = [] submitted_pages = 0 - with tqdm(total=total_pages, desc="OCR", unit="page") as progress: + with _create_progress_bar(total=total_pages, desc="OCR", unit="page") as progress: while True: batch_examples: list[BenchmarkDatasetExample] = [] pages: list[DocumentPage] = [] @@ -346,25 +343,34 @@ async def _predict_texts( text=batch_results[0].text or "", ) has_logged_first_output = True - predicted_texts.extend((result.text or "") for result in batch_results) + predictions.extend( + { + "text": result.text or "", + "metadata": dict(result.metadata), + } + for result in batch_results + ) progress.update(len(batch_results)) progress.set_postfix(submitted=submitted_pages, in_flight=0, refresh=False) - return evaluation_examples, predicted_texts + return evaluation_examples, predictions - async def _predict(index: int, image: Image.Image) -> tuple[int, str]: + async def _predict(index: int, image: Image.Image) -> tuple[int, BenchmarkPrediction]: page = DocumentPage(page_index=index, source_index=0, image=image) if callable(ocr_backend) and not isinstance(ocr_backend, OCRBackend): result = await ocr_backend(page) else: assert isinstance(ocr_backend, OCRBackend) result = await ocr_backend.ocr(page) - return index, (result.text or "") + return index, { + "text": result.text or "", + "metadata": dict(result.metadata), + } dataset_iterator = iter(dataset) evaluation_examples: list[EvaluationExample] = [] - pending_tasks: set[asyncio.Task[tuple[int, str]]] = set() - predicted_texts: list[str] = [] + pending_tasks: set[asyncio.Task[tuple[int, BenchmarkPrediction]]] = set() + predictions: list[BenchmarkPrediction] = [] next_index = 0 wait_poll_seconds = 1.0 @@ -381,7 +387,7 @@ def _progress_heartbeat(progress: tqdm[object], stop_event: threading.Event) -> while not stop_event.wait(wait_poll_seconds): _update_progress_status(progress, force_refresh=True) - with tqdm(total=total_pages, desc="OCR", unit="page") as progress: + with _create_progress_bar(total=total_pages, desc="OCR", unit="page") as progress: heartbeat_stop_event = threading.Event() heartbeat_thread = threading.Thread( target=_progress_heartbeat, @@ -400,7 +406,7 @@ def _progress_heartbeat(progress: tqdm[object], stop_event: threading.Event) -> image = example["image"] assert isinstance(image, Image.Image) evaluation_examples.append(_build_evaluation_example(example)) - predicted_texts.append("") + predictions.append({"text": "", "metadata": {}}) pending_tasks.add(asyncio.create_task(_predict(next_index, image))) next_index += 1 _update_progress_status(progress) @@ -413,10 +419,10 @@ def _progress_heartbeat(progress: tqdm[object], stop_event: threading.Event) -> return_when=asyncio.FIRST_COMPLETED, ) for task in done_tasks: - index, text = await task - predicted_texts[index] = text + index, prediction = await task + predictions[index] = prediction if not has_logged_first_output and index == 0: - _log_first_benchmark_output(options=options, text=text) + _log_first_benchmark_output(options=options, text=prediction["text"]) has_logged_first_output = True progress.update(1) _update_progress_status(progress) @@ -430,7 +436,7 @@ def _progress_heartbeat(progress: tqdm[object], stop_event: threading.Event) -> heartbeat_thread.join(timeout=wait_poll_seconds * 2) _update_progress_status(progress, force_refresh=True) - return evaluation_examples, predicted_texts + return evaluation_examples, predictions async def run(options: BenchmarkOptions) -> int: @@ -447,17 +453,17 @@ async def run(options: BenchmarkOptions) -> int: output_prefix = create_output_prefix(options) start_time = time() - evaluation_examples, predicted_texts = await _predict_texts( + evaluation_examples, predictions = await _predict_texts( dataset, options, total_pages=total_pages, ) elapsed_time = time() - start_time - assert len(evaluation_examples) == len(predicted_texts), ( - f"Mismatch in dataset size ({len(evaluation_examples)}) and predictions ({len(predicted_texts)})." + assert len(evaluation_examples) == len(predictions), ( + f"Mismatch in dataset size ({len(evaluation_examples)}) and predictions ({len(predictions)})." ) - compute_metrics(evaluation_examples, predicted_texts, output_prefix, elapsed_time) + compute_metrics(evaluation_examples, predictions, output_prefix, elapsed_time) return 0 diff --git a/tooling/evaluation/evaluate_page.py b/tooling/evaluation/evaluate_page.py index 587fd16..7c9038d 100644 --- a/tooling/evaluation/evaluate_page.py +++ b/tooling/evaluation/evaluate_page.py @@ -91,6 +91,8 @@ def _compute_text_metrics_core( has_repetition_flag = has_long_repetition(predicted_text) if is_empty != 1.0: + if bleu_metric is None: + initialize_metrics() assert bleu_metric is not None bleu_result = bleu_metric.compute( predictions=[predicted_text], @@ -216,7 +218,7 @@ def batch_evaluate( return aggregate_results(results) processes = min(8, max(1, multiprocessing.cpu_count())) - with multiprocessing.Pool(processes=processes) as pool: + with multiprocessing.Pool(processes=processes, initializer=initialize_metrics) as pool: results = list( tqdm( pool.imap(evaluate_page, zip(dataset, predicted_texts, strict=False)), diff --git a/tooling/evaluation/metrics.py b/tooling/evaluation/metrics.py index e4f016b..ba1ac3c 100644 --- a/tooling/evaluation/metrics.py +++ b/tooling/evaluation/metrics.py @@ -3,13 +3,19 @@ from __future__ import annotations from collections import defaultdict +from collections.abc import Sequence import json from pathlib import Path from typing import Any from churro_ocr._internal.logging import logger from tooling.evaluation.evaluate_page import batch_evaluate -from tooling.evaluation.types import EvaluationExample, PageEvaluationResult +from tooling.evaluation.types import ( + BenchmarkOutputRow, + BenchmarkPrediction, + EvaluationExample, + PageEvaluationResult, +) def _get_llm_total_cost() -> float: @@ -39,7 +45,7 @@ def to_rounded_percentage(metrics: dict[str, Any]) -> dict[str, Any]: def calculate_language_and_type_metrics( - outputs: list[PageEvaluationResult], + outputs: Sequence[PageEvaluationResult], main_metric: str = "normalized_levenshtein_similarity", ) -> tuple[dict[str, float], dict[str, float], dict[str, float]]: """Compute averages grouped by language, document type, and both.""" @@ -70,28 +76,49 @@ def calculate_language_and_type_metrics( return averaged_language, averaged_type, averaged_language_type +def _normalize_prediction(prediction: str | BenchmarkPrediction) -> BenchmarkPrediction: + """Coerce legacy string predictions into the structured benchmark format.""" + if isinstance(prediction, str): + return {"text": prediction, "metadata": {}} + return { + "text": str(prediction["text"]), + "metadata": dict(prediction["metadata"]), + } + + def _build_output_rows( - dataset: list[EvaluationExample], - per_example_outputs: list[PageEvaluationResult], -) -> list[PageEvaluationResult]: + dataset: Sequence[EvaluationExample], + per_example_outputs: Sequence[PageEvaluationResult], + predictions: Sequence[BenchmarkPrediction], +) -> list[BenchmarkOutputRow]: """Attach stable example ids to per-example outputs before writing them.""" return [ - {"example_id": str(example["example_id"]), **evaluation_output} - for evaluation_output, example in zip(per_example_outputs, dataset, strict=False) + { + "example_id": str(example["example_id"]), + "metadata": dict(prediction["metadata"]), + **evaluation_output, + } + for evaluation_output, example, prediction in zip( + per_example_outputs, + dataset, + predictions, + strict=False, + ) ] def compute_metrics( dataset: list[EvaluationExample], - predicted_texts: list[str], + predictions: list[str] | list[BenchmarkPrediction], output_prefix: str | Path, elapsed_time: float, main_metric: str = "normalized_levenshtein_similarity", ) -> dict[str, Any]: """Compute aggregate metrics and write benchmark output files.""" - sanitized_predictions = [prediction or "" for prediction in predicted_texts] + normalized_predictions = [_normalize_prediction(prediction) for prediction in predictions] + sanitized_predictions = [prediction["text"] or "" for prediction in normalized_predictions] aggregate_metrics, per_example_outputs = batch_evaluate(dataset, sanitized_predictions) - outputs = _build_output_rows(dataset, per_example_outputs) + outputs = _build_output_rows(dataset, per_example_outputs, normalized_predictions) output_dir = Path(output_prefix) output_dir.mkdir(parents=True, exist_ok=True) diff --git a/tooling/evaluation/types.py b/tooling/evaluation/types.py index f650679..ede76f1 100644 --- a/tooling/evaluation/types.py +++ b/tooling/evaluation/types.py @@ -1,8 +1,8 @@ -"""Shared dataset and result types for CHURRO tooling evaluation flows.""" +"""Shared dataset, prediction, and result types for CHURRO tooling evaluation flows.""" from __future__ import annotations -from typing import TypedDict +from typing import Any, TypedDict from PIL import Image @@ -61,6 +61,19 @@ class PageEvaluationResult(PageEvaluationMetrics): document_type: str +class BenchmarkPrediction(TypedDict): + """OCR output retained during benchmarking before metrics are computed.""" + + text: str + metadata: dict[str, Any] + + +class BenchmarkOutputRow(PageEvaluationResult): + """Serialized benchmark output row written to ``outputs.json``.""" + + metadata: dict[str, Any] + + def to_evaluation_example(example: BenchmarkDatasetExample) -> EvaluationExample: """Keep only the dataset fields needed after OCR completes.""" return {field_name: example[field_name] for field_name in EVALUATION_EXAMPLE_FIELDS}