File size: 2,753 Bytes
6a0cf5a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
"""BabyloopConfig — HF ``PretrainedConfig`` 互換のモデル設定。

このファイルは ``save_pretrained`` 時に checkpoint ディレクトリへ複製され、
公式 evaluation-pipeline 側(別プロセス)から ``trust_remote_code=True`` で
import される。そのため **transformers と標準ライブラリ以外に依存しない**
(``babyloop`` パッケージ内部を import しない)こと。

フィールドと configs/model/*.yaml の対応は docs/architecture.md §6 を参照。
"""

from __future__ import annotations

from transformers import PretrainedConfig


class BabyloopConfig(PretrainedConfig):
    """ループドTransformer(K=1で標準に縮退)の設定。

    軸A(標準 vs 再帰)は ``k`` のみで切り替える単一実装。視覚(②④)の口も
    持つが、テキストのみ(①③)では ``fusion=None`` で無効。
    """

    model_type = "babyloop"

    def __init__(
        self,
        d_model: int = 768,
        n_layers: int = 12,
        n_heads: int = 12,
        ffn_hidden: int = 2048,
        n_prelude: int = 0,
        n_core: int = 12,
        n_coda: int = 0,
        k: int = 1,
        inject_input: bool = False,
        vocab_size: int = 16000,
        max_seq_len: int = 512,
        rope_base: float = 10000.0,
        rms_eps: float = 1e-5,
        tie_embeddings: bool = True,
        bias: bool = False,
        dropout: float = 0.0,
        fusion: str | None = None,
        pad_token_id: int | None = None,
        bos_token_id: int | None = None,
        eos_token_id: int | None = None,
        **kwargs,
    ):
        self.d_model = d_model
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.ffn_hidden = ffn_hidden
        self.n_prelude = n_prelude
        self.n_core = n_core
        self.n_coda = n_coda
        self.k = k
        self.inject_input = inject_input
        self.vocab_size = vocab_size
        self.max_seq_len = max_seq_len
        self.rope_base = rope_base
        self.rms_eps = rms_eps
        self.tie_embeddings = tie_embeddings
        self.bias = bias
        self.dropout = dropout
        self.fusion = fusion

        # HF 汎用ユーティリティが参照するエイリアス。
        self.hidden_size = d_model
        self.num_attention_heads = n_heads
        self.num_hidden_layers = n_layers
        self.max_position_embeddings = max_seq_len

        super().__init__(
            pad_token_id=pad_token_id,
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            tie_word_embeddings=tie_embeddings,
            **kwargs,
        )

    @property
    def head_dim(self) -> int:
        return self.d_model // self.n_heads