from transformers import PretrainedConfig

class SpeechLLMModelConfig(PretrainedConfig):
    model_type = "custom_model"
    
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.audio_enc_dim = 1280
        self.llm_dim = 2048

        self.audio_processor_name = "facebook/hubert-large-ls960-ft"
        self.audio_encoder_name = 'facebook/hubert-xlarge-ll60k'
        self.llm_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
        self.llm_model_checkpoint = "hf_repo/llm_model_checkpoint"