bug: try-except wrappers with TabPFN imports (#1403)

* Added try-except wrappers to TabPFN imports

* Automated autopep8 fixes

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: martilut

fedot/core/operations/evaluation/tabpfn.py

@@ -1,13 +1,27 @@
+import logging
 from typing import Optional
 
 from fedot.core.data.data import InputData, OutputData
 from fedot.core.operations.evaluation.evaluation_interfaces import EvaluationStrategy
-from fedot.core.operations.evaluation.operation_implementations.models.tabpfn import \
-    FedotTabPFNClassificationImplementation, FedotTabPFNRegressionImplementation
 from fedot.core.operations.operation_parameters import OperationParameters
 from fedot.core.repository.tasks import TaskTypesEnum
 from fedot.utilities.random import ImplementationRandomStateHandler
 
+try:
+    from fedot.core.operations.evaluation.operation_implementations.models.tabpfn import \
+        FedotTabPFNClassificationImplementation, FedotTabPFNRegressionImplementation
+    _TABPFN_AVAILABLE = True
+except ModuleNotFoundError:
+    FedotTabPFNClassificationImplementation = None
+    FedotTabPFNRegressionImplementation = None
+    _ERROR_MESSAGE = (""
+                      "TabPFN is required but not installed. "
+                      "Install with `pip install fedot[extra]` or `pip install tabpfn`."
+                      ""
+                      )
+    _TABPFN_AVAILABLE = False
+    logging.log(100, _ERROR_MESSAGE)
+
 
 class TabPFNStrategy(EvaluationStrategy):
     _operations_by_types = {
@@ -23,6 +37,9 @@ def __init__(self, operation_type: str, params: Optional[OperationParameters] =
         self.max_features = params.get('max_features', 500) if params else 500
 
     def fit(self, train_data: InputData):
+        if not _TABPFN_AVAILABLE:
+            raise ModuleNotFoundError(_ERROR_MESSAGE)
+
         check_data_size(
             data=train_data,
             device=self.device,
@@ -48,6 +65,9 @@ def __init__(self, operation_type: str, params: Optional[OperationParameters] =
         super().__init__(operation_type, params)
 
     def predict(self, trained_operation, predict_data: InputData) -> OutputData:
+        if not _TABPFN_AVAILABLE:
+            raise ModuleNotFoundError(_ERROR_MESSAGE)
+
         if self.output_mode == 'labels':
             output = trained_operation.predict(predict_data)
         elif self.output_mode in ['probs', 'full_probs', 'default']:
@@ -69,6 +89,9 @@ def __init__(self, operation_type: str, params: Optional[OperationParameters] =
         super().__init__(operation_type, params)
 
     def predict(self, trained_operation, predict_data: InputData) -> OutputData:
+        if not _TABPFN_AVAILABLE:
+            raise ModuleNotFoundError(_ERROR_MESSAGE)
+
         return trained_operation.predict(predict_data)
 
 

test/integration/models/test_strategy.py

@@ -85,16 +85,19 @@ def run_tabpfn(
         test_data: pd.DataFrame,
         task: str,
 ):
-    pipeline = PipelineBuilder().add_node(model_name).build()
-    pipeline.fit(train_data)
-    predicted_output = pipeline.predict(test_data, output_mode='labels')
-    if task == 'classification':
-        metric = roc_auc(test_data.target, predicted_output.predict)
-    else:
-        metric = r2_score(test_data.target, predicted_output.predict)
+    try:
+        pipeline = PipelineBuilder().add_node(model_name).build()
+        pipeline.fit(train_data)
+        predicted_output = pipeline.predict(test_data, output_mode='labels')
+        if task == 'classification':
+            metric = roc_auc(test_data.target, predicted_output.predict)
+        else:
+            metric = r2_score(test_data.target, predicted_output.predict)
 
-    assert isinstance(pipeline, Pipeline)
-    assert metric > 0.5
+        assert isinstance(pipeline, Pipeline)
+        assert metric > 0.5
+    except ModuleNotFoundError:
+        pass
 
 
 def test_tabpfn_classification_operation():