Spaces:
Sleeping
Sleeping
File size: 2,887 Bytes
15b4b57 | 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 | import os
import sys
import shutil
import urllib.request
import tarfile
import zipfile
from pathlib import Path
# Configuração via variáveis de ambiente (podem ser definidas no docker-compose.prod.yaml)
BOOTSTRAP_URL = os.getenv("BOOTSTRAP_FAISS_URL")
FAISS_DB_DIR = Path(os.getenv("BASE_FAISS_DIR", "/app/faiss_db"))
def _sincronizar_faiss_do_hub() -> None:
"""
Descarrega o faiss_db/ do Hugging Face Dataset repo para o disco local.
Se HF_TOKEN ou HF_DATASET_REPO não estiverem definidos, salta silenciosamente.
Nunca bloqueia o arranque nem lança excepções não tratadas.
"""
hf_token = os.getenv("HF_TOKEN")
hf_repo = os.getenv("HF_DATASET_REPO")
if not hf_token or not hf_repo:
print("[BOOTSTRAP] HF_TOKEN ou HF_DATASET_REPO não configurados — a usar faiss_db local.")
return
faiss_base_dir = os.getenv("FAISS_BASE_DIR", "faiss_db")
try:
from huggingface_hub import snapshot_download
print(f"[BOOTSTRAP] A sincronizar faiss_db/ de {hf_repo}...")
snapshot_download(
repo_id=hf_repo,
repo_type="dataset",
token=hf_token,
local_dir=faiss_base_dir,
local_dir_use_symlinks=False,
ignore_patterns=[".gitattributes", "README.md", ".git*"],
)
print(f"[BOOTSTRAP] faiss_db/ sincronizado com sucesso de {hf_repo}.")
except Exception as exc:
print(f"[BOOTSTRAP] ⚠️ Falha ao sincronizar do HF Hub ({type(exc).__name__}: {exc}). A continuar com faiss_db local...")
def bootstrap():
# Sincronização com HF Hub (best-effort)
_sincronizar_faiss_do_hub()
if not BOOTSTRAP_URL:
print("[BOOTSTRAP] BOOTSTRAP_FAISS_URL não definida. A saltar download.")
return
if any(FAISS_DB_DIR.iterdir()):
print(f"[BOOTSTRAP] Directório {FAISS_DB_DIR} não está vazio. A saltar bootstrap.")
return
print(f"[BOOTSTRAP] A iniciar download de {BOOTSTRAP_URL}...")
temp_file = "/tmp/faiss_bootstrap.archive"
try:
urllib.request.urlretrieve(BOOTSTRAP_URL, temp_file)
print("[BOOTSTRAP] Download concluído. A extrair...")
if tarfile.is_tarfile(temp_file):
with tarfile.open(temp_file) as tar:
tar.extractall(path=FAISS_DB_DIR)
elif zipfile.is_zipfile(temp_file):
with zipfile.open(temp_file) as zip_ref:
zip_ref.extractall(path=FAISS_DB_DIR)
else:
print("[BOOTSTRAP] Erro: Formato de arquivo desconhecido.")
sys.exit(1)
print(f"[BOOTSTRAP] Extração concluída para {FAISS_DB_DIR}")
except Exception as e:
print(f"[BOOTSTRAP] Erro durante o bootstrap: {e}")
# Em produção, podemos decidir se falhamos ou se continuamos vazio
# sys.exit(1)
if __name__ == "__main__":
bootstrap()
|