NV-Embed, NVIDIA가 공개한 임베딩 모델 (영어용, 32k Input -> 4096 dim)

읽을거리&정보공유
, ,
1

NV-Embed, NVIDIA가 공개한 임베딩 모델 (영어/비상업용, 32k Input -> 4096 dim)

소개

읽어야 하는 이유

NV-Embed 모델은 MTEB(Massive-Text Embedding Benchmark) 벤치마크에서 56개의 다양한 작업에서 최고 성적을 기록(2024년 5월 24일 기준)한 새로운 임베딩 모델입니다. NV-Embed 모델은 NVIDIA가 발표한 모델로, 검색, 재정렬, 분류, 클러스터링, 의미적 텍스트 유사성 작업에서 뛰어난 성능을 자랑합니다.

NV-Embed는 Mistral-7B-v0.1을 기반으로 한 디코더 전용 LLM을 사용하여, 더 나은 풀링 임베딩 출력을 생성하도록 합니다. 이 모델은 두 단계의 인스트럭션 튜닝 방법을 사용하여 검색 및 비검색 작업의 정확도를 높입니다. 또한, 최대 32k 토큰 입력을 지원하며, 임베딩 차원은 4096입니다.

주요 기능

  • Latent-Attention 풀링: LLM이 잠재 벡터를 참조하여 풀링 임베딩 출력 생성

  • 대규모 입력 지원: 최대 32k 토큰 입력 가능

  • 고차원 임베딩: 4096 임베딩 차원

사용 방법

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

# 작업명에 따라 인스트럭션을 설정합니다.
task_name_to_instruct = {"example": "Given a question, retrieve passages that answer the question",}

# 쿼리 프리픽스를 설정합니다.
query_prefix = "Instruct: "+task_name_to_instruct["example"]+"\nQuery: "

# 검색할 쿼리를 설정합니다.
queries = [
    'are judo throws allowed in wrestling?', 
    'how to become a radiology technician in michigan?'
]

# 패시지 프리픽스를 설정합니다.
passage_prefix = ""

# 검색할 패시지를 설정합니다.
passages = [
    "Since you're reading this, you are probably someone from a judo background...",
    "Below are the basic steps to becoming a radiologic technologist in Michigan..."
]

# 모델을 불러옵니다.
model = AutoModel.from_pretrained('nvidia/NV-Embed-v1', trust_remote_code=True)

# 최대 입력 길이를 설정합니다.
max_length = 4096

# 쿼리 임베딩을 생성합니다.
query_embeddings = model.encode(queries, instruction=query_prefix, max_length=max_length)

# 패시지 임베딩을 생성합니다.
passage_embeddings = model.encode(passages, instruction=passage_prefix, max_length=max_length)

# 임베딩을 정규화합니다.
query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
passage_embeddings = F.normalize(passage_embeddings, p=2, dim=1)

# 쿼리와 패시지 간의 유사도를 계산합니다.
scores = (query_embeddings @ passage_embeddings.T) * 100
print(scores.tolist())

라이선스

이 프로젝트는 CC BY-NC 4.0 라이선스로 공개 및 배포되어 있어, 상업적 용도로는 사용할 수 없습니다.

NV-Embed 모델 관련 논문

NV-Embed 모델 가중치 (:hugs: Hugging Face)

임베딩(Embedding) 관련 더 읽어보기



이 글은 GPT 모델로 정리한 글을 바탕으로 한 것으로, 원문의 내용 또는 의도와 다르게 정리된 내용이 있을 수 있습니다. 관심있는 내용이시라면 원문도 함께 참고해주세요! 읽으시면서 어색하거나 잘못된 내용을 발견하시면 덧글로 알려주시기를 부탁드립니다. :hugs:

:pytorch:파이토치 한국 사용자 모임:kr:이 정리한 이 글이 유용하셨나요? 회원으로 가입하시면 주요 글들을 이메일:love_letter:로 보내드립니다! (기본은 Weekly지만 Daily로 변경도 가능합니다.)

:gift: 아래:arrow_lower_right:쪽에 좋아요:heart:를 눌러주시면 새로운 소식들을 정리하고 공유하는데 힘이 됩니다~ :star_struck: