from transformers import PreTrainedModel, PretrainedConfig
import json

class SetuTranslationConfig(PretrainedConfig):
    """Configuration class for SETU Translation model.
    
    This class handles the configuration for the SETU (Script-agnostic English Translation Unifier) model
    which translates multiscript, multilingual, and informal text into clean, formal English.
    """
    
    model_type = "setu_translation"
    
    def __init__(
        self,
        model_name: str = "SETU",
        full_name: str = "Script-agnostic English Translation Unifier", 
        description: str = "A neural translation model that unifies multiscript, multilingual, and informal text into clean, formal English",
        version: str = "1.0.0",
        architecture: str = "transformer_iwslt_de_en",
        src_vocab_size: int = 40253,
        tgt_vocab_size: int = 40253,
        bos_idx: int = 0,
        eos_idx: int = 2,
        pad_idx: int = 1,
        unk_idx: int = 3,
        beam_size: int = 5,
        max_len: int = 200,
        len_penalty: float = 1.0,
        capabilities: list = None,
        **kwargs
    ):
        super().__init__(**kwargs)
        
        self.model_name = model_name
        self.full_name = full_name
        self.description = description
        self.version = version
        self.architecture = architecture
        self.src_vocab_size = src_vocab_size
        self.tgt_vocab_size = tgt_vocab_size
        self.bos_idx = bos_idx
        self.eos_idx = eos_idx
        self.pad_idx = pad_idx
        self.unk_idx = unk_idx
        self.beam_size = beam_size
        self.max_len = max_len
        self.len_penalty = len_penalty
        
        if capabilities is None:
            capabilities = [
                "Romanized Nepali to English",
                "Devanagari Nepali to English", 
                "Code-mixed text to English",
                "Informal/slang to formal English"
            ]
        self.capabilities = capabilities