aimodelshare 0.1.54__py3-none-any.whl → 0.1.59__py3-none-any.whl

aimodelshare/__init__.py

@@ -1,18 +1,98 @@
-from .playground import ModelPlayground, Competition,Experiment, Data
-from .preprocessormodules import export_preprocessor,upload_preprocessor,import_preprocessor
-from .data_sharing.download_data import download_data, import_quickstart_data
-from .reproducibility import export_reproducibility_env, import_reproducibility_env
+"""
+Top-level aimodelshare package initializer.
 
+Refactored to:
+- Avoid eager importing of heavy optional dependencies (networkx, torch, tensorflow, etc.)
+- Allow lightweight submodules (e.g. aimodelshare.moral_compass) to work in minimal environments.
+- Preserve backward compatibility where possible without forcing heavy installs.
+"""
+
+from importlib import import_module
+import warnings
 
 __all__ = [
-    # Object Oriented
-    ModelPlayground,
-    Competition,
-    Data,
-    Experiment,
-    # Preprocessor
-    upload_preprocessor,
-    import_preprocessor,
-    export_preprocessor,
-    download_data
+    # Heavy/high-level objects (added only if imported successfully)
+    "ModelPlayground",
+    "Competition",
+    "Data",
+    "Experiment",
+    # Preprocessor helpers (optional)
+    "upload_preprocessor",
+    "import_preprocessor",
+    "export_preprocessor",
+    "download_data",
+    # Moral Compass client (lightweight)
+    "MoralcompassApiClient",
+    "MoralcompassTableMeta",
+    "MoralcompassUserStats",
+    "ApiClientError",
+    "NotFoundError",
+    "ServerError",
+    "get_api_base_url",
 ]
+
+def _safe_import(module_name, symbols, remove_if_missing=True):
+    """
+    Attempt to import given symbols from module_name.
+    On failure, emit a warning and optionally remove them from __all__.
+    """
+    try:
+        mod = import_module(module_name)
+        for sym in symbols:
+            try:
+                globals()[sym] = getattr(mod, sym)
+            except AttributeError:
+                warnings.warn(
+                    f"aimodelshare: symbol '{sym}' not found in module '{module_name}'."
+                )
+                if remove_if_missing and sym in __all__:
+                    __all__.remove(sym)
+    except Exception as exc:
+        warnings.warn(
+            f"aimodelshare: optional module '{module_name}' not loaded ({exc}). "
+            "Lightweight submodules remain usable."
+        )
+        if remove_if_missing:
+            for sym in symbols:
+                if sym in __all__ and sym not in globals():
+                    __all__.remove(sym)
+
+# Attempt optional imports (silently skip if deps missing)
+_safe_import(
+    "aimodelshare.playground",
+    ["ModelPlayground", "Competition", "Experiment", "Data"],
+)
+
+_safe_import(
+    "aimodelshare.preprocessormodules",
+    ["export_preprocessor", "upload_preprocessor", "import_preprocessor"],
+)
+
+_safe_import(
+    "aimodelshare.data_sharing.download_data",
+    ["download_data", "import_quickstart_data"],  # import_quickstart_data not in __all__
+    remove_if_missing=True,
+)
+# If import_quickstart_data is missing, we don't expose it; if present we leave it accessible.
+if "import_quickstart_data" not in globals():
+    # Ensure it's not accidentally in __all__
+    if "import_quickstart_data" in __all__:
+        __all__.remove("import_quickstart_data")
+
+# Moral Compass submodule (expected always present in new branch)
+_safe_import(
+    "aimodelshare.moral_compass",
+    [
+        "MoralcompassApiClient",
+        "MoralcompassTableMeta",
+        "MoralcompassUserStats",
+        "ApiClientError",
+        "NotFoundError",
+        "ServerError",
+        "get_api_base_url",
+    ],
+    remove_if_missing=False,  # If this fails, keep names so import errors surface clearly later
+)
+
+# Ensure __all__ only contains names actually available (except intentional exposure for debug)
+__all__ = [name for name in __all__ if name in globals()]

aimodelshare/aimsonnx.py

@@ -2,28 +2,31 @@
 import pandas as pd
 import numpy as np 
 
+# Import optional dependency checker
+from aimodelshare.utils.optional_deps import check_optional
+
 # ml frameworks
 try:
     import sklearn
     from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
 except:
