smaller vocab

config.json

@@ -23,5 +23,5 @@
   "torch_dtype": "float16",
   "transformers_version": "4.9.0.dev0",
   "type_vocab_size": 2,
-  "vocab_size": 30522
+  "vocab_size": 5120
 }

make-tiny-electra.py

@@ -63,15 +63,30 @@
 import sys
 import os
 
-from transformers import ElectraTokenizer, ElectraTokenizerFast, ElectraConfig, ElectraForMaskedLM
+from transformers import ElectraTokenizerFast, ElectraConfig, ElectraForMaskedLM
 
 mname_orig = "google/electra-small-generator"
 mname_tiny = "tiny-electra"
 
+
 ### Tokenizer
 
-tokenizer_fast_tiny = ElectraTokenizerFast.from_pretrained(mname_orig)
-tokenizer_tiny = ElectraTokenizer.from_pretrained(mname_orig)
+# Shrink the orig vocab to keep things small (just enough to tokenize any word, so letters+symbols)
+# ElectraTokenizerFast is fully defined by a tokenizer.json, which contains the vocab and the ids, so we just need to truncate it wisely
+import subprocess
+tokenizer_fast = ElectraTokenizerFast.from_pretrained(mname_orig)
+vocab_keep_items = 5120
+tmp_dir = f"/tmp/{mname_tiny}"
+tokenizer_fast.save_pretrained(tmp_dir)
+# resize tokenizer.json (vocab.txt will be automatically resized on save_pretrained)
+# perl -pi -e 's|(2999).*|$1}}}|' tokenizer.json # 0-indexed, so vocab_keep_items-1!
+closing_pat = "}}}"
+cmd = (f"perl -pi -e s|({vocab_keep_items-1}).*|$1{closing_pat}| {tmp_dir}/tokenizer.json").split()
+result = subprocess.run(cmd, capture_output=True, text=True)
+# reload with modified tokenizer
+tokenizer_fast_tiny = ElectraTokenizerFast.from_pretrained(tmp_dir)
+# it seems that ElectraTokenizer is not needed and ElectraTokenizerFast does the job
+
 
 ### Config
 
@@ -85,20 +100,17 @@ config_tiny.update(dict(
     max_position_embeddings=512,
     num_attention_heads=2,
     num_hidden_layers=2,
+    vocab_size=vocab_keep_items,
 ))
 print("New config", config_tiny)
 
+
 ### Model
 
 model_tiny = ElectraForMaskedLM(config_tiny)
-print(f"{mname_tiny}: num of params {model_tiny.num_parameters()}")
-model_tiny.resize_token_embeddings(len(tokenizer_tiny))
-
 
-inputs = tokenizer_tiny("The capital of France is [MASK].", return_tensors="pt")
-labels = tokenizer_tiny("The capital of France is Paris.", return_tensors="pt")["input_ids"]
-outputs = model_tiny(**inputs, labels=labels)
-print("Test with normal tokenizer:", len(outputs.logits[0]))
+print(f"{mname_tiny}: num of params {model_tiny.num_parameters()}")
+model_tiny.resize_token_embeddings(len(tokenizer_fast_tiny))
 
 inputs = tokenizer_fast_tiny("The capital of France is [MASK].", return_tensors="pt")
 labels = tokenizer_fast_tiny("The capital of France is Paris.", return_tensors="pt")["input_ids"]
@@ -108,9 +120,10 @@ print("Test with normal tokenizer:", len(outputs.logits[0]))
 # Save
 model_tiny.half() # makes it smaller
 model_tiny.save_pretrained(".")
-tokenizer_tiny.save_pretrained(".")
 tokenizer_fast_tiny.save_pretrained(".")
 
+#print(model_tiny)
+
 readme = "README.md"
 if not os.path.exists(readme):
     with open(readme, "w") as f:

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:16490204d439317f20cd9f9376475bec315b0085d79908523dabcf03e512654f
-size 4163300
+oid sha256:329472d5ca2d08a2af5798faadab3bebe29fdba4c26ba737240f7bb711f48080
+size 861028

tokenizer_config.json

@@ -1 +1 @@
-{"do_lower_case": true, "unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]", "tokenize_chinese_chars": true, "strip_accents": null, "model_max_length": 512, "special_tokens_map_file": null, "name_or_path": "google/electra-small-generator", "tokenizer_class": "ElectraTokenizer"}
+{"do_lower_case": true, "unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]", "tokenize_chinese_chars": true, "strip_accents": null, "model_max_length": 512, "special_tokens_map_file": null, "name_or_path": "/tmp/tiny-electra", "tokenizer_class": "ElectraTokenizer"}