Spaces:
Sleeping
Sleeping
File size: 1,721 Bytes
9831ced 4554e0a 9831ced | 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 | import threading
from collections import OrderedDict
from models.base_model import BaseForecastModel
_ORDER_SEARCH_MODELS = {'ARIMA', 'Auto-ARIMA', 'ARIMA-GARCH'}
_CACHE = OrderedDict()
_MAX_CACHE = 32
_CACHE_LOCK = threading.Lock()
def _build_model(name: str) -> BaseForecastModel:
if name == 'ARIMA':
from models.arima_model import ArimaModel
return ArimaModel()
elif name == 'Auto-ARIMA':
from models.auto_arima_model import AutoArimaModel
return AutoArimaModel()
elif name == 'ARIMA-GARCH':
from models.arima_garch_model import ArimaGarchModel
return ArimaGarchModel()
elif name == 'Moirai':
from models.moirai_model import MoiraiModel
return MoiraiModel()
elif name == 'TimesFM':
from models.timesfm_model import TimesFMModel
return TimesFMModel()
elif name == 'Ridge':
from models.ridge_model import RidgeModel
return RidgeModel()
raise ValueError(f"Unknown model '{name}'. Expected one of: ARIMA, Auto-ARIMA, ARIMA-GARCH, Moirai, TimesFM, Ridge.")
def get_model(name: str, symbol: str=None, timeframe: str=None, use_features: bool=False) -> BaseForecastModel:
if name in _ORDER_SEARCH_MODELS:
variant = 'featured' if use_features else 'base'
key = (name, symbol, timeframe, variant)
else:
key = name
with _CACHE_LOCK:
if key in _CACHE:
_CACHE.move_to_end(key)
return _CACHE[key]
model = _build_model(name)
if len(_CACHE) >= _MAX_CACHE:
_CACHE.popitem(last=False)
_CACHE[key] = model
return model
def fresh_model(name: str) -> BaseForecastModel:
return _build_model(name)
|