-    print("Warning: Please install sklearn to enable sklearn features")
+    check_optional("sklearn", "Scikit-learn")
 
 try:
     import torch
 except:
-    print("Warning: Please install pytorch to enable pytorch features")
+    check_optional("torch", "PyTorch")
 
 try:
     import xgboost
 except:
-    print("Warning: Please install xgboost to enable xgboost features")
+    check_optional("xgboost", "XGBoost")
 
 try:
     import tensorflow as tf
     import keras
 except:
-    print("Warning: Please install tensorflow/keras to enable tensorflow/keras features")
+    check_optional("tensorflow", "TensorFlow/Keras")
 
 try:
     import pyspark
@@ -32,14 +35,17 @@ try:
     from pyspark.ml.tuning import CrossValidatorModel, TrainValidationSplitModel
     from onnxmltools import convert_sparkml
 except:
-    print("Warning: Please install pyspark to enable pyspark features")
+    check_optional("pyspark", "PySpark")
 
 
 # onnx modules
 import onnx
 import skl2onnx
 from skl2onnx import convert_sklearn
-import tf2onnx
+# tf2onnx import is lazy-loaded to avoid requiring TensorFlow for non-TF workflows
+_TF2ONNX_AVAILABLE = None
+_tf2onnx_module = None
+_tensorflow_module = None
 try:
     from torch.onnx import export
 except:
@@ -71,18 +77,59 @@ import wget
 from copy import copy
 import psutil
 from pympler import asizeof
-from IPython.core.display import display, HTML, SVG
+from IPython.display import display, HTML, SVG
 import absl.logging
 import networkx as nx
 import warnings
 from pathlib import Path
 import time
 import signal
-from scikeras.wrappers import KerasClassifier, KerasRegressor
+
+# scikeras imports keras which requires TensorFlow - lazy load it
+try:
+    from scikeras.wrappers import KerasClassifier, KerasRegressor
+    _SCIKERAS_AVAILABLE = True
+except ImportError:
+    _SCIKERAS_AVAILABLE = False
+    KerasClassifier = None
+    KerasRegressor = None
 
 
 absl.logging.set_verbosity(absl.logging.ERROR)
 
+def _check_tf2onnx_available():
+    """Check if tf2onnx and TensorFlow are available, and load them if needed.
+    
+    Returns:
+        tuple: (tf2onnx_module, tensorflow_module) on success
+        
+    Raises:
+        RuntimeError: If TensorFlow or tf2onnx are not installed
+    """
+    global _TF2ONNX_AVAILABLE, _tf2onnx_module, _tensorflow_module
+    
+    if _TF2ONNX_AVAILABLE is None:
+        try:
+            import tf2onnx as tf2onnx_temp
+            import tensorflow as tf_temp
+            _tf2onnx_module = tf2onnx_temp
+            _tensorflow_module = tf_temp
+            _TF2ONNX_AVAILABLE = True
+        except ImportError as e:
+            _TF2ONNX_AVAILABLE = False
+            raise RuntimeError(
+                "TensorFlow and tf2onnx are required for Keras model conversion to ONNX. "
+                "Please install them with: pip install tensorflow tf2onnx"
+            ) from e
+    
+    if not _TF2ONNX_AVAILABLE:
+        raise RuntimeError(
+            "TensorFlow and tf2onnx are required for Keras model conversion to ONNX. "
+            "Please install them with: pip install tensorflow tf2onnx"
+        )
+    
+    return _tf2onnx_module, _tensorflow_module
+
 def _extract_onnx_metadata(onnx_model, framework):
     '''Extracts model metadata from ONNX file.'''
 
