(
hidden_size: int,
num_hidden_layers: int,
intermediate_size: float,
parallel_attn: bool = True,
vocab_size: int = 32768,
model_type: ModelType = ModelType.mlm,
sliding_window: int = -1,
global_attn_every_n_layers: int = -1,
normalization: str = "layernorm",
compile_model: bool = True,
masked_prediction: bool = False,
)
| 43 | |
| 44 | |
| 45 | def get_model( |
| 46 | hidden_size: int, |
| 47 | num_hidden_layers: int, |
| 48 | intermediate_size: float, |
| 49 | parallel_attn: bool = True, |
| 50 | vocab_size: int = 32768, |
| 51 | model_type: ModelType = ModelType.mlm, |
| 52 | sliding_window: int = -1, |
| 53 | global_attn_every_n_layers: int = -1, |
| 54 | normalization: str = "layernorm", |
| 55 | compile_model: bool = True, |
| 56 | masked_prediction: bool = False, |
| 57 | ): |
| 58 | config = FlexBertConfig( |
| 59 | num_attention_heads=hidden_size // 64, |
| 60 | hidden_size=hidden_size, |
| 61 | num_hidden_layers=num_hidden_layers, |
| 62 | intermediate_size=intermediate_size, |
| 63 | vocab_size=vocab_size, |
| 64 | attention_layer="rope_parallel" if parallel_attn else "rope", |
| 65 | attention_probs_dropout_prob=0.0, |
| 66 | attn_out_bias=False, |
| 67 | attn_out_dropout_prob=0.0, |
| 68 | attn_qkv_bias=False, |
| 69 | bert_layer="parallel_prenorm" if parallel_attn else "prenorm", |
| 70 | decoder_bias=False, |
| 71 | embed_dropout_prob=0.0, |
| 72 | embed_norm=True, |
| 73 | final_norm=False, |
| 74 | embedding_layer="sans_pos", |
| 75 | encoder_layer="base", |
| 76 | hidden_act="gelu", |
| 77 | loss_function="fa_cross_entropy", |
| 78 | loss_kwargs={"reduction": "mean"}, |
| 79 | mlp_dropout_prob=0.0, |
| 80 | mlp_in_bias=False, |
| 81 | mlp_layer="parallel_glu" if parallel_attn else "glu", |
| 82 | mlp_out_bias=False, |
| 83 | norm_kwargs={"eps": 1e-5}, |
| 84 | normalization=normalization, |
| 85 | padding="padded", |
| 86 | head_class_act="silu", |
| 87 | head_class_bias=False, |
| 88 | head_class_dropout=0.0, |
| 89 | head_class_norm=False, |
| 90 | head_pred_act="gelu", |
| 91 | head_pred_bias=False, |
| 92 | head_pred_dropout=0.0, |
| 93 | head_pred_norm=True, |
| 94 | pooling_type="cls", |
| 95 | rotary_emb_dim=None, |
| 96 | rotary_emb_base=10000.0, |
| 97 | rotary_emb_scale_base=None, |
| 98 | rotary_emb_interleaved=False, |
| 99 | use_fa2=True, |
| 100 | use_sdpa_attn_mask=False, |
| 101 | allow_embedding_resizing=False, |
| 102 | init_method="default", |
no test coverage detected