| """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 |
|
|
| |
| 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 |
|
|