| """Script to download quantized models for DocNeedle. |
| |
| - Embedding: unsloth/embeddinggemma-2-GGUF (embeddinggemma-2-UD-Q4_K_XL.gguf + mmproj-F16.gguf) |
| - Reranker: mradermacher/Qwen3-Reranker-0.6B-GGUF (Qwen3-Reranker-0.6B.Q4_K_M.gguf) |
| - LLM: unsloth/Qwen3.5-4B-GGUF (Qwen3.5-4B-Q4_K_M.gguf + mmproj-Qwen3.5-4B-F16.gguf) |
| """ |
|
|
| import argparse |
| import os |
| import sys |
| import urllib.request |
|
|
| MODELS_CONFIG = { |
| "embedding": { |
| "name": "embeddinggemma-2-UD-Q4_K_XL", |
| "repo_id": "unsloth/embeddinggemma-2-GGUF", |
| "filename": "embeddinggemma-2-UD-Q4_K_XL.gguf", |
| "url": ( |
| "https://huggingface.co/unsloth/embeddinggemma-2-GGUF/resolve/main/" |
| "embeddinggemma-2-UD-Q4_K_XL.gguf" |
| ), |
| "subpath": os.path.join("embedding", "embeddinggemma-2-UD-Q4_K_XL.gguf"), |
| }, |
| "embedding_mmproj": { |
| "name": "embeddinggemma-2-mmproj-F16", |
| "repo_id": "unsloth/embeddinggemma-2-GGUF", |
| "filename": "mmproj-F16.gguf", |
| "url": ( |
| "https://huggingface.co/unsloth/embeddinggemma-2-GGUF/resolve/main/mmproj-F16.gguf" |
| ), |
| "subpath": os.path.join("embedding", "mmproj-F16.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": "Qwen3.5-4B-Q4_K_M", |
| "repo_id": "unsloth/Qwen3.5-4B-GGUF", |
| "filename": "Qwen3.5-4B-Q4_K_M.gguf", |
| "url": ( |
| "https://huggingface.co/unsloth/Qwen3.5-4B-GGUF/resolve/main/Qwen3.5-4B-Q4_K_M.gguf" |
| ), |
| "subpath": os.path.join("llm", "Qwen3.5-4B-Q4_K_M.gguf"), |
| }, |
| "llm_mmproj": { |
| "name": "mmproj-Qwen3.5-4B-F16", |
| "repo_id": "unsloth/Qwen3.5-4B-GGUF", |
| "filename": "mmproj-F16.gguf", |
| "url": ("https://huggingface.co/unsloth/Qwen3.5-4B-GGUF/resolve/main/mmproj-F16.gguf"), |
| "subpath": os.path.join("llm", "mmproj-Qwen3.5-4B-F16.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 |
|
|
| downloaded = hf_hub_download( |
| repo_id=repo_id, |
| filename=filename, |
| local_dir=os.path.dirname(dest_path), |
| ) |
| dest_abs = os.path.abspath(dest_path) |
| down_abs = os.path.abspath(downloaded) |
| if down_abs != dest_abs and os.path.exists(down_abs): |
| import shutil |
|
|
| shutil.copy2(down_abs, dest_abs) |
| 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_embedding_models(models_dir: str) -> None: |
| """Download EmbeddingGemma-2 GGUF model and SigLIP vision projector.""" |
| for key in ("embedding", "embedding_mmproj"): |
| cfg = MODELS_CONFIG[key] |
| dest_path = os.path.join(models_dir, cfg["subpath"]) |
| print(f"\nChecking/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"\nChecking/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_models(models_dir: str) -> None: |
| """Download Qwen3.5-4B GGUF LLM and mmproj vision projector.""" |
| for key in ("llm", "llm_mmproj"): |
| cfg = MODELS_CONFIG[key] |
| dest_path = os.path.join(models_dir, cfg["subpath"]) |
| print(f"\nChecking/Downloading {cfg['name']} into {dest_path}...") |
| download_file( |
| cfg["url"], |
| dest_path, |
| repo_id=cfg.get("repo_id"), |
| filename=cfg.get("filename"), |
| ) |
|
|
|
|
| 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 |
|
|
| try: |
| AutoTokenizer.from_pretrained(tokenizer_model, local_files_only=True) |
| return |
| except Exception: |
| pass |
|
|
| 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_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 EmbeddingGemma-2 GGUF & SigLIP vision projector", |
| ) |
| 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 Qwen3.5-4B Q4_K_M GGUF & vision projector", |
| ) |
|
|
| args = parser.parse_args() |
| download_all = not args.embedding and not args.reranker and not args.llm |
|
|
| if download_all or args.embedding: |
| download_embedding_models(args.models_dir) |
|
|
| if download_all or args.reranker: |
| download_reranker_model(args.models_dir) |
|
|
| if download_all or args.llm: |
| download_llm_models(args.models_dir) |
|
|
| download_chunker_tokenizer() |
| print("\nAll requested model downloads are complete.") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
