from speechbrain.pretrained import EncoderClassifier
from transformers import AutoProcessor, Wav2Vec2Processor


class CustomEncoderWav2vec2Classifier(EncoderClassifier):
    """
    This class inherits from `EncoderClassifier` to handle custom models.

    The only difference is that it uses `AutoProcessor` to load the tokenizer.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.tokenizer = AutoProcessor.from_pretrained(
            self.hparams.source, trust_remote_code=True
        )