| from transformers import PretrainedConfig, AutoConfig |
|
|
| class DiscreteDiffusionConfig(PretrainedConfig): |
| model_type = "discrete_diffusion" |
| |
| def __init__( |
| self, |
| backbone_config=None, |
| num_diffusion_timesteps=50, |
| diffusion_type="absorbing", |
| attention_strategy="full", |
| vocab_pad_to_multiple=1, |
| lora=False, |
| lora_target_modules=["query", "value"], |
| lora_alpha=16, |
| lora_rank=16, |
| lora_bias="none", |
| lora_dropout=0, |
| mask_token_id=None, |
| bos_token_id=None, |
| eos_token_id=None, |
| pad_token_id=None, |
| argmax_decoding=True, |
| **kwargs |
| ): |
| super().__init__(**kwargs) |
| self.backbone_config = backbone_config |
| self.num_diffusion_timesteps = num_diffusion_timesteps |
| self.diffusion_type = diffusion_type |
| self.attention_strategy = attention_strategy |
| self.vocab_pad_to_multiple = vocab_pad_to_multiple |
| self.lora = lora |
| self.lora_target_modules = lora_target_modules |
| self.lora_alpha = lora_alpha |
| self.lora_rank = lora_rank |
| self.lora_bias = lora_bias |
| self.lora_dropout = lora_dropout |
| self.mask_token_id = mask_token_id |
| self.bos_token_id = bos_token_id |
| self.eos_token_id = eos_token_id |
| self.pad_token_id = pad_token_id |
| self.argmax_decoding = argmax_decoding |
|
|
| if backbone_config is None: |
| self.backbone_config = AutoConfig.from_pretrained("FacebookAI/xlm-roberta-large").to_dict() |
| elif isinstance(backbone_config, PretrainedConfig): |
| self.backbone_config = backbone_config.to_dict() |
| else: |
| self.backbone_config = backbone_config |
| |
| |
| self.hidden_size = self.backbone_config.get("hidden_size", 1024) |
| self.num_attention_heads = self.backbone_config.get("num_attention_heads", 16) |
| self.intermediate_size = self.backbone_config.get("intermediate_size", 4096) |
| self.max_position_embeddings = self.backbone_config.get("max_position_embeddings", 514) |
| self.tie_word_embeddings = self.backbone_config.get("tie_word_embeddings", True) |
|
|
|
|
