Spaces:
Sleeping
Sleeping
| 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) | |