Spaces:
Building
Building
File size: 2,529 Bytes
0af0f05 | 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 | """
Preload all models into HuggingFace Hub cache at Docker build time.
This avoids cold-start downloads on the first request in production.
"""
from transformers import (
AutoFeatureExtractor,
AutoModelForAudioClassification,
AutoModelForSequenceClassification,
AutoModelForImageClassification,
AutoTokenizer,
)
import sys
MODEL_GROUPS = {
"Audio": [
("AutoFeatureExtractor", "MelodyMachine/Deepfake-audio-detection-V2"),
("AutoModelForAudioClassification", "MelodyMachine/Deepfake-audio-detection-V2"),
],
"Text": [
("AutoTokenizer", "fakespot-ai/roberta-base-ai-text-detection-v1"),
("AutoModelForSequenceClassification", "fakespot-ai/roberta-base-ai-text-detection-v1"),
("AutoTokenizer", "Hello-SimpleAI/chatgpt-detector-roberta"),
("AutoModelForSequenceClassification", "Hello-SimpleAI/chatgpt-detector-roberta"),
("AutoTokenizer", "vikram71198/distilroberta-base-finetuned-fake-news-detection"),
("AutoModelForSequenceClassification", "vikram71198/distilroberta-base-finetuned-fake-news-detection"),
("AutoTokenizer", "jy46604790/Fake-News-Bert-Detect"),
("AutoModelForSequenceClassification", "jy46604790/Fake-News-Bert-Detect"),
],
"Image": [
("AutoModelForImageClassification", "Ateeqq/ai-vs-human-image-detector"),
("AutoModelForImageClassification", "prithivMLmods/AI-vs-Deepfake-vs-Real"),
("AutoModelForImageClassification", "prithivMLmods/Deep-Fake-Detector-Model"),
],
}
LOADERS = {
"AutoFeatureExtractor": AutoFeatureExtractor,
"AutoModelForAudioClassification": AutoModelForAudioClassification,
"AutoModelForSequenceClassification": AutoModelForSequenceClassification,
"AutoModelForImageClassification": AutoModelForImageClassification,
"AutoTokenizer": AutoTokenizer,
}
errors = []
for group, models in MODEL_GROUPS.items():
print(f"\n── {group} ──")
for loader_name, model_name in models:
try:
print(f" Downloading {model_name} ({loader_name})...", end=" ", flush=True)
LOADERS[loader_name].from_pretrained(model_name)
print("OK")
except Exception as e:
print(f"FAILED: {e}")
errors.append((model_name, str(e)))
if errors:
print(f"\n⚠️ {len(errors)} model(s) failed to preload (will download on first request):")
for name, err in errors:
print(f" - {name}: {err}")
else:
print("\nAll models preloaded successfully.")
|