fix torch likelihood imports for missing torch dependency (#2761)

dennisbader · web-flow · commit 513537e0db6e · 2025-04-04T17:36:58.000+02:00
* fix torch likelihood imports for missing torch dependency

* move NotImportedModule to utils/utils.py
diff --git a/darts/explainability/__init__.py b/darts/explainability/__init__.py
@@ -10,7 +10,7 @@
 )
 from darts.explainability.shap_explainer import ShapExplainer
 from darts.logging import get_logger
-from darts.models.utils import NotImportedModule
+from darts.utils.utils import NotImportedModule
 
 logger = get_logger(__name__)
 try:
diff --git a/darts/models/__init__.py b/darts/models/__init__.py
@@ -7,7 +7,7 @@
 
 logger = get_logger(__name__)
 
-from darts.models.utils import NotImportedModule
+from darts.utils.utils import NotImportedModule
 
 try:
     # `lightgbm` needs to be imported first to avoid segmentation fault
diff --git a/darts/models/forecasting/conformal_models.py b/darts/models/forecasting/conformal_models.py
@@ -38,7 +38,6 @@
 from darts.logging import get_logger, raise_log
 from darts.metrics.metrics import METRIC_TYPE
 from darts.models.forecasting.forecasting_model import GlobalForecastingModel
-from darts.models.utils import TORCH_AVAILABLE
 from darts.utils import _build_tqdm_iterator, _with_sanity_checks
 from darts.utils.historical_forecasts.utils import (
     _adjust_historical_forecasts_time_index,
@@ -50,6 +49,7 @@
     series2seq,
 )
 from darts.utils.utils import (
+    TORCH_AVAILABLE,
     _check_quantiles,
     generate_index,
     n_steps_between,
diff --git a/darts/models/forecasting/ensemble_model.py b/darts/models/forecasting/ensemble_model.py
@@ -21,9 +21,9 @@
     GlobalForecastingModel,
     LocalForecastingModel,
 )
-from darts.models.utils import TORCH_AVAILABLE
 from darts.timeseries import TimeSeries, concatenate
 from darts.utils.ts_utils import series2seq
+from darts.utils.utils import TORCH_AVAILABLE
 
 if TORCH_AVAILABLE:
     from darts.models.forecasting.torch_forecasting_model import TorchForecastingModel
diff --git a/darts/models/utils.py b/darts/models/utils.py
diff --git a/darts/tests/explainability/test_shap_explainer.py b/darts/tests/explainability/test_shap_explainer.py
@@ -20,11 +20,11 @@
     ExponentialSmoothing,
     LightGBMModel,
     LinearRegressionModel,
-    NotImportedModule,
     RegressionModel,
     XGBModel,
 )
 from darts.utils.timeseries_generation import linear_timeseries
+from darts.utils.utils import NotImportedModule
 
 lgbm_available = not isinstance(LightGBMModel, NotImportedModule)
 cb_available = not isinstance(CatBoostModel, NotImportedModule)
diff --git a/darts/tests/models/forecasting/test_local_forecasting_models.py b/darts/tests/models/forecasting/test_local_forecasting_models.py
@@ -27,7 +27,6 @@
     NaiveMean,
     NaiveMovingAverage,
     NaiveSeasonal,
-    NotImportedModule,
     Prophet,
     RandomForest,
     RegressionModel,
@@ -44,7 +43,13 @@
 )
 from darts.timeseries import TimeSeries
 from darts.utils import timeseries_generation as tg
-from darts.utils.utils import ModelMode, SeasonalityMode, TrendMode, generate_index
+from darts.utils.utils import (
+    ModelMode,
+    NotImportedModule,
+    SeasonalityMode,
+    TrendMode,
+    generate_index,
+)
 
 logger = get_logger(__name__)
 
