star

venv/lib/python3.9/site-packages/mlflow/keras/utils.py

@@ -0,0 +1,38 @@
+import numpy as np
+
+from mlflow.models import ModelSignature
+from mlflow.types import Schema, TensorSpec
+
+
+def get_model_signature(model):
+    def replace_none_in_shape(shape):
+        return [-1 if dim_size is None else dim_size for dim_size in shape]
+
+    input_shape = model.input_shape
+    input_dtype = model.input_dtype
+    output_shape = model.output_shape
+    output_dtype = model.compute_dtype
+
+    if isinstance(input_shape, list):
+        input_schema = Schema(
+            [
+                TensorSpec(np.dtype(input_dtype), replace_none_in_shape(shape))
+                for shape in input_shape
+            ]
+        )
+    else:
+        input_schema = Schema(
+            [TensorSpec(np.dtype(input_dtype), replace_none_in_shape(input_shape))]
+        )
+    if isinstance(output_shape, list):
+        output_schema = Schema(
+            [
+                TensorSpec(np.dtype(output_dtype), replace_none_in_shape(shape))
+                for shape in output_shape
+            ]
+        )
+    else:
+        output_schema = Schema(
+            [TensorSpec(np.dtype(output_dtype), replace_none_in_shape(output_shape))]
+        )
+    return ModelSignature(inputs=input_schema, outputs=output_schema)