Spaces:
Sleeping
Sleeping
File size: 4,199 Bytes
07476a1 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """Download models from Hugging Face Hub during Docker build"""
import os
import time
import shutil
from huggingface_hub import snapshot_download, login, hf_hub_download
def download_models():
token = os.environ.get("HF_TOKEN")
if token:
login(token=token)
print("Logged in to Hugging Face")
repo_id = "ridzki-nrzngr/gapura-ai-models"
max_retries = 3
for attempt in range(max_retries):
try:
print(f"Downloading models (attempt {attempt + 1}/{max_retries})...")
cache_dir = snapshot_download(
repo_id=repo_id,
token=token,
resume_download=True,
etag_timeout=30,
max_workers=1,
allow_patterns=[
"regression/*",
"multi_task_transformer.onnx",
"models/nlp/severity_bert/*",
"models/nlp/severity_classifier/*",
],
)
print(f"Models downloaded to: {cache_dir}")
dst_base = "/app/models"
# Copy regression models
regression_files = [
"resolution_predictor_latest.pkl",
"resolution_predictor.ubj",
"resolution_predictor.onnx",
"scaler.json",
"feature_names.json",
"label_encoders.json",
"anomaly_stats.pkl",
]
os.makedirs(os.path.join(dst_base, "regression"), exist_ok=True)
for f in regression_files:
src = os.path.join(cache_dir, "regression", f)
dst = os.path.join(dst_base, "regression", f)
if os.path.exists(src):
shutil.copy2(src, dst)
print(f"Copied: regression/{f}")
# Copy multi-task transformer
onnx_src = os.path.join(cache_dir, "models", "multi_task_transformer.onnx")
if os.path.exists(onnx_src):
shutil.copy2(
onnx_src, os.path.join(dst_base, "multi_task_transformer.onnx")
)
print("Copied: multi_task_transformer.onnx")
# Copy severity_bert (fine-tuned model)
severity_bert_src = os.path.join(
cache_dir, "models", "nlp", "severity_bert"
)
if os.path.exists(severity_bert_src):
dst_sev = os.path.join(dst_base, "nlp", "severity_bert")
os.makedirs(dst_sev, exist_ok=True)
for f in os.listdir(severity_bert_src):
src_file = os.path.join(severity_bert_src, f)
dst_file = os.path.join(dst_sev, f)
if os.path.isfile(src_file):
shutil.copy2(src_file, dst_file)
print(f"Copied: nlp/severity_bert/{f}")
# Copy severity_classifier (TF-IDF fallback)
severity_classifier_src = os.path.join(
cache_dir, "models", "nlp", "severity_classifier"
)
if os.path.exists(severity_classifier_src):
dst_cls = os.path.join(dst_base, "nlp", "severity_classifier")
os.makedirs(dst_cls, exist_ok=True)
for f in os.listdir(severity_classifier_src):
src_file = os.path.join(severity_classifier_src, f)
dst_file = os.path.join(dst_cls, f)
if os.path.isfile(src_file):
shutil.copy2(src_file, dst_file)
print(f"Copied: nlp/severity_classifier/{f}")
print("All models downloaded and copied successfully")
return True
except Exception as e:
print(f"Download failed (attempt {attempt + 1}): {e}")
if attempt < max_retries - 1:
wait_time = (attempt + 1) * 10
print(f"Retrying in {wait_time} seconds...")
time.sleep(wait_time)
else:
print("All retries failed, will use fallback models")
return False
return False
if __name__ == "__main__":
download_models()
|