Spaces:
Runtime error
Runtime error
File size: 10,709 Bytes
6f7e932 | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 | """
Pre-trained Model Loader
Downloads and manages pre-trained models from HuggingFace.
No training required - just load and predict!
Available models:
- Podos Transformer: 276K params, trained on 100K games
- FootballerModel: Classification model for match outcomes
- XGBoost (optional): Loaded from local trained model
"""
import os
import json
import pickle
from pathlib import Path
from typing import Dict, Optional, Any, List
from dataclasses import dataclass
import logging
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Model storage directory
MODELS_DIR = Path(__file__).parent.parent.parent / "models"
PRETRAINED_DIR = MODELS_DIR / "pretrained"
TRAINED_DIR = MODELS_DIR / "trained"
CONFIG_DIR = MODELS_DIR / "config"
# Ensure directories exist
for dir_path in [PRETRAINED_DIR, TRAINED_DIR, CONFIG_DIR]:
dir_path.mkdir(parents=True, exist_ok=True)
@dataclass
class ModelInfo:
"""Information about a loaded model"""
name: str
version: str
source: str # 'huggingface', 'kaggle', 'local'
params: int
loaded: bool
path: Optional[str] = None
error: Optional[str] = None
class PretrainedModelLoader:
"""
Downloads and loads pre-trained models from HuggingFace.
Usage:
loader = PretrainedModelLoader()
loader.download_all()
podos = loader.get_model('podos')
"""
# Model registry with HuggingFace repo info
MODELS = {
'podos': {
'repo_id': 'podos/soccer-match-predictor',
'description': 'Transformer model trained on 100K games',
'params': 276000,
'type': 'transformer',
'fallback_repo': None # Will use mock if not available
},
'footballer': {
'repo_id': 'AmjadKha/FootballerModel',
'description': 'Match outcome classifier',
'params': 50000,
'type': 'classifier',
'fallback_repo': None
}
}
def __init__(self, cache_dir: Optional[Path] = None):
self.cache_dir = cache_dir or PRETRAINED_DIR
self.cache_dir.mkdir(parents=True, exist_ok=True)
self._models: Dict[str, Any] = {}
self._model_info: Dict[str, ModelInfo] = {}
self._hf_available = self._check_huggingface()
def _check_huggingface(self) -> bool:
"""Check if HuggingFace libraries are available"""
try:
import torch
from huggingface_hub import hf_hub_download
return True
except ImportError:
logger.warning("HuggingFace/PyTorch not installed. Using fallback models.")
return False
def download_model(self, model_name: str, force: bool = False) -> bool:
"""
Download a specific model from HuggingFace.
Args:
model_name: Name of model ('podos', 'footballer')
force: Force re-download even if cached
Returns:
True if successful
"""
if model_name not in self.MODELS:
logger.error(f"Unknown model: {model_name}")
return False
model_config = self.MODELS[model_name]
model_path = self.cache_dir / f"{model_name}_model.pt"
# Check cache
if model_path.exists() and not force:
logger.info(f"Model {model_name} already cached at {model_path}")
return True
if not self._hf_available:
logger.warning(f"Creating mock model for {model_name} (HuggingFace not available)")
self._create_mock_model(model_name, model_path)
return True
try:
from huggingface_hub import hf_hub_download, HfHubHTTPError
logger.info(f"Downloading {model_name} from HuggingFace...")
try:
# Try to download from HuggingFace
downloaded_path = hf_hub_download(
repo_id=model_config['repo_id'],
filename="pytorch_model.bin",
cache_dir=self.cache_dir / "hf_cache",
local_dir=self.cache_dir
)
# Copy/move to standard location
import shutil
shutil.copy(downloaded_path, model_path)
logger.info(f"Model {model_name} downloaded to {model_path}")
return True
except Exception as e:
logger.warning(f"Could not download from HuggingFace: {e}")
logger.info(f"Creating local mock model for {model_name}")
self._create_mock_model(model_name, model_path)
return True
except ImportError:
logger.warning("huggingface_hub not installed")
self._create_mock_model(model_name, model_path)
return True
def _create_mock_model(self, model_name: str, model_path: Path):
"""Create a mock model when HuggingFace is unavailable"""
from .mock_models import create_mock_predictor
mock_model = create_mock_predictor(model_name)
# Save mock model info
mock_info = {
'name': model_name,
'type': 'mock',
'version': '0.1.0',
'description': f'Mock {model_name} model (HuggingFace unavailable)'
}
with open(model_path.with_suffix('.json'), 'w') as f:
json.dump(mock_info, f)
# Save the mock model
with open(model_path, 'wb') as f:
pickle.dump(mock_model, f)
logger.info(f"Created mock model at {model_path}")
def download_all(self, force: bool = False) -> Dict[str, bool]:
"""Download all available models"""
results = {}
for model_name in self.MODELS:
results[model_name] = self.download_model(model_name, force)
return results
def load_model(self, model_name: str) -> Optional[Any]:
"""
Load a model into memory.
Args:
model_name: Name of model to load
Returns:
Loaded model object or None
"""
if model_name in self._models:
return self._models[model_name]
model_path = self.cache_dir / f"{model_name}_model.pt"
if not model_path.exists():
logger.info(f"Model {model_name} not found, downloading...")
if not self.download_model(model_name):
return None
try:
# Try PyTorch load first
if self._hf_available:
import torch
try:
model = torch.load(model_path, map_location='cpu')
self._models[model_name] = model
self._model_info[model_name] = ModelInfo(
name=model_name,
version='1.0.0',
source='huggingface',
params=self.MODELS[model_name]['params'],
loaded=True,
path=str(model_path)
)
return model
except Exception:
pass
# Fall back to pickle (mock models)
with open(model_path, 'rb') as f:
model = pickle.load(f)
self._models[model_name] = model
self._model_info[model_name] = ModelInfo(
name=model_name,
version='0.1.0',
source='mock',
params=self.MODELS.get(model_name, {}).get('params', 0),
loaded=True,
path=str(model_path)
)
return model
except Exception as e:
logger.error(f"Error loading model {model_name}: {e}")
self._model_info[model_name] = ModelInfo(
name=model_name,
version='0.0.0',
source='error',
params=0,
loaded=False,
error=str(e)
)
return None
def get_model(self, model_name: str) -> Optional[Any]:
"""Get a loaded model, loading it if necessary"""
if model_name not in self._models:
self.load_model(model_name)
return self._models.get(model_name)
def get_model_info(self, model_name: str) -> Optional[ModelInfo]:
"""Get info about a model"""
if model_name not in self._model_info:
self.load_model(model_name)
return self._model_info.get(model_name)
def list_available_models(self) -> List[Dict]:
"""List all available models with their status"""
models = []
for name, config in self.MODELS.items():
model_path = self.cache_dir / f"{name}_model.pt"
info = self._model_info.get(name)
models.append({
'name': name,
'description': config['description'],
'params': config['params'],
'type': config['type'],
'downloaded': model_path.exists(),
'loaded': name in self._models,
'info': info.__dict__ if info else None
})
return models
def unload_model(self, model_name: str):
"""Unload a model from memory"""
if model_name in self._models:
del self._models[model_name]
logger.info(f"Unloaded model: {model_name}")
def clear_cache(self):
"""Clear all cached models"""
import shutil
if self.cache_dir.exists():
shutil.rmtree(self.cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
self._models.clear()
self._model_info.clear()
logger.info("Model cache cleared")
# Global loader instance
_loader: Optional[PretrainedModelLoader] = None
def get_loader() -> PretrainedModelLoader:
"""Get the global model loader instance"""
global _loader
if _loader is None:
_loader = PretrainedModelLoader()
return _loader
def download_all() -> Dict[str, bool]:
"""Download all pre-trained models"""
return get_loader().download_all()
def get_model(model_name: str) -> Optional[Any]:
"""Get a pre-trained model by name"""
return get_loader().get_model(model_name)
def list_models() -> List[Dict]:
"""List all available models"""
return get_loader().list_available_models()
|