Spaces:
Running on Zero
Running on Zero
| """Optional ZeroGPU query encoding; CPU ranking and Gemini stay outside the GPU lease.""" | |
| import json | |
| import os | |
| from pathlib import Path | |
| # Import spaces before torch/transformers so ZeroGPU can install its CUDA hooks. | |
| if os.environ.get("VDR_QUERY_DEVICE", "cpu") == "zerogpu": | |
| import spaces | |
| else: | |
| spaces = None | |
| _encoder = None | |
| def ensure_models(data_dir, model_info): | |
| """Space repositories omit weights; fetch only the pinned official snapshots.""" | |
| from huggingface_hub import snapshot_download | |
| for kind in ("adapter", "base"): | |
| root = data_dir / "models" / model_info[f"{kind}_path"] | |
| if not list(root.glob("*.safetensors")): | |
| snapshot_download( | |
| repo_id=model_info["model_id" if kind == "adapter" else "base_model_id"], | |
| revision=model_info[f"{kind}_revision"], | |
| local_dir=root, | |
| allow_patterns=["*.json", "*.safetensors", "*.jinja", "README.md", "LICENSE*"], | |
| token=False, | |
| ) | |
| def encode_query(query): | |
| if _encoder is None: | |
| raise RuntimeError("Initialize the ZeroGPU encoder at application startup.") | |
| return _encoder.encode(query=query) | |
| if spaces is not None: | |
| encode_query = spaces.GPU(duration=10)(encode_query) | |
| def create_service(data_dir=Path("data")): | |
| from visual_doc_rag.serve import AnswerService | |
| device = os.environ.get("VDR_QUERY_DEVICE", "cpu") | |
| if device == "cpu": | |
| manifest_path = data_dir / "hr/indexes/visual.json" | |
| if (data_dir / "hr/demo.json").exists() and manifest_path.exists(): | |
| manifest = json.loads(manifest_path.read_text(encoding="utf-8")) | |
| ensure_models(data_dir, manifest["signature"]["model"]) | |
| return AnswerService(data_dir=data_dir) | |
| if device != "zerogpu" or spaces is None: | |
| raise ValueError("VDR_QUERY_DEVICE must be cpu or zerogpu (requires the hosting extra).") | |
| from visual_doc_rag.retrieve import PageRetriever | |
| from visual_doc_rag.visual_index import Encoder | |
| global _encoder | |
| manifest = json.loads((data_dir / "hr/indexes/visual.json").read_text(encoding="utf-8")) | |
| ensure_models(data_dir, manifest["signature"]["model"]) | |
| # Official ZeroGPU guidance: place the model on CUDA at startup, before a GPU call. | |
| _encoder = Encoder(data_dir, manifest["signature"]["model"], "cuda") | |
| retriever = PageRetriever(data_dir=data_dir, query_encoder=encode_query) | |
| return AnswerService(data_dir=data_dir, retriever=retriever) | |