@@ -265,28 +312,8 @@ def _sklearn_to_onnx(model, initial_types=None, transfer_learning=None,
 
     onx = convert_sklearn(model, initial_types=initial_types,target_opset={'': 15, 'ai.onnx.ml': 2})
     
-    ## Dynamically set model ir_version to ensure sklearn opsets work properly
-    from onnx.helper import VERSION_TABLE
-    import onnx
-    import numpy as np
-
-    indexlocationlist=[]
-    for i in VERSION_TABLE:
-      indexlocationlist.append(str(i).find(str(onnx.__version__)))
-
-
-    arr = np.array(indexlocationlist)
-
-    def condition(x): return x > -1
-
-    bool_arr = condition(arr)
-
-    output = np.where(bool_arr)[0]
-
-    ir_version=VERSION_TABLE[output[0]][1]
-
-    #add to model object before saving
-    onx.ir_version = ir_version
+    ## set model ir_version to ensure sklearn opsets work properly
+    onx.ir_version = 8
     
     # generate metadata dict 
     metadata = {}
@@ -555,8 +582,9 @@ def _keras_to_onnx(model, transfer_learning=None,
                    deep_learning=None, task_type=None, epochs=None):
     '''Converts a Keras model to ONNX and extracts metadata.'''
 
-    import tf2onnx
-    import tensorflow as tf
+    # Check and load tf2onnx and TensorFlow lazily (only when needed)
+    tf2onnx, tf = _check_tf2onnx_available()
+    
     import numpy as np
     import onnx
     import pickle
@@ -934,7 +962,7 @@ def model_to_onnx(model, framework=None, model_input=None, initial_types=None,
             from pyspark.ml.tuning import CrossValidatorModel, TrainValidationSplitModel
             from onnxmltools import convert_sparkml
         except:
-            print("Warning: Please install pyspark to enable pyspark features")
+            check_optional("pyspark", "PySpark")
         onnx = _pyspark_to_onnx(model, initial_types=initial_types, 
                                 transfer_learning=transfer_learning, 
                                 deep_learning=deep_learning, 
@@ -989,23 +1017,39 @@ def model_to_onnx_timed(model_filepath, force_onnx=False, timeout=60, model_inpu
 
             except:
                 print("Timeout: Model to ONNX conversion is taking longer than expected. This can be the case for big models.")
-                response = ''
-                while response not in {"1", "2"}:
-                    response = input("Do you want to keep trying (1) or submit predictions only (2)? ")
-
-                if response == "1":
-                    try:
-                        import torch
-                        if isinstance(model_filepath, torch.nn.Module):
+                
+                # Detect CI/testing environment for non-interactive fallback
+                is_non_interactive = (
+                    os.environ.get("PYTEST_CURRENT_TEST") is not None or
+                    os.environ.get("AIMS_NON_INTERACTIVE") == "1"
+                )
+                
+                if is_non_interactive:
+                    # Auto-fallback to predictions-only in CI/testing environment
+                    print("Non-interactive environment detected. Falling back to predictions-only submission.")
+                    model_filepath = None
+                else:
+                    # Interactive prompt for manual runs
+                    response = ''
+                    while response not in {"1", "2"}:
+                        response = input("Do you want to keep trying (1) or submit predictions only (2)? ")
+
+                    if response == "1":
+                        try:
+                            import torch
+                            if isinstance(model_filepath, torch.nn.Module):
+                                onnx_model = model_to_onnx(model_filepath, model_input=model_input)
+                            else:
+                                onnx_model = model_to_onnx(model_filepath)
+                        except Exception as e:
+                            # Final fallback - if torch-specific handling failed, try generic conversion
+                            # This handles cases where torch module detection fails but conversion might still work
+                            warnings.warn(f"PyTorch-specific ONNX conversion failed ({e}), attempting generic conversion")
                             onnx_model = model_to_onnx(model_filepath, model_input=model_input)
-                        else:
-                            onnx_model = model_to_onnx(model_filepath)
-                    except:
-                        onnx_model = model_to_onnx(model_filepath)
-                    model_filepath = onnx_model
+                        model_filepath = onnx_model
 
-                elif response == "2":
-                    model_filepath = None
+                    elif response == "2":
+                        model_filepath = None
 
             finally:
                 print()
@@ -1024,6 +1068,12 @@ def _get_metadata(onnx_model):
     #assert(isinstance(onnx_model, onnx.onnx_ml_pb2.ModelProto)), \
      #"Please pass a onnx model object."
     
+    # Handle None input gracefully - always return a dict
+    if onnx_model is None:
+        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+            print("[DEBUG] _get_metadata: onnx_model is None, returning empty dict")
+        return {}
+    
     try: 
         onnx_meta = onnx_model.metadata_props
 
@@ -1034,36 +1084,121 @@ def _get_metadata(onnx_model):
 
         onnx_meta_dict = ast.literal_eval(onnx_meta_dict['model_metadata'])
         
+        # Handle case where metadata is stored as a list instead of dict
+        if isinstance(onnx_meta_dict, list):
+            if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+                print(f"[DEBUG] _get_metadata: metadata is a list of length {len(onnx_meta_dict)}")
+            if len(onnx_meta_dict) > 0 and isinstance(onnx_meta_dict[0], dict):
+                onnx_meta_dict = onnx_meta_dict[0]
+                if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+                    print("[DEBUG] _get_metadata: Extracted first dict from list")
+            else:
+                # Return empty dict if list doesn't contain valid dicts
+                if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+                    print("[DEBUG] _get_metadata: List does not contain valid dicts, returning empty dict")
+                return {}
+        
+        # Ensure we have a dict at this point
+        if not isinstance(onnx_meta_dict, dict):
+            if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+                print(f"[DEBUG] _get_metadata: Unexpected metadata type {type(onnx_meta_dict)}, returning empty dict")
+            return {}
+        
         #if onnx_meta_dict['model_config'] != None and \
         #onnx_meta_dict['ml_framework'] != 'pytorch':
         #    onnx_meta_dict['model_config'] = ast.literal_eval(onnx_meta_dict['model_config'])
         
-        if onnx_meta_dict['model_architecture'] != None:
-            onnx_meta_dict['model_architecture'] = ast.literal_eval(onnx_meta_dict['model_architecture'])
+        # Attempt to parse nested fields only if they are string representations of dicts
+        if 'model_architecture' in onnx_meta_dict and onnx_meta_dict['model_architecture'] != None:
+            try:
+                if isinstance(onnx_meta_dict['model_architecture'], str):
+                    onnx_meta_dict['model_architecture'] = ast.literal_eval(onnx_meta_dict['model_architecture'])
+            except (ValueError, SyntaxError):
+                # Keep as-is if parsing fails
+                pass
+            
+        if 'model_config' in onnx_meta_dict and onnx_meta_dict['model_config'] != None:
+            try:
+                if isinstance(onnx_meta_dict['model_config'], str):
+                    onnx_meta_dict['model_config'] = ast.literal_eval(onnx_meta_dict['model_config'])
+            except (ValueError, SyntaxError):
+                # Keep as-is if parsing fails
+                pass
             
-        if onnx_meta_dict['metadata_onnx'] != None:
-            onnx_meta_dict['metadata_onnx'] = ast.literal_eval(onnx_meta_dict['metadata_onnx'])
+        if 'metadata_onnx' in onnx_meta_dict and onnx_meta_dict['metadata_onnx'] != None:
+            try:
+                if isinstance(onnx_meta_dict['metadata_onnx'], str):
+                    onnx_meta_dict['metadata_onnx'] = ast.literal_eval(onnx_meta_dict['metadata_onnx'])
+            except (ValueError, SyntaxError):
+                # Keep as-is if parsing fails
+                pass
         
         # onnx_meta_dict['model_image'] = onnx_to_image(onnx_model)
 
     except Exception as e:
     
-        print(e)
+        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+            print(f"[DEBUG] _get_metadata: Exception during metadata extraction: {e}")
         
-        onnx_meta_dict = ast.literal_eval(onnx_meta_dict)
+        try:
+            onnx_meta_dict = ast.literal_eval(onnx_meta_dict)
+            # Handle list case in exception path as well
+            if isinstance(onnx_meta_dict, list) and len(onnx_meta_dict) > 0 and isinstance(onnx_meta_dict[0], dict):
+                onnx_meta_dict = onnx_meta_dict[0]
+            elif not isinstance(onnx_meta_dict, dict):
+                onnx_meta_dict = {}
+        except:
+            onnx_meta_dict = {}
+    
+    # Final safety check: ensure we always return a dict
+    if not isinstance(onnx_meta_dict, dict):
+        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+            print(f"[DEBUG] _get_metadata: Final check failed, returning empty dict instead of {type(onnx_meta_dict)}")
+        return {}
         
     return onnx_meta_dict
 
 
 
 def _get_leaderboard_data(onnx_model, eval_metrics=None):
+    '''Extract leaderboard data from ONNX model or return defaults.
     
+    This function performs single-pass normalization and safely handles:
+    - None onnx_model (returns defaults)
+    - Invalid metadata structures
+    - Missing keys in metadata
+    '''
+    
+    # Start with eval_metrics if provided, otherwise empty dict
     if eval_metrics is not None:
-        metadata = eval_metrics
+        metadata = dict(eval_metrics) if isinstance(eval_metrics, dict) else {}
     else:
-        metadata = dict()
+        metadata = {}
+    
+    # Handle None onnx_model gracefully
+    if onnx_model is None:
+        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+            print("[DEBUG] _get_leaderboard_data: onnx_model is None, using default metadata")
+        # Return metadata with safe defaults injected
+        metadata['ml_framework'] = metadata.get('ml_framework', None)
+        metadata['transfer_learning'] = metadata.get('transfer_learning', None)
+        metadata['deep_learning'] = metadata.get('deep_learning', None)
+        metadata['model_type'] = metadata.get('model_type', None)
+        metadata['depth'] = metadata.get('depth', 0)
+        metadata['num_params'] = metadata.get('num_params', 0)
+        return metadata
         
+    # Get metadata from ONNX - _get_metadata now always returns a dict
     metadata_raw = _get_metadata(onnx_model)
+    
+    if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+        print(f"[DEBUG] _get_leaderboard_data: metadata_raw type={type(metadata_raw)}, keys={list(metadata_raw.keys()) if isinstance(metadata_raw, dict) else 'N/A'}")
+    
+    # Single-pass normalization: ensure metadata_raw is a dict
+    if not isinstance(metadata_raw, dict):
+        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+            print(f"[DEBUG] _get_leaderboard_data: metadata_raw is not a dict (type={type(metadata_raw)}), using empty dict")
+        metadata_raw = {}
 
     # get list of current layer types 
     layer_list_keras, activation_list_keras = _get_layer_names()
@@ -1072,46 +1207,55 @@ def _get_leaderboard_data(onnx_model, eval_metrics=None):
     layer_list = list(set(layer_list_keras + layer_list_pytorch))
     activation_list =  list(set(activation_list_keras + activation_list_pytorch))
 
-    # get general model info
-    metadata['ml_framework'] = metadata_raw['ml_framework']
-    metadata['transfer_learning'] = metadata_raw['transfer_learning']
-    metadata['deep_learning'] = metadata_raw['deep_learning']
-    metadata['model_type'] = metadata_raw['model_type']
+    # get general model info - use .get() for safety
+    metadata['ml_framework'] = metadata_raw.get('ml_framework')
+    metadata['transfer_learning'] = metadata_raw.get('transfer_learning')
+    metadata['deep_learning'] = metadata_raw.get('deep_learning')
+    metadata['model_type'] = metadata_raw.get('model_type')
 
 
     # get neural network metrics
-    if metadata_raw['ml_framework'] in ['keras', 'pytorch'] or metadata_raw['model_type'] in ['MLPClassifier', 'MLPRegressor']:
-        metadata['depth'] = metadata_raw['model_architecture']['layers_number']
-        metadata['num_params'] = sum(metadata_raw['model_architecture']['layers_n_params'])
+    # Add isinstance check for model_architecture to prevent TypeError
+    if (metadata_raw.get('ml_framework') in ['keras', 'pytorch'] or 
+        metadata_raw.get('model_type') in ['MLPClassifier', 'MLPRegressor']) and \
+        isinstance(metadata_raw.get('model_architecture'), dict):
+        
+        metadata['depth'] = metadata_raw['model_architecture'].get('layers_number', 0)
+        metadata['num_params'] = sum(metadata_raw['model_architecture'].get('layers_n_params', []))
 
         for i in layer_list:
-            if i in metadata_raw['model_architecture']['layers_summary']:
-                metadata[i.lower()+'_layers'] = metadata_raw['model_architecture']['layers_summary'][i]
+            layers_summary = metadata_raw['model_architecture'].get('layers_summary', {})
+            if i in layers_summary:
+                metadata[i.lower()+'_layers'] = layers_summary[i]
             elif i.lower()+'_layers' not in metadata.keys():
                 metadata[i.lower()+'_layers'] = 0
 
         for i in activation_list:
-            if i in metadata_raw['model_architecture']['activations_summary']:
+            activations_summary = metadata_raw['model_architecture'].get('activations_summary', {})
+            if i in activations_summary:
                 if i.lower()+'_act' in metadata:
-                    metadata[i.lower()+'_act'] += metadata_raw['model_architecture']['activations_summary'][i]
+                    metadata[i.lower()+'_act'] += activations_summary[i]
                 else:    
-                    metadata[i.lower()+'_act'] = metadata_raw['model_architecture']['activations_summary'][i]
+                    metadata[i.lower()+'_act'] = activations_summary[i]
             else:
                 if i.lower()+'_act' not in metadata:
                     metadata[i.lower()+'_act'] = 0
                 
-        metadata['loss'] = metadata_raw['model_architecture']['loss']
-        metadata['optimizer'] = metadata_raw['model_architecture']["optimizer"]
-        metadata['model_config'] = metadata_raw['model_config']
-        metadata['epochs'] = metadata_raw['epochs']
-        metadata['memory_size'] = metadata_raw['memory_size']
+        metadata['loss'] = metadata_raw['model_architecture'].get('loss')
+        metadata['optimizer'] = metadata_raw['model_architecture'].get('optimizer')
+        metadata['model_config'] = metadata_raw.get('model_config')
+        metadata['epochs'] = metadata_raw.get('epochs')
+        metadata['memory_size'] = metadata_raw.get('memory_size')
 
     # get sklearn & pyspark model metrics
-    elif metadata_raw['ml_framework'] in ['sklearn', 'xgboost', 'pyspark']:
+    elif metadata_raw.get('ml_framework') in ['sklearn', 'xgboost', 'pyspark']:
         metadata['depth'] = 0
 
         try:
-            metadata['num_params'] = sum(metadata_raw['model_architecture']['layers_n_params'])
+            if isinstance(metadata_raw.get('model_architecture'), dict):
+                metadata['num_params'] = sum(metadata_raw['model_architecture'].get('layers_n_params', []))
+            else:
+                metadata['num_params'] = 0
         except:
             metadata['num_params'] = 0
 
@@ -1124,15 +1268,29 @@ def _get_leaderboard_data(onnx_model, eval_metrics=None):
         metadata['loss'] = None
 
         try:
-            metadata['optimizer'] = metadata_raw['model_architecture']['optimizer']
+            if isinstance(metadata_raw.get('model_architecture'), dict):
+                metadata['optimizer'] = metadata_raw['model_architecture'].get('optimizer')
+            else:
+                metadata['optimizer'] = None
         except:
             metadata['optimizer'] = None
 
         try:
-            metadata['model_config'] = metadata_raw['model_config']
+            metadata['model_config'] = metadata_raw.get('model_config')
         except:
             metadata['model_config'] = None
     
+    # Default handling for unknown frameworks
+    else:
+        if os.environ.get("AIMODELSHARE_DEBUG_METADATA"):
+            print(f"[DEBUG] _get_leaderboard_data: Unknown framework '{metadata_raw.get('ml_framework')}', using defaults")
+        metadata.setdefault('depth', 0)
+        metadata.setdefault('num_params', 0)
+        for i in layer_list:
+            metadata.setdefault(i.lower()+'_layers', 0)
+        for i in activation_list:
+            metadata.setdefault(i.lower()+'_act', 0)
+    
     return metadata
     
 
@@ -1553,7 +1711,8 @@ def _get_sklearn_modules():
 
     sklearn_modules = ['ensemble', 'gaussian_process', 'isotonic',
                        'linear_model', 'mixture', 'multiclass', 'naive_bayes',
-                       'neighbors', 'neural_network', 'svm', 'tree']
+                       'neighbors', 'neural_network', 'svm', 'tree',
+                       'discriminant_analysis', 'calibration']
 
     models_modules_dict = {}
 
@@ -1569,9 +1728,31 @@ def _get_sklearn_modules():
 
 def model_from_string(model_type):
     models_modules_dict = _get_sklearn_modules()
-    module = models_modules_dict[model_type]
-    model_class = getattr(importlib.import_module(module), model_type)
-    return model_class
+    try:
+        module = models_modules_dict[model_type]
+        model_class = getattr(importlib.import_module(module), model_type)
+        return model_class
+    except KeyError:
+        # Return a placeholder class if estimator not found
+        import warnings
+        warnings.warn(f"Model type '{model_type}' not found in sklearn modules. Returning placeholder class.")
+        
+        # Create a minimal placeholder class that can be instantiated
+        class PlaceholderModel:
+            def __init__(self, **kwargs):
+                self._model_type = model_type
+                self._params = kwargs
+            
+            def get_params(self, deep=True):
+                return self._params
+            
+            def __str__(self):
+                return f"PlaceholderModel({self._model_type})"
+            
+            def __repr__(self):
+                return f"PlaceholderModel({self._model_type})"
+        
+        return PlaceholderModel
 
 def _get_pyspark_modules():
     try: