| """Script to download quantized models for DocNeedle. |
| |
| - Embedding: mradermacher/harrier-oss-v1-0.6b-GGUF (harrier-oss-v1-0.6b.Q4_K_M.gguf) |
| - Reranker: mradermacher/Qwen3-Reranker-0.6B-GGUF (Qwen3-Reranker-0.6B.Q4_K_M.gguf) |
| - LLM: unsloth/gemma-4-E2B-it-GGUF (gemma-4-E2B-it-qat-UD-Q4_K_XL.gguf) |
| """ |
|
|
| import argparse |
| import os |
| import sys |
| import urllib.request |
|
|
| MODELS_CONFIG = { |
| "embedding": { |
| "name": "harrier-oss-v1-0.6b.Q4_K_M", |
| "repo_id": "mradermacher/harrier-oss-v1-0.6b-GGUF", |
| "filename": "harrier-oss-v1-0.6b.Q4_K_M.gguf", |
| "url": ( |
| "https://huggingface.co/mradermacher/harrier-oss-v1-0.6b-GGUF/resolve/main/" |
| "harrier-oss-v1-0.6b.Q4_K_M.gguf" |
| ), |
| "subpath": os.path.join("embedding", "harrier-oss-v1-0.6b.Q4_K_M.gguf"), |
| }, |
| "reranker": { |
| "name": "Qwen3-Reranker-0.6B.Q4_K_M", |
| "repo_id": "mradermacher/Qwen3-Reranker-0.6B-GGUF", |
| "filename": "Qwen3-Reranker-0.6B.Q4_K_M.gguf", |
| "url": ( |
| "https://huggingface.co/mradermacher/Qwen3-Reranker-0.6B-GGUF/resolve/main/" |
| "Qwen3-Reranker-0.6B.Q4_K_M.gguf" |
| ), |
| "subpath": os.path.join("reranker", "Qwen3-Reranker-0.6B.Q4_K_M.gguf"), |
| }, |
| "llm": { |
| "name": "gemma-4-E2B-it-qat-UD-Q4_K_XL", |
| "repo_id": "unsloth/gemma-4-E2B-it-qat-GGUF", |
| "filename": "gemma-4-E2B-it-qat-UD-Q4_K_XL.gguf", |
| "url": ( |
| "https://huggingface.co/unsloth/gemma-4-E2B-it-qat-GGUF/resolve/main/" |
| "gemma-4-E2B-it-qat-UD-Q4_K_XL.gguf" |
| ), |
| "subpath": os.path.join("llm", "gemma-4-E2B-it-qat-UD-Q4_K_XL.gguf"), |
| }, |
| } |
|
|
|
|
| def download_file( |
| url: str, |
| dest_path: str, |
| repo_id: str | None = None, |
| filename: str | None = None, |
| ) -> None: |
| """Download a file with progress reporting.""" |
| os.makedirs(os.path.dirname(dest_path), exist_ok=True) |
| if os.path.exists(dest_path): |
| print(f"File already exists: {dest_path}") |
| return |
|
|
| print(f"Downloading: {url} -> {dest_path}") |
| if repo_id and filename: |
| try: |
| from huggingface_hub import hf_hub_download |
|
|
| hf_hub_download( |
| repo_id=repo_id, |
| filename=filename, |
| local_dir=os.path.dirname(dest_path), |
| ) |
| print("\n Download complete.") |
| return |
| except Exception as e: |
| print(f"hf_hub_download failed ({e}), falling back to standard download...") |
|
|
| def report_progress(block_num: int, block_size: int, total_size: int) -> None: |
| downloaded = block_num * block_size |
| if total_size > 0: |
| percent = downloaded / total_size * 100 |
| mb_downloaded = downloaded / (1024 * 1024) |
| mb_total = total_size / (1024 * 1024) |
| sys.stdout.write( |
| f"\r Progress: {percent:.1f}% ({mb_downloaded:.1f}MB / {mb_total:.1f}MB)" |
| ) |
| else: |
| mb_downloaded = downloaded / (1024 * 1024) |
| sys.stdout.write(f"\r Downloaded: {mb_downloaded:.1f}MB") |
| sys.stdout.flush() |
|
|
| headers = {"User-Agent": "DocNeedle-ModelDownloader/1.0"} |
| req = urllib.request.Request(url, headers=headers) |
| with urllib.request.urlopen(req) as response, open(dest_path, "wb") as out_file: |
| total_size = int(response.headers.get("Content-Length", 0)) |
| block_size = 1024 * 1024 |
| downloaded = 0 |
| block_num = 0 |
| while True: |
| chunk = response.read(block_size) |
| if not chunk: |
| break |
| out_file.write(chunk) |
| downloaded += len(chunk) |
| block_num += 1 |
| report_progress(block_num, block_size, total_size) |
| print("\n Download complete.") |
|
|
|
|
| def download_onnx_embedding_model(models_dir: str) -> None: |
| """Download Harrier-OSS-v1-0.6B ONNX Q4 model and tokenizer from onnx-community.""" |
| from huggingface_hub import hf_hub_download |
|
|
| dest_dir = os.path.join(models_dir, "embedding", "harrier-onnx") |
| os.makedirs(dest_dir, exist_ok=True) |
| repo_id = "onnx-community/harrier-oss-v1-0.6b-ONNX" |
|
|
| files = [ |
| "onnx/model_q4.onnx", |
| "onnx/model_q4.onnx_data", |
| "tokenizer.json", |
| "tokenizer_config.json", |
| ] |
| print(f"\nChecking/Downloading Harrier-OSS ONNX model into {dest_dir}...") |
| for f in files: |
| target = os.path.join(dest_dir, f) |
| if not os.path.exists(target): |
| print(f" Downloading {f}...") |
| hf_hub_download(repo_id=repo_id, filename=f, local_dir=dest_dir) |
| else: |
| print(f" {f} already exists.") |
|
|
|
|
| def download_embedding_model(models_dir: str) -> None: |
| """Download Harrier-OSS-v1-0.6B Q4_K_M GGUF.""" |
| cfg = MODELS_CONFIG["embedding"] |
| dest_path = os.path.join(models_dir, cfg["subpath"]) |
| print(f"\n[1/3] Checking/Downloading {cfg['name']} into {dest_path}...") |
| download_file( |
| cfg["url"], |
| dest_path, |
| repo_id=cfg.get("repo_id"), |
| filename=cfg.get("filename"), |
| ) |
|
|
|
|
| def download_reranker_model(models_dir: str) -> None: |
| """Download Qwen3-Reranker-0.6B Q4_K_M GGUF.""" |
| cfg = MODELS_CONFIG["reranker"] |
| dest_path = os.path.join(models_dir, cfg["subpath"]) |
| print(f"\n[2/3] Checking/Downloading {cfg['name']} into {dest_path}...") |
| download_file( |
| cfg["url"], |
| dest_path, |
| repo_id=cfg.get("repo_id"), |
| filename=cfg.get("filename"), |
| ) |
|
|
|
|
| def download_llm_model(models_dir: str) -> None: |
| """Download Gemma-4-E2B-it-qat-UD-Q4_K_XL GGUF.""" |
| cfg = MODELS_CONFIG["llm"] |
| dest_path = os.path.join(models_dir, cfg["subpath"]) |
| print(f"\n[3/3] Checking/Downloading {cfg['name']} into {dest_path}...") |
| download_file( |
| cfg["url"], |
| dest_path, |
| repo_id=cfg.get("repo_id"), |
| filename=cfg.get("filename"), |
| ) |
|
|
|
|
| def download_rapidocr_models() -> None: |
| """Pre-download RapidOCR detection and recognition ONNX models without importing onnxruntime.""" |
| try: |
| from pathlib import Path |
|
|
| import rapidocr |
| from rapidocr.inference_engine.base import FileInfo |
| from rapidocr.utils.download_models import download_task |
| from rapidocr.utils.typings import ( |
| EngineType, |
| LangDet, |
| LangRec, |
| ModelType, |
| OCRVersion, |
| TaskType, |
| ) |
|
|
| print("\nChecking/Downloading RapidOCR PP-OCRv6 models...") |
| ocr_models_dir = Path(rapidocr.__file__).resolve().parent / "models" |
| ocr_models_dir.mkdir(parents=True, exist_ok=True) |
| for model_tier in (ModelType.SMALL, ModelType.MEDIUM): |
| try: |
| download_task( |
| ocr_models_dir, |
| FileInfo( |
| EngineType.ONNXRUNTIME, |
| OCRVersion.PPOCRV6, |
| TaskType.DET, |
| LangDet.CH, |
| model_tier, |
| ), |
| ) |
| except ValueError: |
| pass |
| try: |
| download_task( |
| ocr_models_dir, |
| FileInfo( |
| EngineType.ONNXRUNTIME, |
| OCRVersion.PPOCRV6, |
| TaskType.REC, |
| LangRec.CH, |
| model_tier, |
| ), |
| ) |
| except ValueError: |
| pass |
| print("RapidOCR models are verified and present.") |
| except Exception as e: |
| print(f"RapidOCR models check bypassed: {e}") |
|
|
|
|
| def download_chunker_tokenizer(tokenizer_model: str = "Qwen/Qwen3-Reranker-0.6B") -> None: |
| """Pre-download Hugging Face tokenizer used by DoclingChunker so GPU workers do not wait.""" |
| try: |
| from transformers import AutoTokenizer |
|
|
| print(f"\nChecking/Downloading chunker tokenizer ({tokenizer_model})...") |
| AutoTokenizer.from_pretrained(tokenizer_model) |
| print("Chunker tokenizer is verified and cached.") |
| except Exception as e: |
| print(f"Chunker tokenizer check bypassed: {e}") |
|
|
|
|
| def ensure_all_models_present(models_dir: str | None = None) -> None: |
| """Ensure all required quantized models exist, downloading any missing ones.""" |
| if models_dir is None: |
| models_dir = os.environ.get("MODEL_STORAGE_DIR", "models") |
|
|
| for _, cfg in MODELS_CONFIG.items(): |
| dest = os.path.join(models_dir, cfg["subpath"]) |
| if not os.path.exists(dest): |
| print(f"Missing model [{cfg['name']}]. Starting automatic bootstrap download...") |
| download_file( |
| cfg["url"], |
| dest, |
| repo_id=cfg.get("repo_id"), |
| filename=cfg.get("filename"), |
| ) |
| download_rapidocr_models() |
| download_chunker_tokenizer() |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description="Download DocNeedle quantized models") |
| parser.add_argument( |
| "--models-dir", |
| default="models", |
| help="Target directory for downloaded models (default: models)", |
| ) |
| parser.add_argument( |
| "--embedding", |
| action="store_true", |
| help="Download only Harrier-OSS-v1-0.6B Q4_K_M GGUF", |
| ) |
| parser.add_argument( |
| "--reranker", |
| action="store_true", |
| help="Download only Qwen3-Reranker-0.6B Q4_K_M GGUF", |
| ) |
| parser.add_argument( |
| "--llm", |
| action="store_true", |
| help="Download only Gemma-4-E2B-it-qat-UD-Q4_K_XL GGUF", |
| ) |
| parser.add_argument( |
| "--onnx", |
| action="store_true", |
| help="Download Harrier-OSS-v1-0.6B ONNX Q4 model and tokenizer", |
| ) |
|
|
| args = parser.parse_args() |
| download_all = not args.embedding and not args.reranker and not args.llm and not args.onnx |
|
|
| if download_all or args.embedding: |
| download_embedding_model(args.models_dir) |
|
|
| if download_all or args.onnx: |
| download_onnx_embedding_model(args.models_dir) |
|
|
| if download_all or args.reranker: |
| download_reranker_model(args.models_dir) |
|
|
| if download_all or args.llm: |
| download_llm_model(args.models_dir) |
|
|
| print("\nAll requested model downloads are complete.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
