aimodelshare 0.1.29__py3-none-any.whl → 0.1.64__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 = {}
@@ -549,197 +579,178 @@ def _pyspark_to_onnx(model, initial_types, spark_session,
     return onx
 
 def _keras_to_onnx(model, transfer_learning=None,
-                  deep_learning=None, task_type=None, epochs=None):
-    '''Extracts metadata from keras model object.'''
+                   deep_learning=None, task_type=None, epochs=None):
+    '''Converts a Keras model to ONNX and extracts metadata.'''
 
-    # check whether this is a fitted keras model
-    # isinstance...
+    # Check and load tf2onnx and TensorFlow lazily (only when needed)
+    tf2onnx, tf = _check_tf2onnx_available()
+    
+    import numpy as np
+    import onnx
+    import pickle
+    import psutil
+    import warnings
+    from pympler import asizeof
+    import logging
+    import os
+    import sys
+    from contextlib import contextmanager
+
+    # -- Helper to suppress tf2onnx stderr (NumPy warnings etc.)
+    @contextmanager
+    def suppress_stderr():
+        with open(os.devnull, "w") as devnull:
+            old_stderr = sys.stderr
+            sys.stderr = devnull
+            try:
+                yield
+            finally:
+                sys.stderr = old_stderr
+
+    # Reduce logging output
+    tf2onnx_logger = logging.getLogger("tf2onnx")
+    tf2onnx_logger.setLevel(logging.CRITICAL)
+
+    # Unwrap scikeras, sklearn pipelines etc.
+    from sklearn.model_selection import GridSearchCV, RandomizedSearchCV
+    from sklearn.pipeline import Pipeline
+    from scikeras.wrappers import KerasClassifier, KerasRegressor
 
-    # handle keras models in sklearn wrapper
     if isinstance(model, (GridSearchCV, RandomizedSearchCV)):
         model = model.best_estimator_
-
-    if isinstance(model, sklearn.pipeline.Pipeline):
+    if isinstance(model, Pipeline):
         model = model.steps[-1][1]
-
-    sklearn_wrappers = (KerasClassifier,KerasRegressor)
-
-    if isinstance(model, sklearn_wrappers):
+    if isinstance(model, (KerasClassifier, KerasRegressor)):
         model = model.model
-    
-    # convert to onnx
-    #onx = convert_keras(model)
-    # generate tempfile for onnx object 
-    temp_dir = tempfile.mkdtemp()
-
-
 
-    
-    tf.get_logger().setLevel('ERROR') # probably not good practice
-    output_path = os.path.join(temp_dir, 'temp.onnx')
-    
-    
-    model.save(temp_dir)
-
-    # # Convert the model
-    try:
-            modelstringtest="python -m tf2onnx.convert --saved-model  "+temp_dir+" --output "+output_path+" --opset 13"
-            resultonnx=os.system(modelstringtest)
-            resultonnx2=1
-            if resultonnx==0:
-              pass
-            else:
-              raise Exception('Model conversion to onnx unsuccessful.  Please try different model or submit predictions to leaderboard without submitting preprocessor or model files.')
-    except:
-            converter = tf.lite.TFLiteConverter.from_saved_model(temp_dir) # path to the SavedModel directory
-            converter.target_spec.supported_ops = [
-                tf.lite.OpsSet.TFLITE_BUILTINS, # enable TensorFlow Lite ops.
-                tf.lite.OpsSet.SELECT_TF_OPS # enable TensorFlow ops.
-              ]
-            tflite_model = converter.convert()
-
-            # Save the model.
-            with open(os.path.join(temp_dir,'tempmodel.tflite'), 'wb') as f:
-              f.write(tflite_model)
-
-            modelstringtest="python -m tf2onnx.convert --tflite "+os.path.join(temp_dir,'tempmodel.tflite')+" --output "+output_path+" --opset 13"
-            resultonnx2=os.system(modelstringtest)
-            pass
-
-    if any([resultonnx==0, resultonnx2==0]):
-      pass
-    else:
-      return print("Model conversion to onnx unsuccessful.  Please try different model or submit\npredictions to leaderboard without submitting preprocessor or model files.")
-
-    onx = onnx.load(output_path)
-
-
-    # generate metadata dict 
-    metadata = {}
-
-    # placeholders, need to be generated elsewhere
-    metadata['model_id'] = None
-    metadata['data_id'] = None
-    metadata['preprocessor_id'] = None
-
-    # infer ml framework from function call
-    metadata['ml_framework'] = 'keras'
-
-    # get model type from model object
-    metadata['model_type'] =  str(model.__class__.__name__)
-
-    # get transfer learning bool from user input
-    metadata['transfer_learning'] = transfer_learning
-
-    # get deep learning bool from user input
-    metadata['deep_learning'] = deep_learning
-
-    # get task type from user input
-    metadata['task_type'] = task_type
-
-    # placeholders, need to be inferred from data 
-    metadata['target_distribution'] = None
-    metadata['input_type'] = None
-    metadata['input_shape'] = None
-    metadata['input_dtypes'] = None       
-    metadata['input_distribution'] = None
+    # Input signature
+    input_shape = model.input_shape
+    if isinstance(input_shape, list):
+        input_shape = input_shape[0]
+    input_signature = [tf.TensorSpec(input_shape, tf.float32, name="input")]
+
+    # Wrap model in tf.function
+    @tf.function(input_signature=input_signature)
+    def model_fn(x):
+        return model(x)
+
+    concrete_func = model_fn
+
+    # Convert to ONNX
+    with suppress_stderr():
+        onx_model, _ = tf2onnx.convert.from_function(
+            concrete_func,
+            input_signature=input_signature,
+            opset=13,
+            output_path=None
+        )
 
