Commit 6403401f authored by Sachin Kadyan's avatar Sachin Kadyan
Browse files

Added configuration options for the new PreembeddingEmbedder.

- Added `preembedding_embedder` config dictionary in `config`
- Added `preemb_dim_size` property in `config` for specifying single seq embedding size.
parent 75889e9a
......@@ -189,6 +189,7 @@ c_m = mlc.FieldReference(256, field_type=int)
c_t = mlc.FieldReference(64, field_type=int)
c_e = mlc.FieldReference(64, field_type=int)
c_s = mlc.FieldReference(384, field_type=int)
preemb_dim_size = mlc.FieldReference(1280, field_type=int)
blocks_per_ckpt = mlc.FieldReference(None, field_type=int)
chunk_size = mlc.FieldReference(4, field_type=int)
aux_distogram_bins = mlc.FieldReference(64, field_type=int)
......@@ -403,6 +404,13 @@ config = mlc.ConfigDict(
"c_m": c_m,
"relpos_k": 32,
},
"preembedding_embedder": {
"tf_dim": 22,
"preembedding_dim": preemb_dim_size,
"c_z": c_z,
"c_m": c_m,
"relpos_k": 32,
},
"recycling_embedder": {
"c_z": c_z,
"c_m": c_m,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment