File size: 476 Bytes
b599481
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def get_special_tokens_dict(dataset):
    if dataset.startswith("redial"):
        movie_token = "<movie>"
    elif dataset.startswith("opendialkg"):
        movie_token = "<mask>"
    gpt2_special_tokens_dict = {
        "pad_token": "<pad>",
        "additional_special_tokens": [movie_token],
    }

    prompt_special_tokens_dict = {
        "additional_special_tokens": [movie_token],
    }

    return gpt2_special_tokens_dict, prompt_special_tokens_dict