| """DLC SEDD config for Hugging Face.""" | |
| import transformers | |
| class DLCSEDDGraphConfig(transformers.PretrainedConfig): | |
| def __init__(self, graph_type="absorb", **kwargs): | |
| super().__init__(**kwargs) | |
| self.type = graph_type | |
| self.file = "data" | |
| self.report_all = False | |
| class DLCSEDDNoiseConfig(transformers.PretrainedConfig): | |
| def __init__(self, noise_type="loglinear", sigma_min=0.0001, sigma_max=20, **kwargs): | |
| super().__init__(*kwargs) | |
| self.type = noise_type | |
| self.sigma_min = sigma_min | |
| self.sigma_max = sigma_max | |
| class DLCSEDDConfig(transformers.PretrainedConfig): | |
| """Hugging Face configuration class for dlc sedd.""" | |
| model_type = "dlc_sedd" | |
| def __init__( | |
| self, | |
| length: int = 512, | |
| hidden_size: int = 1024, | |
| cond_dim: int = 128, | |
| n_blocks: int = 24, | |
| n_heads: int = 16, | |
| dropout: float = 0.1, | |
| time_conditioning: bool = True, | |
| scale_by_sigma: bool = True, | |
| shared_embedding: bool = True, | |
| graph_type: str = "absorb", | |
| tokens: float = 256, | |
| **kwargs, | |
| ): | |
| super().__init__(**kwargs) | |
| self.length = length | |
| self.hidden_size = hidden_size | |
| self.cond_dim = cond_dim | |
| self.n_blocks = n_blocks | |
| self.n_heads = n_heads | |
| self.dropout = dropout | |
| self.time_conditioning = time_conditioning | |
| self.scale_by_sigma = scale_by_sigma | |
| self.shared_embedding = shared_embedding | |
| self.graph = DLCSEDDGraphConfig() | |
| self.noise = DLCSEDDNoiseConfig() | |
| self.tokens = tokens | |