updates to comments

IDAES · Oct 18, 2024 · d4fb28d · d4fb28d
1 parent 05ff135
commit d4fb28d
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/idaes/core/surrogate/onnx_surrogate.py b/idaes/core/surrogate/onnx_surrogate.py
@@ -69,9 +69,9 @@ def __init__(
         Args:
            onnx_model: Onnx model file to be loaded.
            input_labels: list of str
-              The ordered list of labels corresponding to the inputs in the keras model
+              The ordered list of labels corresponding to the inputs in the onnx model
            output_labels: list of str
-              The ordered list of labels corresponding to the outputs in the keras model
+              The ordered list of labels corresponding to the outputs in the onnx model
            input_bounds: None of dict of tuples
               Keys correspond to each of the input labels and values are the tuples of
               bounds (lb, ub)
@@ -105,7 +105,7 @@ def populate_block(self, block, additional_options=None):
               The block to be populated with variables and/or constraints.
            additional_options: dict or None
               If not None, then should be a dict with the following keys;
-              'formulation': KerasSurrogate.Formulation
+              'formulation': ONNXSurrogate.Formulation
               The formulation to use with OMLT. Possible values are FULL_SPACE,
               REDUCED_SPACE, RELU_BIGM, or RELU_COMPLEMENTARITY (default is FULL_SPACE)
         """
@@ -146,7 +146,7 @@ def populate_block(self, block, additional_options=None):
         else:
             raise ValueError(
                 'An unrecognized formulation "{}" was passed to '
-                "KerasSurrogate.populate_block. Please pass a valid "
+                "ONNXSurrogate.populate_block. Please pass a valid "
                 "formulation.".format(formulation)
             )
         self.populate_block_with_net(block, formulation_object)

diff --git a/idaes/core/surrogate/tests/test_onnx_surrogate.py b/idaes/core/surrogate/tests/test_onnx_surrogate.py
@@ -43,8 +43,6 @@
 def load_onnx_model_data(
     name="net_st_net_5000_STM_100_s_2000000_60_5_tanh_1e-06_4096_tr_15481_Calcite_ST",
 ):
-    #    keras_folder_name = os.path.join(this_file_dir(), 'data', 'keras_models', name)
-    #    keras_model = keras.models.load_model(keras_folder_name)
     onnx_folder_name = os.path.join(this_file_dir(), "data", "onnx_models")
     onnx_model = onnx.load(os.path.join(onnx_folder_name, "{}.onnx".format(name)))
     with open(os.path.join(onnx_folder_name, "{}_idaes_info.json".format(name))) as fd: