277 lines
11 KiB
Python
277 lines
11 KiB
Python
|
"""
|
|||
|
|
custom_modeling.py – model-agnostic toxicity and prompt injection wrapper
|
|||
|
|
--------------------------------------------------------------------------
|
|||
|
|
Place in repo root together with:
|
|||
|
|
• toxic.keras
|
|||
|
|
• PI.keras
|
|||
|
|
Add to config.json:
|
|||
|
|
"auto_map": { "AutoModelForCausalLM": "custom_modeling.SafeGenerationModel" }
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import importlib
|
|||
|
|
import os
|
|||
|
|
import logging
|
|||
|
|
from functools import lru_cache
|
|||
|
|
|
|||
|
|
import torch
|
|||
|
|
import transformers
|
|||
|
|
import tensorflow as tf
|
|||
|
|
import keras
|
|||
|
|
from huggingface_hub import hf_hub_download
|
|||
|
|
|
|||
|
|
# Configure logging
|
|||
|
|
logger = logging.getLogger(__name__)
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------ #
|
|||
|
|
# 1) MIXIN – toxicity and prompt injection filtering logic #
|
|||
|
|
# ------------------------------------------------------------------ #
|
|||
|
|
class _SafeGenerationMixin:
|
|||
|
|
_toxicity_model = None
|
|||
|
|
_pi_model = None
|
|||
|
|
_tox_threshold = 0.8
|
|||
|
|
_pi_threshold = 0.8
|
|||
|
|
|
|||
|
|
# Safety messages
|
|||
|
|
_safe_in_msg = "Sorry, I can't help with that toxic input."
|
|||
|
|
_safe_out_msg = "I'm sorry, but I can't continue with that response."
|
|||
|
|
_pi_in_msg = "PI detected at Input level"
|
|||
|
|
_pi_out_msg = "PI detected at output level"
|
|||
|
|
|
|||
|
|
_tokenizer = None
|
|||
|
|
|
|||
|
|
# ---- helpers ----------------------------------------------------
|
|||
|
|
def _device(self):
|
|||
|
|
return next(self.parameters()).device
|
|||
|
|
|
|||
|
|
def _is_local_path(self, path_or_repo):
|
|||
|
|
"""Check if the path is a local directory rather than a HF repo ID"""
|
|||
|
|
return os.path.isdir(path_or_repo) or os.path.isabs(path_or_repo)
|
|||
|
|
|
|||
|
|
def _get_model_file_path(self, filename):
|
|||
|
|
"""Get path to model file, supporting both local and remote repositories"""
|
|||
|
|
if self._is_local_path(self.config.name_or_path):
|
|||
|
|
# Local path - look for file in the model directory
|
|||
|
|
local_path = os.path.join(self.config.name_or_path, filename)
|
|||
|
|
if os.path.exists(local_path):
|
|||
|
|
return local_path
|
|||
|
|
else:
|
|||
|
|
# Fallback: try to download from HF if local file doesn't exist
|
|||
|
|
# This handles cases where someone has a local model but the .keras files are on HF
|
|||
|
|
try:
|
|||
|
|
# Extract just the repo name from the path for HF download
|
|||
|
|
repo_name = os.path.basename(self.config.name_or_path.rstrip('/'))
|
|||
|
|
if '/' in repo_name or len(repo_name.split('-')) > 1:
|
|||
|
|
# Try using the directory name as repo_id
|
|||
|
|
return hf_hub_download(repo_id=repo_name, filename=filename)
|
|||
|
|
else:
|
|||
|
|
raise FileNotFoundError(f"Could not find {filename} locally or determine HF repo")
|
|||
|
|
except Exception as hf_error:
|
|||
|
|
raise FileNotFoundError(
|
|||
|
|
f"{filename} not found at {local_path}. "
|
|||
|
|
f"Also failed to download from HF: {hf_error}"
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
# Remote repo - download from HF Hub
|
|||
|
|
try:
|
|||
|
|
return hf_hub_download(
|
|||
|
|
repo_id=self.config.name_or_path,
|
|||
|
|
filename=filename,
|
|||
|
|
)
|
|||
|
|
except Exception as hf_error:
|
|||
|
|
# Fallback: check if it's actually a local path that wasn't detected
|
|||
|
|
if os.path.exists(self.config.name_or_path):
|
|||
|
|
local_path = os.path.join(self.config.name_or_path, filename)
|
|||
|
|
if os.path.exists(local_path):
|
|||
|
|
return local_path
|
|||
|
|
raise FileNotFoundError(
|
|||
|
|
f"Could not download {filename} from HF repo '{self.config.name_or_path}': {hf_error}"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def _tox_model(self):
|
|||
|
|
if self._toxicity_model is None:
|
|||
|
|
try:
|
|||
|
|
path = self._get_model_file_path("toxic.keras")
|
|||
|
|
# Load .keras format directly
|
|||
|
|
self._toxicity_model = keras.models.load_model(path, compile=False)
|
|||
|
|
logger.info("Toxicity model loaded successfully")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Failed to load toxicity model: {e}")
|
|||
|
|
raise RuntimeError(f"Could not load required toxicity model: {e}")
|
|||
|
|
return self._toxicity_model
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def _prompt_injection_model(self):
|
|||
|
|
if self._pi_model is None:
|
|||
|
|
try:
|
|||
|
|
path = self._get_model_file_path("Proooo33_fine_tuned.keras")
|
|||
|
|
# Load .keras format directly
|
|||
|
|
self._pi_model = keras.models.load_model(path, compile=False)
|
|||
|
|
logger.info("Prompt injection model loaded successfully")
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Failed to load prompt injection model: {e}")
|
|||
|
|
raise RuntimeError(f"Could not load required prompt injection model: {e}")
|
|||
|
|
return self._pi_model
|
|||
|
|
|
|||
|
|
def _ensure_tokenizer(self):
|
|||
|
|
if self._tokenizer is None:
|
|||
|
|
try:
|
|||
|
|
self._tokenizer = transformers.AutoTokenizer.from_pretrained(
|
|||
|
|
self.config.name_or_path, trust_remote_code=True
|
|||
|
|
)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Failed to load tokenizer: {e}")
|
|||
|
|
|
|||
|
|
def _is_toxic(self, text: str) -> bool:
|
|||
|
|
if not text.strip():
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# Ensure CPU execution for compatibility
|
|||
|
|
with tf.device('/CPU:0'):
|
|||
|
|
inputs = tf.constant([text], dtype=tf.string)
|
|||
|
|
|
|||
|
|
# Handle both Keras models and SavedModel formats
|
|||
|
|
if hasattr(self._tox_model, 'predict'):
|
|||
|
|
prob = float(self._tox_model.predict(inputs, verbose=0)[0, 0])
|
|||
|
|
else:
|
|||
|
|
# For SavedModel format
|
|||
|
|
prob = float(self._tox_model(inputs).numpy()[0, 0])
|
|||
|
|
|
|||
|
|
return prob >= self._tox_threshold
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Toxicity prediction failed: {e}")
|
|||
|
|
# Don't fallback to rule-based - let it fail if models don't work
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
def _has_prompt_injection(self, text: str) -> bool:
|
|||
|
|
if not text.strip():
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# Ensure CPU execution for compatibility
|
|||
|
|
with tf.device('/CPU:0'):
|
|||
|
|
inputs = tf.constant([text], dtype=tf.string)
|
|||
|
|
|
|||
|
|
# Handle both Keras models and SavedModel formats
|
|||
|
|
if hasattr(self._prompt_injection_model, 'predict'):
|
|||
|
|
prob = float(self._prompt_injection_model.predict(inputs, verbose=0)[0, 0])
|
|||
|
|
else:
|
|||
|
|
# For SavedModel format
|
|||
|
|
prob = float(self._prompt_injection_model(inputs).numpy()[0, 0])
|
|||
|
|
|
|||
|
|
return prob >= self._pi_threshold
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error(f"Prompt injection prediction failed: {e}")
|
|||
|
|
# Don't fallback to rule-based - let it fail if models don't work
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
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(message, return_tensors="pt")["input_ids"][0]
|
|||
|
|
if length is not None:
|
|||
|
|
pad_id = (
|
|||
|
|
self.config.eos_token_id
|
|||
|
|
if self.config.eos_token_id is not None
|
|||
|
|
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
|
|||
|
|
)
|
|||
|
|
else:
|
|||
|
|
ids = ids[:length]
|
|||
|
|
return ids.to(self._device())
|
|||
|
|
|
|||
|
|
# ---- main override ---------------------------------------------
|
|||
|
|
def generate(self, *args, **kwargs):
|
|||
|
|
self._ensure_tokenizer()
|
|||
|
|
|
|||
|
|
# 1) Extract prompt text
|
|||
|
|
prompt_txt = None
|
|||
|
|
if self._tokenizer is not None:
|
|||
|
|
if "input_ids" in kwargs:
|
|||
|
|
prompt_txt = self._tokenizer.decode(
|
|||
|
|
kwargs["input_ids"][0].tolist(), skip_special_tokens=True
|
|||
|
|
)
|
|||
|
|
elif args:
|
|||
|
|
prompt_txt = self._tokenizer.decode(
|
|||
|
|
args[0][0].tolist(), skip_special_tokens=True
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
# 2) Check input for prompt injection (higher priority)
|
|||
|
|
if prompt_txt and self._has_prompt_injection(prompt_txt):
|
|||
|
|
return self._safe_ids(self._pi_in_msg).unsqueeze(0)
|
|||
|
|
|
|||
|
|
# 3) Check input for toxicity
|
|||
|
|
if prompt_txt and self._is_toxic(prompt_txt):
|
|||
|
|
return self._safe_ids(self._safe_in_msg).unsqueeze(0)
|
|||
|
|
|
|||
|
|
# 4) Normal generation
|
|||
|
|
outputs = super().generate(*args, **kwargs)
|
|||
|
|
|
|||
|
|
# 5) Check outputs for safety violations
|
|||
|
|
if self._tokenizer is None:
|
|||
|
|
return outputs
|
|||
|
|
|
|||
|
|
new_seqs = []
|
|||
|
|
for seq in outputs.detach().cpu():
|
|||
|
|
txt = self._tokenizer.decode(seq.tolist(), skip_special_tokens=True)
|
|||
|
|
|
|||
|
|
# Check for prompt injection first (higher priority)
|
|||
|
|
if self._has_prompt_injection(txt):
|
|||
|
|
new_seqs.append(self._safe_ids(self._pi_out_msg, length=seq.size(0)))
|
|||
|
|
# Then check for toxicity
|
|||
|
|
elif 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 base class & cache subclass #
|
|||
|
|
# ------------------------------------------------------------------ #
|
|||
|
|
@lru_cache(None)
|
|||
|
|
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)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@lru_cache(None)
|
|||
|
|
def _make_safe_subclass(base_cls):
|
|||
|
|
return type(
|
|||
|
|
f"SafeGeneration_{base_cls.__name__}",
|
|||
|
|
(_SafeGenerationMixin, base_cls),
|
|||
|
|
{},
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ------------------------------------------------------------------ #
|
|||
|
|
# 3) Dispatcher class – referenced by auto_map #
|
|||
|
|
# ------------------------------------------------------------------ #
|
|||
|
|
class SafeGenerationModel:
|
|||
|
|
@classmethod
|
|||
|
|
def from_pretrained(cls, repo_id, *model_args, **kwargs):
|
|||
|
|
kwargs.setdefault("trust_remote_code", True)
|
|||
|
|
if kwargs.get("torch_dtype") == "auto":
|
|||
|
|
kwargs.pop("torch_dtype")
|
|||
|
|
|
|||
|
|
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]
|
|||
|
|
|
|||
|
|
Base = _get_base_cls(arch_str)
|
|||
|
|
Safe = _make_safe_subclass(Base)
|
|||
|
|
|
|||
|
|
kwargs.pop("config", None) # avoid duplicate
|
|||
|
|
return Safe.from_pretrained(repo_id, *model_args, config=config, **kwargs)
|