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