Spaces:
Running on Zero
Running on Zero
File size: 2,488 Bytes
e834698 d627632 e834698 d627632 e834698 d627632 e834698 d627632 e834698 d627632 e834698 d627632 7144723 d627632 7144723 d627632 e834698 7144723 aaccd1d 7144723 aaccd1d e834698 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | from __future__ import annotations
from pathlib import Path
from threading import Lock
from typing import Optional
import time
import requests
from huggingface_hub import hf_hub_download
from config import (
CACHE_DIR,
FAISS_LOCAL,
FAISS_URL,
IDLE_UNLOAD_SECONDS,
MANIFEST_LOCAL,
MANIFEST_URL,
RERANKER_FILENAME,
RERANKER_REPO_ID,
VECTORS_LOCAL,
VECTORS_URL,
)
from index.store import IndexStore
def _download_to(url: str, dst: Path) -> Path:
dst.parent.mkdir(parents=True, exist_ok=True)
if dst.exists() and dst.stat().st_size > 0:
return dst
with requests.get(url, stream=True, timeout=600) as r:
r.raise_for_status()
with open(dst, "wb") as f:
for chunk in r.iter_content(chunk_size=1024 * 1024):
if chunk:
f.write(chunk)
return dst
class RuntimeState:
def __init__(self) -> None:
self.store: Optional[IndexStore] = None
self.reranker_model_path: Optional[Path] = None
self.loading = False
self.last_used_ts = 0.0
self._lock = Lock()
def ensure_assets(self) -> None:
_download_to(FAISS_URL, FAISS_LOCAL)
_download_to(MANIFEST_URL, MANIFEST_LOCAL)
_download_to(VECTORS_URL, VECTORS_LOCAL)
local_path = hf_hub_download(
repo_id=RERANKER_REPO_ID,
filename=RERANKER_FILENAME,
)
self.reranker_model_path = Path(local_path)
def maybe_unload(self) -> None:
if IDLE_UNLOAD_SECONDS <= 0:
return
now = time.time()
if self.store is not None and (now - self.last_used_ts) > IDLE_UNLOAD_SECONDS:
self.store = None
def ensure_loaded(self) -> None:
self.maybe_unload()
if self.store is not None:
self.last_used_ts = time.time()
return
with self._lock:
if self.store is not None:
self.last_used_ts = time.time()
return
self.loading = True
try:
self.ensure_assets()
self.store = IndexStore(root=CACHE_DIR).load_all()
self.last_used_ts = time.time()
finally:
self.loading = False
def status_text(self) -> str:
if self.loading:
return "Loading assets..."
if self.store is None:
return "Cold state"
return "Ready"
RUNTIME = RuntimeState() |