timorobrecht/10k_initials_gpu1

1"""Configuration for the Transformers-compatible pinyin-code causal LM."""
2
3from __future__ import annotations
4
5from transformers import PretrainedConfig
6
7
8class PinyinCodeConfig(PretrainedConfig):
9    """Configuration for the compact GPT-style pinyin-code decoder."""
10
11    model_type = "pinyin_code"
12
13    def __init__(
14        self,
15        vocab_size: int = 8000,
16        block_size: int = 128,
17        n_layer: int = 6,
18        n_head: int = 8,
19        n_embd: int = 256,
20        dropout: float = 0.1,
21        bos_token_id: int | None = None,
22        eos_token_id: int | None = None,
23        pad_token_id: int | None = None,
24        unk_token_id: int | None = None,
25        **kwargs,
26    ) -> None:
27        super().__init__(
28            bos_token_id=bos_token_id,
29            eos_token_id=eos_token_id,
30            pad_token_id=pad_token_id,
31            unk_token_id=unk_token_id,
32            **kwargs,
33        )
34        self.vocab_size = vocab_size
35        self.block_size = block_size
36        self.n_layer = n_layer
37        self.n_head = n_head
38        self.n_embd = n_embd
39        self.dropout = dropout
40        self.num_hidden_layers = n_layer
41        self.num_attention_heads = n_head
42        self.hidden_size = n_embd
43        self.max_position_embeddings = block_size
44        self.is_decoder = True
45        self.is_encoder_decoder = False
46        self.use_cache = False
47