Update custom_modeling.py

custom_modeling.py

@@ -1,13 +1,9 @@
 """
-custom_modeling.py
-------------------
-A single model-agnostic toxicity wrapper for any causal-LM on the Hugging
-Face Hub.
-
-Keep this in the repo **alongside**:
-  • toxic.keras            -- your TF/Keras classifier file
-
-Make sure config.json contains:
+custom_modeling.py  – model-agnostic toxicity wrapper
+----------------------------------------------------
+Place in repo root together with:
+  • toxic.keras
+Add to config.json:
   "auto_map": { "AutoModelForCausalLM": "custom_modeling.SafeGenerationModel" }
 """
 
@@ -21,12 +17,16 @@ from huggingface_hub import hf_hub_download
 
 
 # ------------------------------------------------------------------ #
-# 1)  MIXIN  – all toxicity filtering lives here                     #
+# 1)  MIXIN – toxicity filtering logic                               #
 # ------------------------------------------------------------------ #
 class _SafeGenerationMixin:
     _toxicity_model = None
-    _tox_threshold = 0.6
-    _safe_message = "Response is toxic, please be kind to yourself and others."
+    _tox_threshold   = 0.6
+
+    # Separate messages
+    _safe_in_msg  = "Sorry, I can’t help with that request."
+    _safe_out_msg = "I’m sorry, but I can’t continue with that."
+
     _tokenizer = None
 
     # ---- helpers ----------------------------------------------------
@@ -55,18 +55,17 @@ class _SafeGenerationMixin:
     def _is_toxic(self, text: str) -> bool:
         if not text.strip():
             return False
-        inputs = tf.constant([text], dtype=tf.string)   # <= proper tensor
+        inputs = tf.constant([text], dtype=tf.string)
         prob   = float(self._tox_model.predict(inputs)[0, 0])
         return prob >= self._tox_threshold
 
-    def _safe_ids(self, length: int | None = None) -> torch.LongTensor:
-        """Return token IDs for the safe message, padded / truncated to *length*."""
+    def _safe_ids(self, message: str, length: int | None = None):
+        """Encode *message* and pad/truncate to *length* tokens (if given)."""
         self._ensure_tokenizer()
         if self._tokenizer is None:
             raise RuntimeError("Tokenizer unavailable for safe-message encoding.")
 
-        ids = self._tokenizer(self._safe_message, return_tensors="pt")["input_ids"][0]
-
+        ids = self._tokenizer(message, return_tensors="pt")["input_ids"][0]
         if length is not None:
             pad_id = (
                 self.config.eos_token_id
@@ -79,14 +78,13 @@ class _SafeGenerationMixin:
                 )
             else:
                 ids = ids[:length]
-
         return ids.to(self._device())
 
     # ---- main override ---------------------------------------------
     def generate(self, *args, **kwargs):
         self._ensure_tokenizer()
 
-        # 1)  Prompt toxicity check
+        # 1) prompt toxicity
         prompt_txt = None
         if self._tokenizer is not None:
             if "input_ids" in kwargs:
@@ -99,81 +97,63 @@ class _SafeGenerationMixin:
                 )
 
         if prompt_txt and self._is_toxic(prompt_txt):
-            return self._safe_ids().unsqueeze(0)
+            return self._safe_ids(self._safe_in_msg).unsqueeze(0)
 
-        # 2)  Normal generation
+        # 2) normal generation
         outputs = super().generate(*args, **kwargs)
 
-        # 3)  Output toxicity check
+        # 3) output toxicity
         if self._tokenizer is None:
             return outputs
 
-        cleaned_seqs = []
+        new_seqs = []
         for seq in outputs.detach().cpu():
             txt = self._tokenizer.decode(seq.tolist(), skip_special_tokens=True)
-            cleaned_seqs.append(
-                self._safe_ids(length=seq.size(0)) if self._is_toxic(txt) else seq
-            )
-
-        return torch.stack(cleaned_seqs, dim=0).to(self._device())
+            if self._is_toxic(txt):
+                new_seqs.append(self._safe_ids(self._safe_out_msg, length=seq.size(0)))
+            else:
+                new_seqs.append(seq)
+        return torch.stack(new_seqs, dim=0).to(self._device())
 
 
 # ------------------------------------------------------------------ #
-# 2)  Utilities: resolve & cache the real base class                 #
+# 2)  utilities: resolve base class & cache subclass                 #
 # ------------------------------------------------------------------ #
 @lru_cache(None)
-def _get_base_cls(arch_name: str):
-    """Return the actual Transformers class for the given architecture string."""
-    if hasattr(transformers, arch_name):
-        return getattr(transformers, arch_name)
-
-    # Fallback: import based on naming convention
-    stem = arch_name.replace("ForCausalLM", "").lower()
+def _get_base_cls(arch: str):
+    if hasattr(transformers, arch):
+        return getattr(transformers, arch)
+    stem = arch.replace("ForCausalLM", "").lower()
     module = importlib.import_module(f"transformers.models.{stem}.modeling_{stem}")
-    return getattr(module, arch_name)
+    return getattr(module, arch)
 
 
 @lru_cache(None)
 def _make_safe_subclass(base_cls):
-    """Create (and cache) SafeGeneration_<Base> = (Mixin, Base)."""
-    return type(f"SafeGeneration_{base_cls.__name__}", (_SafeGenerationMixin, base_cls), {})
+    return type(
+        f"SafeGeneration_{base_cls.__name__}",
+        (_SafeGenerationMixin, base_cls),
+        {},
+    )
 
 
 # ------------------------------------------------------------------ #
-# 3)  Dispatcher class referenced in auto_map                        #
+# 3)  Dispatcher class – referenced by auto_map                      #
 # ------------------------------------------------------------------ #
 class SafeGenerationModel:
-    """
-    Lightweight dispatcher so that `AutoModelForCausalLM` can load the
-    wrapped model transparently.
-    """
-
     @classmethod
-    def from_pretrained(cls, model_name_or_path, *model_args, **kwargs):
-        # Ensure custom code execution is allowed
+    def from_pretrained(cls, repo_id, *model_args, **kwargs):
         kwargs.setdefault("trust_remote_code", True)
-
-        # Remove literal "auto" dtype to avoid downstream dtype bugs
         if kwargs.get("torch_dtype") == "auto":
             kwargs.pop("torch_dtype")
 
-        # Load config first to discover architecture
-        config = transformers.AutoConfig.from_pretrained(model_name_or_path, **kwargs)
+        config = transformers.AutoConfig.from_pretrained(repo_id, **kwargs)
         if not getattr(config, "architectures", None):
             raise ValueError("`config.architectures` missing in config.json.")
         arch_str = config.architectures[0]
 
-        # Build / fetch dynamic subclass
-        BaseCLS = _get_base_cls(arch_str)
-        SafeCLS = _make_safe_subclass(BaseCLS)
-
-        # Avoid duplicate 'config' key
-        kwargs.pop("config", None)
+        Base = _get_base_cls(arch_str)
+        Safe = _make_safe_subclass(Base)
 
-        # Delegate full loading to the safe subclass
-        return SafeCLS.from_pretrained(
-            model_name_or_path,
-            *model_args,
-            config=config,
-            **kwargs,
-        )
+        kwargs.pop("config", None)    # avoid duplicate
+        return Safe.from_pretrained(repo_id, *model_args, config=config, **kwargs)