Update custom_modeling.py

custom_modeling.py

@@ -45,7 +45,9 @@ class _SafeGenerationMixin:
                 repo_id=self.config.name_or_path,
                 filename="toxic.keras",
             )
-            self._toxicity_model = tf.keras.models.load_model(path, compile=False)
+            # ✅ Force CPU loading for better performance
+            with tf.device('/CPU:0'):
+                self._toxicity_model = tf.keras.models.load_model(path, compile=False)
         return self._toxicity_model
 
     @property
@@ -55,7 +57,9 @@ class _SafeGenerationMixin:
                 repo_id=self.config.name_or_path,
                 filename="PI.keras",
             )
-            self._pi_model = tf.keras.models.load_model(path, compile=False)
+            # ✅ Force CPU loading for better performance
+            with tf.device('/CPU:0'):
+                self._pi_model = tf.keras.models.load_model(path, compile=False)
         return self._pi_model
 
     def _ensure_tokenizer(self):
@@ -70,15 +74,19 @@ class _SafeGenerationMixin:
     def _is_toxic(self, text: str) -> bool:
         if not text.strip():
             return False
-        inputs = tf.constant([text], dtype=tf.string)
-        prob = float(self._tox_model.predict(inputs)[0, 0])
+        # ✅ Ensure CPU execution
+        with tf.device('/CPU:0'):
+            inputs = tf.constant([text], dtype=tf.string)
+            prob = float(self._tox_model.predict(inputs, verbose=0)[0, 0])
         return prob >= self._tox_threshold
 
     def _has_prompt_injection(self, text: str) -> bool:
         if not text.strip():
             return False
-        inputs = tf.constant([text], dtype=tf.string)
-        prob = float(self._prompt_injection_model.predict(inputs)[0, 0])
+        # ✅ Ensure CPU execution
+        with tf.device('/CPU:0'):
+            inputs = tf.constant([text], dtype=tf.string)
+            prob = float(self._prompt_injection_model.predict(inputs, verbose=0)[0, 0])
         return prob >= self._pi_threshold
 
     def _safe_ids(self, message: str, length: int | None = None):