HuggingFace GPT2 tokenizer special token 추가하기 ( feat. adding HuggingFace Bert tokenizer special token )

 

 

 

안녕하십니까 다제입니다. 

 

오늘은 Hugginface Tokenizer에 special token을 추가하는 것에 대해서 한번 배워보고자 합니다.

 

언어모델은 크게 Encoder 계열, Decoder 계열, Encoder-Decoder 계열로 구분할 수 있습니다. 

그러나, Tokenizer의 기준으로는 Encoder 계열, Decoder 계열로 나눌 수 있으며, 필요에 따라서는 special token을 추가하여 학습을 진행할 수도 있습니다. 

 

그런데, Encoder 계열, Decoder 계열의 Tokenizer에 따라 special token을 추가하는 방법이 약간 상이 합니다. 

아래와 같이 추가하는 방법을 정리해 놓았으니, 참고하시어 연구하시는데 조금이나 시간을 절약하시기 바랍니다. 

 

 

1. Decoder(GPT2) 계열

gpt2에 special token을 추가하는 방법 
def load_tokenizer(args):
    tokenizer = GPT2Tokenizer.from_pretrained(args['model_name'])
    special_tokens = ['<speaker1>', '<speaker2>']
    tokenizer.add_special_tokens({
        'bos_token': '<bos>',
        'additional_special_tokens': special_tokens
    })

    # add new token ids to args
    special_tokens += ['<bos>', '<eos>']
    sp1_id, sp2_id, bos_id, eos_id = tokenizer.encode(special_tokens)
    args['sp1_id'] = sp1_id
    args['sp2_id'] = sp2_id
    args['bos_id'] = bos_id
    args['eos_id'] = eos_id

    return tokenizer


def load_model(args, tokenizer, device):
    model = GPT2LMHeadModel.from_pretrained(args['model_name']).to(device)
    model.resize_token_embeddings(len(tokenizer))
    return model
    
   tokenizer = load_tokenizer(args)
   model = load_model(args, tokenizer, device)

args에 내가 원하는 모델의 이름을 넣어주기만 하면 바로 사용이 가능합니다. 

 

2. Eecoder(Bert) 계열

Bert에 special Token을 추가하는 방법 
def load_tokenizer(args):
    tokenizer = BertTokenizerFast.from_pretrained(args['model_name'])
    special_tokens = ['<speaker1>', '<speaker2>', '<bos>', '<eos>']
    tokenizer.add_special_tokens({
        'bos_token': '<bos>',
        'eos_token': '<eos>',
        'additional_special_tokens': special_tokens
    })

    # add new token ids to args
    sp1_id, sp2_id, bos_id, eos_id = tokenizer.convert_tokens_to_ids(special_tokens)
    args['sp1_id'] = sp1_id
    args['sp2_id'] = sp2_id
    args['bos_id'] = bos_id
    args['eos_id'] = eos_id

    return tokenizer


def load_model(args, tokenizer, device):
    model = GPT2LMHeadModel.from_pretrained(args['model_name']).to(device)
    model.resize_token_embeddings(len(tokenizer))
    return model

 

참고 부탁드립니다. 

감사합니다.