| from typing import Optional | |
| from transformers import PretrainedConfig | |
| class LUARConfig(PretrainedConfig): | |
| model_type = "LUAR" | |
| def __init__(self, | |
| embedding_size: int = 512, | |
| use_memory_efficient_attention=False, | |
| q_bucket_size=512, | |
| k_bucket_size=1024, | |
| upstream_transformer_revision: Optional[str] = None, | |
| **kwargs, | |
| ): | |
| self.embedding_size = embedding_size | |
| self.use_memory_efficient_attention = use_memory_efficient_attention | |
| self.q_bucket_size = q_bucket_size | |
| self.k_bucket_size = k_bucket_size | |
| self.upstream_transformer_revision = upstream_transformer_revision | |
| super().__init__(**kwargs) | |