File size: 2,826 Bytes
2475d58 92e4399 2475d58 92e4399 2475d58 92e4399 2475d58 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | """Prompt-injection classifier on top of the axiotic/ogma-base encoder.
Loadable with AutoModelForSequenceClassification / pipeline via trust_remote_code.
The base encoder declares its rotary position caches as non-persistent buffers,
so a vanilla from_pretrained leaves them uninitialised (NaN/zeros, nondeterministic
per process). We rebuild them from inv_freq AND re-register them as PERSISTENT, so
the correct values are saved into this model's weights and loaded reliably — the
NaN bug cannot reach downstream users.
"""
import torch
from torch import nn
from transformers import PreTrainedModel, AutoModel, AutoConfig
from transformers.modeling_outputs import SequenceClassifierOutput
try:
from .configuration_ogma_classifier import OgmaClassifierConfig
except ImportError: # loaded standalone (not as a package)
from configuration_ogma_classifier import OgmaClassifierConfig
class OgmaForPromptInjection(PreTrainedModel):
config_class = OgmaClassifierConfig
def __init__(self, config):
super().__init__(config)
base_cfg = AutoConfig.from_pretrained(config.base_model_id, trust_remote_code=True)
self.encoder = AutoModel.from_config(base_cfg, trust_remote_code=True)
self.dropout = nn.Dropout(0.1)
self.head = nn.Linear(config.hidden, config.num_labels)
self._rope_ready = False
self.post_init()
def _ensure_rope(self):
"""Rebuild the encoder's rotary caches from inv_freq. The base declares
cos/sin as non-persistent buffers, so from_pretrained leaves them
uninitialised; rebuilding here makes inference correct and deterministic.
Done lazily on first forward so it survives any instantiation path
(from_config, meta/fast-init, etc.).
The encoder's _build_cache builds its position grid on CPU (torch.arange
with no device), so a rebuild while the module sits on MPS/CUDA raises a
device-mismatch. Rebuild each tiny RoPE module on CPU, then move it back
to its original device — correct on every backend, not just CPU."""
for m in self.modules():
if hasattr(m, "_build_cache") and hasattr(m, "cos_cached"):
dev = m.cos_cached.device
m.to("cpu")
m._build_cache(m.cos_cached.shape[0])
m.to(dev)
self._rope_ready = True
def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
if not self._rope_ready:
self._ensure_rope()
emb = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
logits = self.head(self.dropout(emb))
loss = None
if labels is not None:
loss = nn.functional.cross_entropy(logits, labels)
return SequenceClassifierOutput(loss=loss, logits=logits)
|