Initial commit

.gitattributes

@@ -1,6 +1,7 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
@@ -9,14 +10,10 @@
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
@@ -27,5 +24,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoint*/** filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+__pycache__/

config.json

@@ -0,0 +1,32 @@
+{
+  "_name_or_path": "Finnish-NLP/t5-small-nl16-finnish",
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 512,
+  "decoder_start_token_id": 0,
+  "dense_act_fn": "gelu_new",
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "n_positions": 512,
+  "num_decoder_layers": 16,
+  "num_heads": 8,
+  "num_layers": 16,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.20.1",
+  "use_cache": true,
+  "vocab_size": 32128
+}

convert_t5x_checkpoint_to_flax.py

@@ -0,0 +1,157 @@
+# https://gist.github.com/stefan-it/30e4998ef159f33696e377a46f699d9f
+
+import argparse
+
+from t5x import checkpoints
+from transformers import T5Config, FlaxT5ForConditionalGeneration
+
+
+def convert_t5x_checkpoint_to_flax(t5x_checkpoint_path, config_name, flax_dump_folder_path):
+    config = T5Config.from_pretrained(config_name)
+    flax_model = FlaxT5ForConditionalGeneration(config=config)
+    t5x_model = checkpoints.load_t5x_checkpoint(t5x_checkpoint_path)
+
+    split_mlp_wi = "wi_0" in t5x_model["target"]["encoder"]["layers_0"]["mlp"]
+
+    # Encoder
+    for layer_index in range(config.num_layers):
+        layer_name = f"layers_{str(layer_index)}"
+        
+        # Self-Attention
+        t5x_attention_key = t5x_model["target"]["encoder"][layer_name]["attention"]["key"]["kernel"]
+        t5x_attention_out = t5x_model["target"]["encoder"][layer_name]["attention"]["out"]["kernel"]
+        t5x_attention_query = t5x_model["target"]["encoder"][layer_name]["attention"]["query"]["kernel"]
+        t5x_attention_value = t5x_model["target"]["encoder"][layer_name]["attention"]["value"]["kernel"]
+
+        ## Layer Normalization
+        t5x_attention_layer_norm = t5x_model["target"]["encoder"][layer_name]["pre_attention_layer_norm"]["scale"]
+        
+        if split_mlp_wi:
+            t5x_mlp_wi_0 = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi_0"]["kernel"]
+            t5x_mlp_wi_1 = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi_1"]["kernel"]
+        else:
+            t5x_mlp_wi = t5x_model["target"]["encoder"][layer_name]["mlp"]["wi"]["kernel"]
+        
+        t5x_mlp_wo = t5x_model["target"]["encoder"][layer_name]["mlp"]["wo"]["kernel"]
+        
+        ## Layer Normalization
+        t5x_mlp_layer_norm = t5x_model["target"]["encoder"][layer_name]["pre_mlp_layer_norm"]["scale"]
+        
+        # Assigning
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["k"]["kernel"] = t5x_attention_key
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["o"]["kernel"] = t5x_attention_out
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["q"]["kernel"] = t5x_attention_query
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["v"]["kernel"] = t5x_attention_value
+        
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["0"]["layer_norm"]["weight"] = t5x_attention_layer_norm
+        
+        if split_mlp_wi:
+            flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["DenseReluDense"]["wi_0"]["kernel"] = t5x_mlp_wi_0
+            flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["DenseReluDense"]["wi_1"]["kernel"] = t5x_mlp_wi_1
+        else:
+            flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["DenseReluDense"]["wi"]["kernel"] = t5x_mlp_wi
+
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["DenseReluDense"]["wo"]["kernel"] = t5x_mlp_wo
+        flax_model.params["encoder"]["block"][str(layer_index)]["layer"]["1"]["layer_norm"]["weight"] = t5x_mlp_layer_norm
+
+    # Only for layer 0:
+    t5x_encoder_rel_embedding = t5x_model["target"]["encoder"]["relpos_bias"]["rel_embedding"].T
+    flax_model.params["encoder"]["block"]["0"]["layer"]["0"]["SelfAttention"]["relative_attention_bias"]["embedding"] = t5x_encoder_rel_embedding
+
+    # Assigning
+    t5x_encoder_norm = t5x_model["target"]["encoder"]["encoder_norm"]["scale"]
+    flax_model.params["encoder"]["final_layer_norm"]["weight"] = t5x_encoder_norm
+
+    # Decoder
+    for layer_index in range(config.num_layers):
+        layer_name = f"layers_{str(layer_index)}"
+
+        # Self-Attention
+        t5x_attention_key = t5x_model["target"]["decoder"][layer_name]["self_attention"]["key"]["kernel"]
+        t5x_attention_out = t5x_model["target"]["decoder"][layer_name]["self_attention"]["out"]["kernel"]
+        t5x_attention_query = t5x_model["target"]["decoder"][layer_name]["self_attention"]["query"]["kernel"]
+        t5x_attention_value = t5x_model["target"]["decoder"][layer_name]["self_attention"]["value"]["kernel"]
+
+        ## Layer Normalization
+        t5x_pre_attention_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_self_attention_layer_norm"]["scale"]
+
+        # Encoder-Decoder-Attention
+        t5x_enc_dec_attention_key = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["key"]["kernel"]
+        t5x_enc_dec_attention_out = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["out"]["kernel"]
+        t5x_enc_dec_attention_query = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["query"]["kernel"]
+        t5x_enc_dec_attention_value = t5x_model["target"]["decoder"][layer_name]["encoder_decoder_attention"]["value"]["kernel"]
+
+        ## Layer Normalization
+        t5x_cross_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_cross_attention_layer_norm"]["scale"]
+
+        # MLP
+        if split_mlp_wi:
+            t5x_mlp_wi_0 = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi_0"]["kernel"]
+            t5x_mlp_wi_1 = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi_1"]["kernel"]
+        else:
+            t5x_mlp_wi = t5x_model["target"]["decoder"][layer_name]["mlp"]["wi"]["kernel"]
+
+        t5x_mlp_wo = t5x_model["target"]["decoder"][layer_name]["mlp"]["wo"]["kernel"]
+
+        ## Layer Normalization
+        tx5_mlp_layer_norm = t5x_model["target"]["decoder"][layer_name]["pre_mlp_layer_norm"]["scale"]
+        
+        # Assigning
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["k"]["kernel"] = t5x_attention_key
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["o"]["kernel"] = t5x_attention_out
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["q"]["kernel"] = t5x_attention_query
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["SelfAttention"]["v"]["kernel"] = t5x_attention_value
+
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["0"]["layer_norm"]["weight"] = t5x_pre_attention_layer_norm
+        
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["k"]["kernel"] = t5x_enc_dec_attention_key
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["o"]["kernel"] = t5x_enc_dec_attention_out
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["q"]["kernel"] = t5x_enc_dec_attention_query
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["EncDecAttention"]["v"]["kernel"] = t5x_enc_dec_attention_value
+        
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["1"]["layer_norm"]["weight"] = t5x_cross_layer_norm
+        
+        if split_mlp_wi:
+            flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_0"]["kernel"] = t5x_mlp_wi_0
+            flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi_1"]["kernel"] = t5x_mlp_wi_1
+        else:
+            flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wi"]["kernel"] = t5x_mlp_wi
+
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["DenseReluDense"]["wo"]["kernel"] = t5x_mlp_wo
+        
+        flax_model.params["decoder"]["block"][str(layer_index)]["layer"]["2"]["layer_norm"]["weight"] = tx5_mlp_layer_norm
+
+    # Decoder Normalization
+    tx5_decoder_norm = t5x_model["target"]["decoder"]["decoder_norm"]["scale"]
+    flax_model.params["decoder"]["final_layer_norm"]["weight"] = tx5_decoder_norm
+
+    # Only for layer 0:
+    t5x_decoder_rel_embedding = t5x_model["target"]["decoder"]["relpos_bias"]["rel_embedding"].T
+    flax_model.params["decoder"]["block"]["0"]["layer"]["0"]["SelfAttention"]["relative_attention_bias"]["embedding"] = t5x_decoder_rel_embedding
+
+    # Token Embeddings
+    tx5_token_embeddings = t5x_model["target"]["token_embedder"]["embedding"]
+    flax_model.params["shared"]["embedding"] = tx5_token_embeddings
+
+    # LM Head
+    flax_model.params["lm_head"]["kernel"] = t5x_model["target"]["decoder"]["logits_dense"]["kernel"]
+
+    flax_model.save_pretrained(flax_dump_folder_path)
+    print("T5X Model was sucessfully converted!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Required parameters
+    parser.add_argument(
+        "--t5x_checkpoint_path", default=None, type=str, required=True, help="Path the TX5 checkpoint."
+    )
+    parser.add_argument(
+        "--config_name", default=None, type=str, required=True, help="Config name of T5 model."
+    )
+    parser.add_argument(
+        "--flax_dump_folder_path", default=None, type=str, required=True, help="Path to the output FLAX model."
+    )
+    args = parser.parse_args()
+    convert_t5x_checkpoint_to_flax(args.t5x_checkpoint_path, args.config_name, args.flax_dump_folder_path)
+    

small_nl16.gin

@@ -0,0 +1,23 @@
+# T5.1.1 Efficient small nl16 model.
+
+import seqio
+include 't5x/examples/scalable_t5/t5_1_1/base.gin' # imports vocab, optimizer and model.
+
+# ------------------- Network specification overrides --------------------------
+network.Transformer.config = @network.T5Config()
+network.T5Config:
+  emb_dim = 512
+  num_heads = 8
+  num_encoder_layers = 16
+  num_decoder_layers = 16
+  head_dim = 64
+  mlp_dim = 2048
+
+# ------------------- Model specification overrides --------------------------
+VOCABULARY = @seqio.SentencePieceVocabulary()
+seqio.SentencePieceVocabulary.sentencepiece_model_file = "spiece.model"
+
+MODEL = @models.EncoderDecoderModel()
+models.EncoderDecoderModel:
+  input_vocabulary = %VOCABULARY
+  output_vocabulary = %VOCABULARY

small_nl16_pretrain.gin

@@ -0,0 +1,26 @@
+# Register necessary SeqIO Tasks/Mixtures.
+from __gin__ import dynamic_registration
+from t5x import utils
+from t5x import partitioning
+import tasks
+import __main__ as train_script
+
+include 'small_nl16.gin'
+include 't5x/configs/runs/pretrain.gin'
+
+
+# ------------------- Training specification overrides --------------------------
+train_script.train:
+  eval_period = 10000
+  use_gda = False
+
+utils.SaveCheckpointConfig:
+  period = 10000
+  keep = 10
+
+MIXTURE_OR_TASK_NAME = "pretrain_finnish"
+USE_CACHED_TASKS = False
+TASK_FEATURE_LENGTHS = {"inputs": 512, "targets": 512}
+TRAIN_STEPS = 500000
+DROPOUT_RATE = 0.0
+BATCH_SIZE = 256

spiece.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55a3645122435e9773fac81fa3f94c1e14149e80311636dfa9245fba3e57a826
+size 824186

start_train.sh

@@ -0,0 +1,12 @@
+# set train hyperparams
+unset LD_PRELOAD
+
+PROJECT_DIR="/researchdisk/t5-small-nl16-finnish"
+T5X_DIR=${HOME}"/t5x"  # directory where the t5x is cloned.
+MODEL_DIR="/researchdisk/t5-small-nl16-finnish"
+export PYTHONPATH=${PROJECT_DIR}
+
+python3 ${T5X_DIR}/t5x/train.py \
+  --gin_search_paths=${PROJECT_DIR} \
+  --gin_file="small_nl16_pretrain.gin" \
+  --gin.MODEL_DIR=\"${MODEL_DIR}\"

tasks.py

@@ -0,0 +1,82 @@
+# adapted from https://huggingface.co/pere/pk-nb-t5x/blob/main/tasks.py
+
+import functools
+
+import seqio
+import tensorflow as tf
+import t5.data
+from datasets import load_dataset, load_from_disk
+from t5.data import postprocessors
+from t5.data import preprocessors
+from t5.evaluation import metrics
+from seqio import FunctionDataSource, utils
+
+TaskRegistry = seqio.TaskRegistry
+
+vocabulary = seqio.SentencePieceVocabulary('spiece.model', extra_ids=0)
+
+DEFAULT_OUTPUT_FEATURES = {
+    "inputs": seqio.Feature(
+        vocabulary=vocabulary, add_eos=True,
+        required=False),
+    "targets": seqio.Feature(
+        vocabulary=vocabulary, add_eos=True)
+}
+
+
+def gen_dataset(split, shuffle=False, seed=None, column="text", dataset=None):
+    if shuffle:
+        if seed:
+            dataset = dataset.shuffle(seed=seed)
+        else:
+            dataset = dataset.shuffle()
+    while True:
+        for item in dataset[str(split)]:
+            yield item[column]
+
+
+def dataset_fn(split, shuffle_files, seed=None, dataset=None):
+    return tf.data.Dataset.from_generator(
+        functools.partial(gen_dataset, split, shuffle_files, seed, dataset=dataset),
+        output_signature=tf.TensorSpec(shape=(), dtype=tf.string, name=dataset_name)
+    )
+
+
+@utils.map_over_dataset
+def target_to_key(x, key_map, target_key):
+    """Assign the value from the dataset to target_key in key_map"""
+    return {**key_map, target_key: x}
+
+
+# Final pretraining task used in Raffel et al., 2019 adaptated to our use
+dataset_name = "/researchdisk/lm_training_dataset_full"
+dataset_params = {"from_disk_path": dataset_name}
+
+if "from_disk_path" in dataset_params:
+    dataset = load_from_disk(dataset_params.get("from_disk_path"))
+else:
+    dataset = load_dataset(**dataset_params)
+
+dataset_shapes = {"train": dataset["train"].num_rows, "validation": dataset["validation"].num_rows}
+TaskRegistry.add(
+    "pretrain_finnish",
+    source=seqio.FunctionDataSource(
+        dataset_fn=functools.partial(dataset_fn, dataset=dataset),
+        splits=("train", "validation"),
+        caching_permitted=False,
+        num_input_examples=dataset_shapes,
+    ),
+    preprocessors=[
+        functools.partial(
+            target_to_key, key_map={
+                "inputs": None,
+                "targets": None,
+            }, target_key="targets"),
+        seqio.preprocessors.tokenize,
+        # seqio.CacheDatasetPlaceholder(),
+        preprocessors.span_corruption, 
+        seqio.preprocessors.append_eos_after_trim,
+    ],
+    output_features={"targets": DEFAULT_OUTPUT_FEATURES["targets"]},
+    metric_fns=[metrics.accuracy]
+)

train_sentencepiece.py

@@ -0,0 +1,6 @@
+import sentencepiece as spm
+
+spm.SentencePieceTrainer.train(input="/researchdisk/lm_training_dataset_full_sentences/train.txt", model_prefix='spiece', vocab_size=32000, character_coverage=1.0,
+                                pad_id=0, unk_id=2, eos_id=1, bos_id=-1,
+                                train_extremely_large_corpus=True,
+                                num_threads=96, input_sentence_size=50000000, shuffle_input_sentence=True)