Upload app.py

app.py

@@ -71,17 +71,21 @@ def load_model():
                 bce = -(y_true * tf.math.log(y_pred) + (1.0 - y_true) * tf.math.log(1.0 - y_pred))
                 return tf.reduce_mean(bce)
             
-            # Provide custom_objects to handle the custom loss function
+            # Provide custom_objects to handle custom loss function and TFRobertaModel
+            # TFRobertaModel is needed because the model architecture uses it
             custom_objects = {
-                'weighted_binary_crossentropy': weighted_binary_crossentropy
+                'weighted_binary_crossentropy': weighted_binary_crossentropy,
+                'TFRobertaModel': TFRobertaModel
             }
             
-            # Load the model using TensorFlow's Keras with custom_objects
+            # Load the model using tf_keras directly (not tf.keras) with custom_objects
+            # Use safe_mode=False to allow loading custom objects
             # This is needed because the model was saved with tf_keras and a custom loss
-            model = tf.keras.models.load_model(
+            model = tf_keras.models.load_model(
                 model_path, 
                 compile=False,
-                custom_objects=custom_objects
+                custom_objects=custom_objects,
+                safe_mode=False
             )
             print("Model loaded successfully!")
         return model