File size: 476 Bytes
2475d58 | 1 2 3 4 5 6 7 8 9 10 11 12 13 | """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)
|