[SentenceTransformer] special token 추가하기

 

공식 문서에서 위 방법을 제시하고 있다.

나의 경우 내 개인모델을 사용중인데 위 방법대로 하면 토크나이징은 잘 되나 model.encode() 코드 부분에서 에러가 났다.

(에러 메시지: torch.nn.modules.module.ModuleAttributeError: 'Transformer' object has no attribute 'encode')

 

그래서 아래 방법을 사용했다.

from sentence_transformers import SentenceTransformer, models

model = SentenceTransformer("MY MODEL")
tokens = ["TOK1", "TOK2"]

word_embedding_model = model._first_module()
word_embedding_model.tokenizer.add_tokens(tokens, special_tokens=True)
word_embedding_model.auto_model.resize_token_embeddings(len(word_embedding_model.tokenizer))
pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
model = SentenceTransformer(modules=[word_embedding_model, pooling_model])

 

 

Tokenize된 결과는 아래 코드로 확인할 수 있다.

 

text = "텍스트를 입력하세요."
model.tokenizer.convert_ids_to_tokens(model.tokenize([text])["input_ids"][0])