Update custom_modeling.py

custom_modeling.py

@@ -1,39 +1,3 @@
-import torch
-import tensorflow as tf
-from transformers import LlamaForCausalLM
-from transformers.utils import cached_file
-import os
-import logging
-
-# Set up logging
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-class ToxicityChecker:
-    def __init__(self, model_path):
-        try:
-            # Check if file exists
-            if not os.path.exists(model_path):
-                raise FileNotFoundError(f"Toxicity model not found at {model_path}")
-                
-            logger.info(f"Loading toxicity model from: {model_path}")
-            self.model = tf.keras.models.load_model(model_path)
-            logger.info("Toxicity model loaded successfully")
-        except Exception as e:
-            logger.error(f"Failed to load toxicity model: {str(e)}")
-            raise
-    
-    def is_toxic(self, text, threshold=0.6):
-        try:
-            # Convert to TensorFlow constant
-            text_tensor = tf.constant([text])
-            prob = self.model.predict(text_tensor, verbose=0)[0][0]
-            logger.debug(f"Toxicity check: '{text[:30]}...' → prob: {prob:.3f}")
-            return prob > threshold
-        except Exception as e:
-            logger.error(f"Toxicity check failed: {str(e)}")
-            return False
-
 import os
 import torch
 import tensorflow as tf
@@ -43,50 +7,50 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-class ToxicityChecker:
-    def __init__(self, model_path):
-        if not os.path.exists(model_path):
-            raise FileNotFoundError(f"Toxicity model not found at {model_path}")
-        logger.info(f"Loading toxicity model from: {model_path}")
-        self.model = tf.keras.models.load_model(model_path)
+class SafeGenerationModel(LlamaForCausalLM):
+    def __init__(self, config):
+        super().__init__(config)
+        
+        # Load toxicity model
+        toxic_path = cached_file(config._name_or_path, "toxic.keras")
+        if not os.path.exists(toxic_path):
+            raise FileNotFoundError(f"Toxicity model not found at {toxic_path}")
+        
+        self.toxicity_model = tf.keras.models.load_model(toxic_path)
+        self.tokenizer = None
         logger.info("Toxicity model loaded successfully")
-    
+
     def is_toxic(self, text, threshold=0.6):
         try:
-            text_tensor = tf.constant([text])
-            prob = self.model.predict(text_tensor, verbose=0)[0][0]
-            logger.debug(f"Toxicity: '{text[:30]}...' → {prob:.3f}")
+            prob = self.toxicity_model.predict([text], verbose=0)[0][0]
             return prob > threshold
-        except Exception:
+        except Exception as e:
+            logger.error(f"Toxicity check failed: {str(e)}")
             return False
 
-class SafeGenerationModel(LlamaForCausalLM):
-    def __init__(self, config):
-        super().__init__(config)
-        toxic_path = cached_file(config.name_or_path, "toxic.keras")
-        self.toxicity_checker = ToxicityChecker(toxic_path)
-        self.tokenizer = None
-    
     def generate(self, *args, **kwargs):
         inputs = kwargs.get("input_ids")
-        if inputs is not None and self.tokenizer:
+        
+        # Check input toxicity
+        if self.tokenizer and inputs is not None:
             input_text = self.tokenizer.decode(inputs[0], skip_special_tokens=True)
-            if self.toxicity_checker.is_toxic(input_text):
-                return self._safe_response_tensor()
+            if self.is_toxic(input_text):
+                return self._safe_response()
         
+        # Generate response
         outputs = super().generate(*args, **kwargs)
         
+        # Check output toxicity
         if self.tokenizer:
             output_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            if self.toxicity_checker.is_toxic(output_text):
-                return self._safe_response_tensor()
+            if self.is_toxic(output_text):
+                return self._safe_response()
         
         return outputs
-    
-    def _safe_response_tensor(self):
-        safe_response = "I'm unable to respond to that request. HAHAHA"
-        return self.tokenizer.encode(safe_response, return_tensors="pt").to(self.device)
-    
+
+    def _safe_response(self):
+        safe_text = "I'm unable to respond to that request. HAHAHA"
+        return self.tokenizer.encode(safe_text, return_tensors="pt").to(self.device)
+
     def set_tokenizer(self, tokenizer):
-        self.tokenizer = tokenizer
-        logger.info("Tokenizer injected into model")
+        self.tokenizer = tokenizer