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) |