diff --git a/darts/tests/models/forecasting/test_probabilistic_models.py b/darts/tests/models/forecasting/test_probabilistic_models.py
@@ -16,11 +16,11 @@
     ExponentialSmoothing,
     LightGBMModel,
     LinearRegressionModel,
-    NotImportedModule,
     XGBModel,
 )
 from darts.tests.conftest import TORCH_AVAILABLE, tfm_kwargs
 from darts.utils import timeseries_generation as tg
+from darts.utils.utils import NotImportedModule
 
 logger = get_logger(__name__)
 
diff --git a/darts/tests/models/forecasting/test_prophet.py b/darts/tests/models/forecasting/test_prophet.py
@@ -6,9 +6,9 @@
 
 from darts import TimeSeries
 from darts.logging import get_logger
-from darts.models import NotImportedModule, Prophet
+from darts.models import Prophet
 from darts.utils import timeseries_generation as tg
-from darts.utils.utils import freqs, generate_index
+from darts.utils.utils import NotImportedModule, freqs, generate_index
 
 logger = get_logger(__name__)
 
diff --git a/darts/tests/models/forecasting/test_regression_models.py b/darts/tests/models/forecasting/test_regression_models.py
@@ -15,6 +15,7 @@
 from sklearn.neighbors import KNeighborsRegressor
 
 import darts
+import darts.utils.utils
 from darts import TimeSeries
 from darts.dataprocessing.encoders import (
     FutureCyclicEncoder,
@@ -26,7 +27,6 @@
     CatBoostModel,
     LightGBMModel,
     LinearRegressionModel,
-    NotImportedModule,
     RandomForest,
     RegressionModel,
     XGBModel,
@@ -35,7 +35,7 @@
 from darts.utils.likelihood_models.base import Likelihood, LikelihoodType
 from darts.utils.likelihood_models.sklearn import _get_likelihood
 from darts.utils.multioutput import MultiOutputRegressor
-from darts.utils.utils import generate_index
+from darts.utils.utils import NotImportedModule, generate_index
 
 logger = get_logger(__name__)
 
@@ -3422,7 +3422,7 @@ def test_get_categorical_features_helper(self):
         (
             darts.models.forecasting.lgbm.lgb.LGBMRegressor
             if lgbm_available
-            else darts.models.utils.NotImportedModule
+            else darts.utils.utils.NotImportedModule
         ),
         "fit",
     )
diff --git a/darts/tests/utils/historical_forecasts/test_historical_forecasts.py b/darts/tests/utils/historical_forecasts/test_historical_forecasts.py
@@ -27,7 +27,6 @@
     LinearRegressionModel,
     NaiveDrift,
     NaiveSeasonal,
-    NotImportedModule,
     StatsForecastAutoARIMA,
 )
 from darts.models.forecasting.forecasting_model import (
@@ -41,6 +40,7 @@
     quantile_names,
 )
 from darts.utils.ts_utils import SeriesType, get_series_seq_type
+from darts.utils.utils import NotImportedModule
 
 if TORCH_AVAILABLE:
     import torch
diff --git a/darts/utils/data/__init__.py b/darts/utils/data/__init__.py
@@ -44,7 +44,7 @@
         TrainingDataset,
     )
 except ImportError:  # Torch is not available
