# Copyright (c) 2026, Bodhan.  All rights reserved.
# Licensed under the Apache License, Version 2.0.
"""IndicCanary model configuration (HF-style port of NeMo EncDecMultiTaskModel)."""

from transformers import PretrainedConfig


class IndicCanaryConfig(PretrainedConfig):
    model_type = "indic_canary"

    def __init__(
        self,
        vocab_size=7152,
        d_model=1024,
        num_mel_bins=128,
        encoder_layers=32,
        encoder_attention_heads=8,
        encoder_ffn_dim=4096,
        conv_kernel_size=9,
        subsampling_factor=8,
        subsampling_conv_channels=256,
        decoder_layers=24,
        decoder_attention_heads=8,
        decoder_ffn_dim=4096,
        max_target_positions=1024,
        max_generation_delta=50,
        pad_token_id=2,
        eos_token_id=3,
        bos_token_id=4,
        # the frozen canary2 prompt starts with <|startofcontext|> (7); generate()
        # prepends decoder_start_token_id when the prompt doesn't start with it
        decoder_start_token_id=7,
        is_encoder_decoder=True,
        tie_word_embeddings=True,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.num_mel_bins = num_mel_bins
        self.encoder_layers = encoder_layers
        self.encoder_attention_heads = encoder_attention_heads
        self.encoder_ffn_dim = encoder_ffn_dim
        self.conv_kernel_size = conv_kernel_size
        self.subsampling_factor = subsampling_factor
        self.subsampling_conv_channels = subsampling_conv_channels
        self.decoder_layers = decoder_layers
        self.decoder_attention_heads = decoder_attention_heads
        self.decoder_ffn_dim = decoder_ffn_dim
        self.max_target_positions = max_target_positions
        self.max_generation_delta = max_generation_delta
        super().__init__(
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
            bos_token_id=bos_token_id,
            decoder_start_token_id=decoder_start_token_id,
            is_encoder_decoder=is_encoder_decoder,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )