[LLM] HuggingFace SentenceTransformer에서 Tokenizer에 새 단어 추가하고 의미 학습하기

HuggingFace의 SentenceTransformer를 사용할 때 기존 tokenizer에 없는 새로운 단어를 추가하고, 그 단어의 의미를 학습시키는 과정이 필요할 때가 있습니다. 이때 라벨이 포함된 학습 데이터를 사용하면 모델이 새 단어가 포함된 문맥을 더 잘 이해할 수 있게 됩니다. 이 포스팅에서는 새 단어를 추가하고 라벨이 있는 데이터를 활용해 학습하는 방법을 단계별로 설명하겠습니다.

 

Step 1. 기존 Tokenizer 불러오기

먼저, 사용할 SentenceTransformer 모델을 선택하고, 이 모델의 tokenizer를 가져옵니다.

from transformers import AutoTokenizer

model_name = "sentence-transformers/all-mpnet-base-v2"
tokenizer = AutoTokenizer.from_pretrained(model_name)

Step 2. Tokenizer에 새로운 단어 추가하기

새로운 단어를 add_tokens를 사용해 추가합니다. 추가할 단어를 리스트로 받아 아래와 같이 실행합니다.

new_words = ["새로운단어1", "새로운단어2"]
num_added_toks = tokenizer.add_tokens(new_words)
print(f"추가된 토큰 개수: {num_added_toks}")

 

추가된 단어가 실제로 tokenizer에 포함되었는지 convert_tokens_to_ids로 확인합니다.

token_id = tokenizer.convert_tokens_to_ids("새로운단어1")
print(f"새로운단어1의 토큰 ID: {token_id}")

Step 3. SentenceTransformer 모델에 Tokenizer 업데이트하기

새 tokenizer를 적용한 후, 모델의 embedding layer에도 반영해야 합니다. resize_token_embeddings 메서드를 사용해 추가된 토큰 크기에 맞춰 token embedding을 조정합니다.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer(model_name)
model.tokenizer = tokenizer  # 새 tokenizer 적용
model.resize_token_embeddings(len(tokenizer))  # 새 토큰 개수에 맞게 embedding 크기 조정

Step 4. 추가된 단어 학습을 위한 라벨 설정하기

새로운 단어는 모델의 embedding layer에 포함되지만 의미를 학습하지 않은 상태입니다. 추가된 단어의 의미를 학습하기 위해 라벨이 포함된 학습 데이터를 준비합니다. 새 단어가 포함된 문장과 유사한 의미의 문장을 쌍으로 제공하고, 해당 쌍을 1로 라벨링하여 모델이 유사도를 학습할 수 있게 합니다.

from sentence_transformers import InputExample

train_examples = [
    InputExample(texts=["이 문장에 새로운단어1이 포함되어 있습니다.", "유사한 의미를 가진 문장입니다."], label=1),
    InputExample(texts=["새로운단어1과 관련된 문장입니다.", "다른 유사한 문장입니다."], label=1),
    # 새 단어가 포함된 문장과 무관한 문장도 추가해 label 0으로 설정합니다.
    InputExample(texts=["새로운단어1이 들어간 문장입니다.", "무관한 내용의 문장입니다."], label=0),
]

 

Step 5. 추가 학습 (Fine-Tuning)

라벨이 포함된 데이터를 사용해 모델을 학습합니다. 이 과정을 통해 모델이 새 단어의 문맥을 학습하고 의미를 더 잘 이해할 수 있습니다.

from sentence_transformers import losses
from torch.utils.data import DataLoader

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=8)
train_loss = losses.CosineSimilarityLoss(model=model)

# 모델 학습 수행
model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=1, warmup_steps=100)

 

Step 6. 학습 후 결과 확인하기

학습이 완료되면, 새 단어가 포함된 문장을 모델에 넣어 embedding 결과를 확인합니다.

sentence = "이 문장에 새로운단어1을 포함해 embedding을 수행합니다."
embedding = model.encode(sentence)
print(embedding)

 

새 단어가 문맥에 맞는 벡터로 변환되었는지, 학습 전후의 embedding을 비교하여 확인하는 것도 좋습니다.

결론

이 과정을 통해 기존 SentenceTransformer 모델에 새로운 단어를 추가하고, 그 단어의 의미를 학습시킬 수 있습니다. 이를 통해 사용자 정의 단어를 추가하여 더욱 정밀한 embedding 결과를 얻을 수 있습니다.

 

Reference

https://www.sbert.net/docs/sentence_transformer/training_overview.html#

Training Overview — Sentence Transformers documentation