-    from darts.models.utils import NotImportedModule
+    from darts.utils.utils import NotImportedModule
 
     HorizonBasedDataset = NotImportedModule(module_name="(Py)Torch", warn=False)
     DualCovariatesInferenceDataset = NotImportedModule(
diff --git a/darts/utils/likelihood_models/__init__.py b/darts/utils/likelihood_models/__init__.py
@@ -3,26 +3,58 @@
 -----------------
 """
 
+from darts.logging import get_logger
+from darts.utils.utils import NotImportedModule
+
+logger = get_logger(__name__)
+
 # allow direct import of torch likelihoods as they are the only ones useful to the user
-from darts.utils.likelihood_models.torch import (
-    BernoulliLikelihood,
-    BetaLikelihood,
-    CauchyLikelihood,
-    ContinuousBernoulliLikelihood,
-    DirichletLikelihood,
-    ExponentialLikelihood,
-    GammaLikelihood,
-    GaussianLikelihood,
-    GeometricLikelihood,
-    GumbelLikelihood,
-    HalfNormalLikelihood,
-    LaplaceLikelihood,
-    LogNormalLikelihood,
-    NegativeBinomialLikelihood,
-    PoissonLikelihood,
-    QuantileRegression,
-    WeibullLikelihood,
-)
+try:
+    from darts.utils.likelihood_models.torch import (
+        BernoulliLikelihood,
+        BetaLikelihood,
+        CauchyLikelihood,
+        ContinuousBernoulliLikelihood,
+        DirichletLikelihood,
+        ExponentialLikelihood,
+        GammaLikelihood,
+        GaussianLikelihood,
+        GeometricLikelihood,
+        GumbelLikelihood,
+        HalfNormalLikelihood,
+        LaplaceLikelihood,
+        LogNormalLikelihood,
+        NegativeBinomialLikelihood,
+        PoissonLikelihood,
+        QuantileRegression,
+        WeibullLikelihood,
+    )
+except ModuleNotFoundError:
+    logger.warning(
+        "Support for PyTorch based likelihood models not available. "
+        'To enable them, install "darts", "u8darts[torch]" or "u8darts[all]" (with pip); '
+        'or "u8darts-torch" or "u8darts-all" (with conda).'
+    )
+    BernoulliLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    BetaLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    CauchyLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    ContinuousBernoulliLikelihood = NotImportedModule(
+        module_name="(Py)Torch", warn=False
+    )
+    DirichletLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    ExponentialLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    GammaLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    GaussianLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    GeometricLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    GumbelLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    HalfNormalLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    LaplaceLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    LogNormalLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    NegativeBinomialLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    PoissonLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+    QuantileRegression = NotImportedModule(module_name="(Py)Torch", warn=False)
+    WeibullLikelihood = NotImportedModule(module_name="(Py)Torch", warn=False)
+
 
 __all__ = [
     "GaussianLikelihood",
diff --git a/darts/utils/utils.py b/darts/utils/utils.py
@@ -19,11 +19,20 @@
 
 from darts.logging import get_logger, raise_if, raise_if_not, raise_log
 
+logger = get_logger(__name__)
+
 try:
     from IPython import get_ipython
 except ModuleNotFoundError:
     get_ipython = None
 
+try:
+    import torch  # noqa: F401
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+
 logger = get_logger(__name__)
 
 MAX_TORCH_SEED_VALUE = (1 << 31) - 1  # to accommodate 32-bit architectures
@@ -55,6 +64,25 @@
 ]
 
 
+class NotImportedModule:
+    """Helper class for handling import errors of optional dependencies."""
+
+    usable = False
+
+    def __init__(self, module_name: str, warn: bool = True):
+        self.error_message = (
+            f"The `{module_name}` module could not be imported. "
+            f"To enable {module_name} support in Darts, follow the detailed "
+            f"instructions in the installation guide: "
+            f"https://github.com/unit8co/darts/blob/master/INSTALL.md"
+        )
+        if warn:
+            logger.warning(self.error_message)
+
+    def __call__(self, *args, **kwargs):
+        raise_log(ImportError(self.error_message), logger=logger)
+
+
 # Enums
 class SeasonalityMode(Enum):
     MULTIPLICATIVE = "multiplicative"

Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@`
`10`	`10`	`)`
`11`	`11`	`from darts.explainability.shap_explainer import ShapExplainer`
`12`	`12`	`from darts.logging import get_logger`
`13`		`-from darts.models.utils import NotImportedModule`
	`13`	`+from darts.utils.utils import NotImportedModule`
`14`	`14`
`15`	`15`	`logger = get_logger(__name__)`
`16`	`16`	`try:`