Update hbw/tasks/inference.py

Co-authored-by: Mathis Frahm <49306645+mafrahm@users.noreply.github.com>

Author: Lara813

hbw/ml/base.py

@@ -335,23 +335,14 @@ def open_model(self, target: law.LocalDirectoryTarget) -> dict[str, Any]:
         models["parameters"] = yaml.load(f_in, Loader=yaml.Loader)
 
         # custom loss needed due to output layer changes for negative weights
-        from hbw.ml.tf_util import cumulated_crossentropy, categorical_crossentropy
+        from hbw.ml.tf_util import cumulated_crossentropy
 
-        # Check for negative weight handling and assign loss function accordingly.
-        if self.negative_weights == "ignore":
-            models["model"] = tf.keras.models.load_model(
-                target["mlmodel"].path, custom_objects={categorical_crossentropy.__name__: categorical_crossentropy},
-            )
-            models["best_model"] = tf.keras.models.load_model(
-                target["checkpoint"].path, custom_objects={categorical_crossentropy.__name__: categorical_crossentropy},
-            )
-        else:
-            models["model"] = tf.keras.models.load_model(
-                target["mlmodel"].path, custom_objects={cumulated_crossentropy.__name__: cumulated_crossentropy},
-            )
-            models["best_model"] = tf.keras.models.load_model(
-                target["checkpoint"].path, custom_objects={cumulated_crossentropy.__name__: cumulated_crossentropy},
-            )
+        models["model"] = tf.keras.models.load_model(
+            target["mlmodel"].path, custom_objects={cumulated_crossentropy.__name__: cumulated_crossentropy},
+        )
+        models["best_model"] = tf.keras.models.load_model(
+            target["checkpoint"].path, custom_objects={cumulated_crossentropy.__name__: cumulated_crossentropy},
+        )
 
         return models
 
@@ -553,7 +544,7 @@ def prepare_ml_model(
 
         from keras.models import Sequential
         from keras.layers import Dense, BatchNormalization
-        from hbw.ml.tf_util import cumulated_crossentropy, categorical_crossentropy
+        from hbw.ml.tf_util import cumulated_crossentropy
 
         n_inputs = len(set(self.input_features))
         n_outputs = len(self.processes)
@@ -576,7 +567,7 @@ def prepare_ml_model(
         optimizer = keras.optimizers.Adam(learning_rate=0.00050)
         if self.negative_weights == "ignore":
             model.compile(
-                loss=categorical_crossentropy,
+                loss="categorical_crossentropy",
                 optimizer=optimizer,
                 weighted_metrics=["categorical_accuracy"],
             )

hbw/ml/mixins.py

@@ -55,7 +55,7 @@ def prepare_ml_model(
         import tensorflow.keras as keras
         from keras.models import Sequential
         from keras.layers import Dense, BatchNormalization
-        from hbw.ml.tf_util import cumulated_crossentropy  # , categorical_crossentropy
+        from hbw.ml.tf_util import cumulated_crossentropy
 
         n_inputs = len(set(self.input_features))
         n_outputs = len(self.processes)

hbw/tasks/inference.py

@@ -32,7 +32,6 @@
 
 # Function copied from Mathis Hist hook commit
 # TODO: define once at central place (hist_util.py)
-# TODO: define once at central place (hist_util.py)
 def apply_rebinning_edges(h: hist.Histogram, axis_name: str, edges: list):
     """
     Generalized rebinning of a single axis from a hist.Histogram, using predefined edges.