aimodelshare 0.1.32__py3-none-any.whl → 0.1.62__py3-none-any.whl

aimodelshare/aimsonnx.py

@@ -1,29 +1,32 @@
 # data wrangling
 import pandas as pd
-import numpy as np
+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.'''
 
@@ -92,11 +139,14 @@ def _extract_onnx_metadata(onnx_model, framework):
     # initialize metadata dict
     metadata_onnx = {}
 
-    # get input shape
-    metadata_onnx["input_shape"] = graph.input[0].type.tensor_type.shape.dim[1].dim_value
-
-    # get output shape
-    metadata_onnx["output_shape"] = graph.output[0].type.tensor_type.shape.dim[1].dim_value 
+    def _get_shape(dims):
+        return [d.dim_value if d.HasField("dim_value") else None for d in dims]
+    
+    input_dims = graph.input[0].type.tensor_type.shape.dim
+    output_dims = graph.output[0].type.tensor_type.shape.dim
+    
+    metadata_onnx["input_shape"] = _get_shape(input_dims)
+    metadata_onnx["output_shape"] = _get_shape(output_dims)
     
     # get layers and activations NEW
     # match layers and nodes and initalizers in sinle object
@@ -262,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 = {}
@@ -552,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
@@ -647,20 +678,52 @@ def _keras_to_onnx(model, transfer_learning=None,
         metadata['model_weights'] = pickle.dumps(model.get_weights())
 
     # Extract architecture
-
+    if not model.built: # add shape outputs if model not built
+        try:
+            model(tf.random.uniform([1] + list(input_shape[1:])))
+        except Exception:
+            pass  # fallback, don't crash conversion
+            
     keras_layers = keras_unpack(model)
+    
+    
+    from tensorflow.python.framework import tensor_shape  # <- place this at the top of your file
+    
     layers = []
     layers_n_params = []
     layers_shapes = []
     activations = []
-
+    
     for layer in keras_layers:
+        # layer name
         layers.append(layer.__class__.__name__)
-        layers_n_params.append(layer.count_params())
-        layers_shapes.append(getattr(layer, 'output_shape', None))
+    
+        # parameter count
+        try:
+            layers_n_params.append(layer.count_params())
+        except:
+            layers_n_params.append(0)
+    
+        # output shape (sanitized for JSON)
+        shape = getattr(layer, 'output_shape', None)
+    
+        if isinstance(shape, tensor_shape.TensorShape):
+            shape = shape.as_list()
+        elif shape is not None:
+            try:
+                shape = list(shape)
+            except:
+                shape = str(shape)
+        else:
+            shape = None
+    
+        layers_shapes.append(shape)
+    
+        # activation
         if hasattr(layer, 'activation'):
             act = getattr(layer.activation, '__name__', None)
-            if act: activations.append(act)
+            if act:
+                activations.append(act)
 
     optimizer = getattr(model.optimizer, '__class__', None)
     loss = getattr(model.loss, '__class__', None)
@@ -899,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, 
@@ -954,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()
@@ -989,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
 
@@ -999,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()
@@ -1037,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
 
@@ -1089,21 +1268,36 @@ 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
     
 
 
 def _model_summary(meta_dict, from_onnx=False):
     '''Creates model summary table from model metadata dict.'''
+    import io
     
     assert(isinstance(meta_dict, dict)), \
     "Please pass valid metadata dict."
@@ -1112,9 +1306,9 @@ def _model_summary(meta_dict, from_onnx=False):
     "Please make sure model architecture data is included."
 
     if from_onnx == True:
-        model_summary = pd.read_json(meta_dict['metadata_onnx']["model_summary"])
+        model_summary = pd.read_json(io.StringIO(meta_dict['metadata_onnx']["model_summary"]))
     else:
-        model_summary = pd.read_json(meta_dict["model_summary"])
+        model_summary = pd.read_json(io.StringIO(meta_dict["model_summary"]))
        
     return model_summary
 
@@ -1517,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 = {}
 
@@ -1533,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:

aimodelshare/api.py

@@ -57,7 +57,7 @@ class create_prediction_api_class():
         self.pyspark_support = pyspark_support
         self.region = os.environ.get("AWS_REGION")
         self.bucket_name = os.environ.get("BUCKET_NAME")
-        self.python_runtime = 'python3.10'
+        self.python_runtime = os.environ.get("PYTHON_RUNTIME", "python3.12")
         #####
 
         self.model_type = self.model_type.lower()
@@ -87,23 +87,24 @@ class create_prediction_api_class():
             "video": 90,
             "custom": 90
         }
-
+        # UPDATED: New eval layer ARNs (python3.12)
         self.eval_layer_map = {
-            "us-east-1": "arn:aws:lambda:us-east-1:517169013426:layer:eval_layer_test:13",
-            "us-east-2": "arn:aws:lambda:us-east-2:517169013426:layer:eval_layer_test:7",
-            "us-west-1": "arn:aws:lambda:us-west-1:517169013426:layer:eval_layer_test:2",
-            "us-west-2": "arn:aws:lambda:us-west-2:517169013426:layer:eval_layer_test:2",
-            "eu-west-1": "arn:aws:lambda:eu-west-1:517169013426:layer:eval_layer_test:2",
-            "eu-west-2": "arn:aws:lambda:eu-west-2:517169013426:layer:eval_layer_test:2",
-            "eu-west-3": "arn:aws:lambda:eu-west-3:517169013426:layer:eval_layer_test:2"
+            "us-east-1": "arn:aws:lambda:us-east-1:585666012274:layer:eval-layer-python3-12:4",
+            "us-east-2": "arn:aws:lambda:us-east-2:517169013426:layer:eval_layer_test:5",
+            "us-west-1": "arn:aws:lambda:us-west-1:517169013426:layer:eval_layer_test:1",
+            "us-west-2": "arn:aws:lambda:us-west-2:517169013426:layer:eval_layer_test:1",
+            "eu-west-1": "arn:aws:lambda:eu-west-1:585666012274:layer:eval-layer-python3-12:1",
+            "eu-west-2": "arn:aws:lambda:eu-west-2:517169013426:layer:eval_layer_test:1",
+            "eu-west-3": "arn:aws:lambda:eu-west-3:517169013426:layer:eval_layer_test:1"
         }
 
+        # UPDATED: New auth layer ARNs (python3.12)
         self.auth_layer_map = {
-            "us-east-1": "arn:aws:lambda:us-east-1:517169013426:layer:aimsauth_layer:2",
+            "us-east-1": "arn:aws:lambda:us-east-1:585666012274:layer:aimsauth-layer-python3-12:6",
             "us-east-2": "arn:aws:lambda:us-east-2:517169013426:layer:aimsauth_layer:9",
             "us-west-1": "arn:aws:lambda:us-west-1:517169013426:layer:aimsauth_layer:1",
             "us-west-2": "arn:aws:lambda:us-west-2:517169013426:layer:aimsauth_layer:1",
-            "eu-west-1": "arn:aws:lambda:eu-west-1:517169013426:layer:aimsauth_layer:1",
+            "eu-west-1": "arn:aws:lambda:eu-west-1:585666012274:layer:aimsauth-layer-python3-12:6",
             "eu-west-2": "arn:aws:lambda:eu-west-2:517169013426:layer:aimsauth_layer:1",
             "eu-west-3": "arn:aws:lambda:eu-west-3:517169013426:layer:aimsauth_layer:1"
         }
@@ -293,7 +294,7 @@ class create_prediction_api_class():
             )
 
         eval_code_source = {'S3Bucket': self.bucket_name, 'S3Key':  self.unique_model_id + "/" + "archiveeval.zip"}
-        eval_layers = [self.eval_layer]
+        eval_layers = [self.eval_layer, self.auth_layer]
         create_lambda_function(lambdaevalfxnname, self.python_runtime, role_arn, handler, eval_code_source, 90, 2048, eval_layers)
 
         auth_code_source = {'S3Bucket': self.bucket_name, 'S3Key':  self.unique_model_id + "/" + "archiveauth.zip"}