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()