from transformers import HubertConfig
from typing import List

class HuBERTECGConfig(HubertConfig):
    
    model_type = "hubert_ecg"
    
    def __init__(self, ensemble_length: int = 1, vocab_sizes: List[int] = [100], **kwargs):
        super().__init__(**kwargs)
        self.ensemble_length = ensemble_length
        self.vocab_sizes = vocab_sizes if isinstance(vocab_sizes, list) else [vocab_sizes]