in summarize_from_feedback/models/transformer.py [0:0]
def __init__(
self,
n_ctx,
n_vocab,
d_model=128,
n_layer=2,
heads=1,
attn_dropout=0.0,
resid_dropout=0.0,
emb_dropout=0.0,
zero_out=False,
init_scale=1.0,
res_scale=False,
m_attn=0.25,
m_mlp=1,
mp_comm: Comm = None,
include_pos_embeddings=True,
include_input_embeddings=True,
include_output_unembeddings=True,
# For e.g. reward model training, we want the final layer norm before the extra head
# but not the output unembeddings, so control this separately.
include_final_layer_norm=True,
afn="quick_gelu",
key_bias=False,
flatten_multi_index_batch_dims=False, # allows the first dims to be batch dims and flattens them, outputs will be flattened
global_idxs_for_resblocks: Optional[List[int]] = None,
**extra_args,