본문 바로가기
AI

how to random masking in Huggingface.

by soypablo 2022. 11. 24.

 

문장에 랜덤마스킹을 하고싶은데, 어떻게 하는지 기억이 안나서 찾은김에 기록하기로 했다.

아주 간단하다.

transformer에서, DataCollatorForLanguageModeling을 임포트한다.
그 뒤 사용할 토크나이저와 mlm, mlm_probability(마스킹 비율), return_tensors(리턴할 텐서 타입)을 넣어주면 끝이다.

# transformer import
from transformers import DataCollatorForLanguageModeling, AutoTokenizer

# 토크나이저 선언
tokenizer = AutoTokenizer.from_pretrained("klue/roberta-large")
# DataCollator선언.
mask_datacollator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=0.15, return_tensors="pt"
)

# 문장 하나 토크나이즈 하기.
sentence = "조선의 2222대 왕은 나다."
inputs = tokenizer(sentence, return_tensors="pt", add_special_tokens=False)
# 문장에 랜덤 마스킹 하기.
masked_input = mask_datacollator([inputs])["input_ids"][0][0]
print(f"원 문장: {sentence}")
print(f"랜덤 마스킹된 문장: {tokenizer.decode(masked_input)}")
출력:
원 문장: 조선의 2222대 왕은 나다.
랜덤 마스킹된 문장: 조선의 2222대 왕 [MASK] 나다.

배치 형태로 전달이 되었을 때 사용할 때랑, 문장 하나만 랜덤으로 마스킹 하는 예시는 조금 다르니 참고하세요.(input을 다르게 해줘야 합니다.)

 

※트레이닝 할때는 Trainer에 parameter로 datacollaor를 넣어주면 됩니다.

예시:

trainer = Trainer(
model=model,
args=training_args,
data_collator=mask_datacollator,
train_dataset=dataset
)