from transformers import Wav2Vec2Config


class Wav2Vec2DualHypersphereConfig(Wav2Vec2Config):
    model_type = "wav2vec2-dual-hypersphere"

    def __init__(
        self,
        classifier_proj_size=256,
        final_dropout=0.1,
        freeze_feat_extract_train=True,
        **kwargs,
    ):
        # Ensure default classification metadata is set before calling super
        if "num_labels" not in kwargs:
            kwargs["num_labels"] = 2

        if "id2label" not in kwargs:
            kwargs["id2label"] = {0: "AI Voice (Fake)", 1: "Human Voice (Real)"}

        if "label2id" not in kwargs:
            kwargs["label2id"] = {"AI Voice (Fake)": 0, "Human Voice (Real)": 1}

        super().__init__(**kwargs)
        self.classifier_proj_size = classifier_proj_size
        self.final_dropout = final_dropout
        self.freeze_feat_extract_train = freeze_feat_extract_train


__all__ = ["Wav2Vec2DualHypersphereConfig"]