[AI] jinaai Embedding 모델 샘플코드

모델 로딩 (huggingface)

import torch
from transformers import AutoModel, AutoTokenizer

model_name = "jinaai/jina-embeddings-v3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

Tokenizer 테스트

# Token 출력 (Dict. id값)
embedding = tokenizer("한 남자가 밧줄을 타고 올라가고 있다.")
print(len(embedding[0]))    # 14
# [0, 3103, 64757, 713, 6, 246028, 60354, 413, 156860, 87179, 108435, 3162, 5, 2]

# Truncation
embedding = tokenizer("한 남자가 밧줄을 타고 올라가고 있다.", truncation=True, max_length=10)
print(len(embedding[0]))    # 10
# [0, 3103, 64757, 713, 6, 246028, 60354, 413, 156860, 2]


# Padding
embedding = tokenizer("한 남자가 밧줄을 타고 올라가고 있다.", padding="max_length", max_length=20)
print(len(embedding[0]))    # 20
# [0, 3103, 64757, 713, 6, 246028, 60354, 413, 156860, 87179, 108435, 3162, 5, 2, 1, 1, 1, 1, 1, 1]

Embedding / Vector 출력 및 유사도 확인

inputs = ["한 남자가 밧줄을 타고 올라가고 있다.", "한 남자가 밧줄을 타고 내려온다.", "한 남자가 사다리를 타고 올라간다.", "한 여성이 플루트를 연주하고 있다."]

with torch.no_grad():
    result = model.encode(inputs, task="text-matching", max_length=2048, truncate_dim=512)

# 첫 문장과 다른 문장들 간의 유사도 비교
print(result[0] @ result[1].T)    # 0.8912654
print(result[0] @ result[2].T)    # 0.80115944
print(result[0] @ result[3].T)    # 0.2859928