-    # get model config dict from keras model object
-    metadata['model_config'] = str(model.get_config())
+    # Extract metadata
+    metadata = {
+        'model_id': None,
+        'data_id': None,
+        'preprocessor_id': None,
+        'ml_framework': 'keras',
+        'model_type': model.__class__.__name__,
+        'transfer_learning': transfer_learning,
+        'deep_learning': deep_learning,
+        'task_type': task_type,
+        'target_distribution': None,
+        'input_type': None,
+        'input_shape': input_shape,
+        'input_dtypes': None,
+        'input_distribution': None,
+        'model_config': str(model.get_config()),
+        'model_state': None,
+        'eval_metrics': None,
+        'model_graph': "",
+        'metadata_onnx': None,
+        'epochs': epochs
+    }
 
-    # get model weights from keras object 
     model_size = asizeof.asizeof(model.get_weights())
     mem = psutil.virtual_memory()
 
-    if model_size > mem.available: 
-
-        warnings.warn(f"Model size ({model_size/1e6} MB) exceeds available memory ({mem.available/1e6} MB). Skipping extraction of model weights.")
-
+    if model_size > mem.available:
+        warnings.warn(f"Model size ({model_size/1e6} MB) exceeds available memory.")
         metadata['model_weights'] = None
-
-    else: 
-        
+    else:
         metadata['model_weights'] = pickle.dumps(model.get_weights())
 
-    # get model state from pytorch model object
-    metadata['model_state'] = None
+    # 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
     
-    # get list of current layer types 
-    layer_list, activation_list = _get_layer_names()
-
-    # extract model architecture metadata 
     layers = []
     layers_n_params = []
     layers_shapes = []
     activations = []
-
-
-    keras_layers = keras_unpack(model)
-
-    for i in keras_layers: 
-        
-        # get layer names 
-        if i.__class__.__name__ in layer_list:
-            layers.append(i.__class__.__name__)
-            layers_n_params.append(i.count_params())
-            layers_shapes.append(i.output_shape)
-        
-        # get activation names
-        if i.__class__.__name__ in activation_list: 
-            activations.append(i.__class__.__name__.lower())
-        if hasattr(i, 'activation') and i.activation.__name__ in activation_list:
-            activations.append(i.activation.__name__)
-
-    if hasattr(model, 'loss'):
-        loss = model.loss.__class__.__name__
-    else:
-        loss = None
-        
-    if hasattr(model, 'optimizer'):
-        optimizer = model.optimizer.__class__.__name__
-    else:
-        optimizer = None 
-
-    model_summary_pd = model_summary_keras(model)
     
-    # insert data into model architecture dict 
-    model_architecture = {'layers_number': len(layers),
-                          'layers_sequence': layers,
-                          'layers_summary': {i:layers.count(i) for i in set(layers)},
-                          'layers_n_params': layers_n_params,
-                          'layers_shapes': layers_shapes,
-                          'activations_sequence': activations,
-                          'activations_summary': {i:activations.count(i) for i in set(activations)},
-                          'loss':loss,
-                          'optimizer': optimizer
-                         }
+    for layer in keras_layers:
+        # layer name
+        layers.append(layer.__class__.__name__)
+    
+        # 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)
+
+    optimizer = getattr(model.optimizer, '__class__', None)
+    loss = getattr(model.loss, '__class__', None)
+
+    model_architecture = {
+        'layers_number': len(layers),
+        'layers_sequence': layers,
+        'layers_summary': {i: layers.count(i) for i in set(layers)},
+        'layers_n_params': layers_n_params,
+        'layers_shapes': layers_shapes,
+        'activations_sequence': activations,
+        'activations_summary': {i: activations.count(i) for i in set(activations)},
+        'loss': loss.__name__ if loss else None,
+        'optimizer': optimizer.__name__ if optimizer else None
+    }
 
     metadata['model_architecture'] = str(model_architecture)
-
-
-    metadata['model_summary'] = model_summary_pd.to_json()
-
+    metadata['model_summary'] = model_summary_keras(model).to_json()
     metadata['memory_size'] = model_size
 
-    metadata['epochs'] = epochs
-
-    # model graph 
-    #G = model_graph_keras(model)
-    #metadata['model_graph'] = G.create_dot().decode('utf-8')
-    metadata['model_graph'] = ""
-    # placeholder, needs evaluation engine
-    metadata['eval_metrics'] = None
-
-    # add metadata from onnx object
-    # metadata['metadata_onnx'] = str(_extract_onnx_metadata(onx, framework='keras'))
-    metadata['metadata_onnx'] = None
-    # add metadata dict to onnx object
-    
-    meta = onx.metadata_props.add()
+    # Embed metadata in ONNX
+    meta = onx_model.metadata_props.add()
     meta.key = 'model_metadata'
     meta.value = str(metadata)
 
-    return onx
+    return onx_model
+
 
 
 def _pytorch_to_onnx(model, model_input, transfer_learning=None, 
@@ -951,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, 
@@ -1006,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()
@@ -1041,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
 
@@ -1051,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()
@@ -1089,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
 
@@ -1141,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."
@@ -1164,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
 
@@ -1569,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 = {}
 
@@ -1585,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:
@@ -1844,22 +2009,12 @@ def torch_unpack(model):
 
     
 def keras_unpack(model):
-    
     layers = []
-    
     for module in model.layers:
-                
         if isinstance(module, (tf.keras.Model, tf.keras.Sequential)):
-            
-            layers_out = keras_unpack(module)
-            
-            layers = layers + layers_out
-            
-            
+            layers += keras_unpack(module)
         else:
-            
             layers.append(module)
-            
     return layers