Update custom_modeling.py

custom_modeling.py

@@ -1,15 +1,8 @@
 """
-custom_modeling.py
-------------------
-Model-agnostic toxicity wrapper for any Hugging Face causal-LM.
-
-Add (or keep) in your config.json:
-  "auto_map": {
-    "AutoModelForCausalLM": "custom_modeling.SafeGenerationModel"
-  }
-
-Files that must live in the repo alongside this script:
-  • toxic.keras  – Keras classifier (sigmoid output: toxic prob)
+custom_modeling.py  – model-agnostic toxicity wrapper
+----------------------------------------------------
+Keep in config.json:
+  "auto_map": {"AutoModelForCausalLM": "custom_modeling.SafeGenerationModel"}
 """
 
 import importlib
@@ -22,19 +15,15 @@ from huggingface_hub import hf_hub_download
 
 
 # ---------------------------------------------------------------------
-# 1)  MIXIN – all toxicity logic lives here
+# 1)  MIXIN – toxicity logic
 # ---------------------------------------------------------------------
 class _SafeGenerationMixin:
-    """Mixin that overrides .generate() to filter toxic prompts / outputs."""
-
     _toxicity_model = None
     _tox_threshold = 0.6
-    _safe_message = (
-        "Response is toxic, please be kind to yourself and others."
-    )
+    _safe_message = "Response is toxic, please be kind to yourself and others."
     _tokenizer = None
 
-    # ----- helper: load classifier on first use -----------------------
+    # ---------- classifier ----------
     @property
     def _tox_model(self):
         if self._toxicity_model is None:
@@ -42,12 +31,10 @@ class _SafeGenerationMixin:
                 repo_id=self.config.name_or_path,
                 filename="toxic.keras",
             )
-            self._toxicity_model = tf.keras.models.load_model(
-                path, compile=False
-            )
+            self._toxicity_model = tf.keras.models.load_model(path, compile=False)
         return self._toxicity_model
 
-    # ----- helper: load tokenizer (once) ------------------------------
+    # ---------- tokenizer ----------
     def _ensure_tokenizer(self):
         if self._tokenizer is None:
             try:
@@ -57,14 +44,13 @@ class _SafeGenerationMixin:
             except Exception:
                 pass
 
-    # ----- helper: tox check -----------------------------------------
+    # ---------- helpers ----------
     def _is_toxic(self, text: str) -> bool:
         if not text.strip():
             return False
         prob = float(self._tox_model.predict([text])[0, 0])
         return prob >= self._tox_threshold
 
-    # ----- helper: safe token ids ------------------------------------
     def _safe_ids(self, length: int | None = None) -> torch.LongTensor:
         self._ensure_tokenizer()
         if self._tokenizer is None:
@@ -77,19 +63,16 @@ class _SafeGenerationMixin:
                 else (self.config.pad_token_id or 0)
             )
             if ids.size(0) < length:
-                ids = torch.cat(
-                    [ids, ids.new_full((length - ids.size(0),), pad_id)],
-                    dim=0,
-                )
+                ids = torch.cat([ids, ids.new_full((length - ids.size(0),), pad_id)], 0)
             else:
                 ids = ids[:length]
         return ids.to(self.device)
 
-    # ----- override generate() ---------------------------------------
+    # ---------- override generate ----------
     def generate(self, *args, **kwargs):
         self._ensure_tokenizer()
 
-        # 1) prompt toxicity
+        # 1) prompt check
         prompt_txt = None
         if self._tokenizer is not None:
             if "input_ids" in kwargs:
@@ -100,81 +83,66 @@ class _SafeGenerationMixin:
                 prompt_txt = self._tokenizer.decode(
                     args[0][0].tolist(), skip_special_tokens=True
                 )
-
         if prompt_txt and self._is_toxic(prompt_txt):
             return self._safe_ids().unsqueeze(0)
 
         # 2) normal generation
-        output = super().generate(*args, **kwargs)
+        out = super().generate(*args, **kwargs)
 
-        # 3) output toxicity
+        # 3) output check
         if self._tokenizer is None:
-            return output
-        seqs = output.detach().cpu()
+            return out
+        seqs = out.detach().cpu()
         safe = []
         for seq in seqs:
-            txt = self._tokenizer.decode(seq.tolist(), skip_special_tokens=True)
-            if self._is_toxic(txt):
+            if self._is_toxic(self._tokenizer.decode(seq.tolist(), skip_special_tokens=True)):
                 safe.append(self._safe_ids(length=seq.size(0)))
             else:
                 safe.append(seq)
-        return torch.stack(safe, dim=0).to(self.device)
+        return torch.stack(safe, 0).to(self.device)
 
 
 # ---------------------------------------------------------------------
-# 2)  Resolve base class for the repo’s architecture string
+# 2)  helpers – resolve base class & cache dynamic subclass
 # ---------------------------------------------------------------------
 @lru_cache(None)
 def _get_base_cls(arch_name: str):
-    # direct attribute
     if hasattr(transformers, arch_name):
         return getattr(transformers, arch_name)
-
-    # heuristic import: e.g. LlamaForCausalLM -> transformers.models.llama.modeling_llama
     stem = arch_name.replace("ForCausalLM", "").lower()
-    module_path = f"transformers.models.{stem}.modeling_{stem}"
-    try:
-        mod = importlib.import_module(module_path)
-        return getattr(mod, arch_name)
-    except Exception as e:
-        raise ValueError(f"Cannot resolve base class for '{arch_name}': {e}") from e
+    mod = importlib.import_module(f"transformers.models.{stem}.modeling_{stem}")
+    return getattr(mod, arch_name)
 
 
 @lru_cache(None)
 def _make_dynamic_cls(base_cls):
-    """Create (and cache) SafeGeneration_<Base> = (Mixin, Base)."""
     return type(f"SafeGeneration_{base_cls.__name__}", (_SafeGenerationMixin, base_cls), {})
 
 
 # ---------------------------------------------------------------------
-# 3)  Dispatcher class – target in `auto_map`
+# 3)  dispatcher
 # ---------------------------------------------------------------------
 class SafeGenerationModel:
-    """
-    Thin dispatcher used by Hugging Face AutoClass.
-
-    It implements only `from_pretrained()`: determine the true base
-    architecture, build the dynamic subclass, and defer loading to it.
-    """
-
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-        # propagate trust_remote_code if caller set it
         kwargs.setdefault("trust_remote_code", True)
 
-        # 1) load config to know arch string
+        # 1) load config to know architecture
         config = transformers.AutoConfig.from_pretrained(
             pretrained_model_name_or_path, **kwargs
         )
         if not getattr(config, "architectures", None):
-            raise ValueError("`config.architectures` missing; cannot wrap model.")
+            raise ValueError("`config.architectures` missing – cannot wrap model.")
         arch_name = config.architectures[0]
 
-        # 2) build / retrieve dynamic subclass
+        # 2) dynamic subclass
         base_cls = _get_base_cls(arch_name)
         SafeCls = _make_dynamic_cls(base_cls)
 
-        # 3) delegate full loading
+        # 3) drop duplicate 'config' if caller already passed one
+        kwargs = {k: v for k, v in kwargs.items() if k != "config"}
+
+        # 4) delegate real loading
         return SafeCls.from_pretrained(
             pretrained_model_name_or_path,
             *model_args,