ogma-prompt-injection / configuration_ogma_classifier.py
sam-at-axiotic's picture
ogma-prompt-injection v2: generated-benign retrain, label smoothing, RoPE fix baked in
2475d58 verified
Raw
History Blame Contribute Delete
476 Bytes
"""Config for the ogma prompt-injection classifier (trust_remote_code)."""
from transformers import PretrainedConfig
class OgmaClassifierConfig(PretrainedConfig):
model_type = "ogma_prompt_injection"
def __init__(self, base_model_id="axiotic/ogma-base", hidden=256, max_len=512, num_labels=2, **kwargs):
self.base_model_id = base_model_id
self.hidden = hidden
self.max_len = max_len
super().__init__(num_labels=num_labels, **kwargs)