Spaces:
Sleeping
Sleeping
| """Single-source-of-truth offline-mode enforcer. | |
| Both inference.py and the analysis scripts import this so the same logic runs | |
| everywhere. If the BAAI bge-m3 + bge-reranker-v2-m3 model directories are | |
| present in the local HuggingFace cache AND each has a usable weights file, | |
| we set HF_HUB_OFFLINE / TRANSFORMERS_OFFLINE so subsequent model loads | |
| NEVER touch the network. | |
| Override with HF_FORCE_NETWORK=1 if you really want to re-pull from the Hub. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| import sys | |
| def _has_usable_weights(model_dir: str) -> bool: | |
| """Walk a HF cache model directory and confirm at least one usable | |
| weight file is present. The directory existing isn't enough — a | |
| partial download (interrupted by network) leaves directories but no | |
| weights. If we flip offline mode on a partial cache, the actual | |
| model load later fails opaquely with `OSError: no file named ...`. | |
| """ | |
| needles = ("model.safetensors", "pytorch_model.bin") | |
| snapshots_dir = os.path.join(model_dir, "snapshots") | |
| if not os.path.isdir(snapshots_dir): | |
| return False | |
| try: | |
| for entry in os.listdir(snapshots_dir): | |
| snap = os.path.join(snapshots_dir, entry) | |
| if not os.path.isdir(snap): | |
| continue | |
| try: | |
| names = os.listdir(snap) | |
| except OSError: | |
| continue | |
| if any(n in names for n in needles): | |
| return True | |
| except OSError: | |
| return False | |
| return False | |
| def enforce_offline_if_cached(verbose: bool = True) -> bool: | |
| """Return True iff we just enforced offline mode (or it was already on). | |
| Requires not just the model directories to exist, but also a | |
| usable weights file inside each. Prevents offline-mode flipping on | |
| a partially downloaded cache (which would later crash at model load). | |
| """ | |
| if os.environ.get("HF_FORCE_NETWORK") == "1": | |
| return False | |
| if os.environ.get("HF_HUB_OFFLINE") == "1" and os.environ.get("TRANSFORMERS_OFFLINE") == "1": | |
| return True | |
| hf_home = os.environ.get("HF_HOME") or os.path.join( | |
| os.path.expanduser("~"), ".cache", "huggingface" | |
| ) | |
| hub_dir = os.path.join(hf_home, "hub") | |
| if not os.path.isdir(hub_dir): | |
| return False | |
| needed = ("models--BAAI--bge-m3", "models--BAAI--bge-reranker-v2-m3") | |
| for name in needed: | |
| path = os.path.join(hub_dir, name) | |
| if not os.path.isdir(path) or not _has_usable_weights(path): | |
| return False | |
| os.environ["HF_HUB_OFFLINE"] = "1" | |
| os.environ["TRANSFORMERS_OFFLINE"] = "1" | |
| if verbose: | |
| print( | |
| "[offline-guard] HF cache complete (weights verified) — enforcing offline mode " | |
| "(HF_HUB_OFFLINE=1, TRANSFORMERS_OFFLINE=1)", | |
| file=sys.stderr, | |
| ) | |
| return True | |