How to create prefix_tree_opt.pkl for new dataset?

[zhiweihu1103]

Hi, fork. Nice work, I need to know how to create prefix_tree_opt.pkl for new dataset? I want to use this thought to another task.

[Senbao-Shi]

Hi, thank you for your question.

We use the initializer function of the Trie class to build the prefix tree. Make sure every title begins with 'bos' and ends with 'eos'.

from trie import Trie

tree = Trie(tittle_ids)
prefix_tree_dict = tree.trie_dict
with open(prefix_tree_file, 'wb') as f:
    pickle.dump(prefix_tree_dict, f)

[zhiweihu1103]

Thanks for your quick reply. If I have a list of new entity names. Would you mind tell me how to get tittle_ids in the Trie(tittle_ids)?

[Senbao-Shi]

Just tokenize all the new entity names and get the input_ids as follows:

tittle_ids = [tokenizer(t)['input_ids'] for t in tittle]

[zhiweihu1103]

Which model I need as tokenizer? If OPT, I need OPT tokenizer, if Llama, I need Llama tokenizer. Right?

[Senbao-Shi]

yep

[zhiweihu1103]

Thx, let me try.

[zhiweihu1103]

Hi, further question. If I have an entity 'Joe Biden', the input to Trie should be 'eos Joe Biden eos', right? I need to add eos and spaces to separate from Joe Biden.

[Senbao-Shi]

Yes, you can use this method, and you can test it with a small amount of data to see if it performs well.

[zhiweihu1103]

Further question is, why for training, you pad right, for dev or test, you pad left?

[zhiweihu1103]

would you mind provide the entity_list that you used to create tree.pkl file, I may need to test whether my code is correct.

[zhiweihu1103]

I may need to give more details on how I generated tree_opt.pkl, and I hope I can get your help.
First, I give some of my data scenarios:

• I use the opt-1.3b as the tokenizer.
• I use the kb_entity.json from WikiDiverse datatset as the entity_names, the kb_entity.json from the MIMIC paper, I uploaded a copy here.

I use the following code to generate the tree_opt.pkl file:

import pickle
import json

from trie import Trie
from transformers import AutoTokenizer

def create_trie_pkl(data_path, tokenizer_path, output_path):
    entity_name_list = []
    with open(data_path, 'r') as file:
        data = json.load(file)
    for single_data in data:
        entity_name_list.append(single_data['entity_name'])

    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_fast=False)
    tittle_ids = [tokenizer(t)['input_ids'] for t in entity_name_list]
    tree = Trie(tittle_ids)
    prefix_tree_dict = tree.trie_dict
    with open(output_path, 'wb') as f:
        pickle.dump(prefix_tree_dict, f)

if __name__ == '__main__':
    data_path = './kb_entity.json'
    tokenizer_path = './opt-1.3b'
    output_path = './tree_opt.pkl'
    create_trie_pkl(data_path, tokenizer_path, output_path)

However, the generated tree_opt.pkl file is very small. I also uploaded it here. The size of your prefix_tree_opt.pkl file is 209M, but the size of the tree_opt.pkl I generated is only 2.9M. I don't understand how your prefix_tree_opt.pkl is generated. In particular, what exactly does the content of the entity_name_list you use look like? I don’t know what happened in the middle? Looking forward for your reply.

[zhiweihu1103]

If it's convenient, can you leave a WeChat ID? Thank you for your help.

[zhiweihu1103]

Sorry, I have solved this problem, I will close this issue.

[Senbao-Shi]

Sorry for the late response. We have provided guidelines on how to build and use a prefix tree for constrained decoding.