File size: 1,847 Bytes
5a7cdd2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
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()
|