SBERT 빠른 시작

SBERT를 처음 쓸 때 가장 빠르게 확인하는 방법은 사전 학습된 모델 하나를 로드해서 세 문장의 임베딩을 만들어 보는 거예요. 설치부터 유사도 계산까지 몇 줄이면 끝나죠.

출처: SBERT Quickstart

먼저 설치부터 할게요. Python 3.10+, PyTorch 1.11.0+, transformers v4.41.0+를 권장해요.

pip install -U sentence-transformers

설치가 끝나면 SentenceTransformer 모델을 로드하고 encode로 임베딩을 만들어요.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

sentences = [
    "The weather is lovely today.",
    "It's so sunny outside!",
    "He drove to the stadium.",
]

embeddings = model.encode(sentences)
print(embeddings.shape)  # [3, 384]

similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[1.0000, 0.6660, 0.1046],
#         [0.6660, 1.0000, 0.1411],
#         [0.1046, 0.1411, 1.0000]])

all-MiniLM-L6-v2는 10억 개가 넘는 학습 쌍으로 파인튜닝된 MiniLM 모델이에요. 모델 하나마다 출력 차원이 달라서, 출력이 [3, 384]처럼 384차원으로 나오는 걸 확인할 수 있어요. 유사도 행렬을 보면 문장 1·2가 서로 0.6660으로 높고, 문장 3과는 낮은 걸 볼 수 있죠. 의미가 비슷한 문장일수록 값이 크다는 걸 그대로 보여주는 부분이에요.

이후에는 검색처럼 더 복잡한 파이프라인을 만들고 싶어질 텐데, 그땐 기본 encode에 더해 query·document 전용 프롬프트를 자동으로 다루는 encode_query()와 encode_document()를 쓰게 돼요.