Upload create_miniature_model.py

create_miniature_model.py

@@ -0,0 +1,46 @@
+import json
+
+import tokenizers
+import torch
+import transformers
+
+
+def shrink_vocab(tokenizer, new_vocab_size):
+    tokenizer_json = json.loads(tokenizer._tokenizer.to_str())
+    vocab = tokenizer_json["model"]["vocab"]
+    if tokenizer_json["model"]["type"] == "BPE":
+        new_vocab = { token: i for token, i in vocab.items() if i < new_vocab_size }
+        merges = tokenizer_json["model"]["merges"]
+        new_merges = []
+        for i in range(len(merges)):
+            if len( merges[i].split()) == 2:
+                a, b = merges[i].split()
+            else:
+                print('skip')
+                continue
+            new_token = "".join((a, b))
+            if a in new_vocab and b in new_vocab and new_token in new_vocab:
+                new_merges.append(merges[i])
+        tokenizer_json["model"]["merges"] = new_merges
+    elif tokenizer_json["model"]["type"] == "Unigram":
+        new_vocab = vocab[:new_vocab_size]
+    elif tokenizer_json["model"]["type"] == "WordPiece" or tokenizer_json["model"]["type"] == "WordLevel":
+        new_vocab = { token: i for token, i in vocab.items() if i < new_vocab_size }
+    else:
+        raise ValueError(f"don't know how to handle {tokenizer_json['model']['type']}")
+    tokenizer_json["model"]["vocab"] = new_vocab
+    tokenizer._tokenizer = tokenizers.Tokenizer.from_str(json.dumps(tokenizer_json))
+
+
+def main():
+    tokenizer = transformers.AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
+    shrink_vocab(tokenizer, new_vocab_size=2000)
+    tokenizer.save_pretrained(".")
+
+    config = transformers.AutoConfig.from_pretrained('noamwies/llama-test-gqa-with-better-transformer')
+    model = transformers.AutoModelForCausalLM.from_config(config, torch_dtype=config.torch_dtype)
+    torch.save(model.state_dict(), 'pytorch_model.bin')
+
+
+if __name__ == '__main__':
+    main()