ReactFileStructure/부트캠프 과정

벡터 검색 엔진의 설정 - 질문에 대한 유사 문장의 탐색을 위해

명징직조지훈 2024. 10. 25. 11:07

벡터 검색 엔진은 문서 임베딩을 저장하고, 

입력된 질문에 대한 관련된 문서 임베딩을 빠르게 찾을 수 있도록 도와준다.

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer

# 임베딩 모델 로드
model = SentenceTransformer('all-MiniLM-L6-v2')

# 예시 문서 리스트 (임베딩할 문서들)
list_of_documents = [
    "This is the first document about AI.",
    "This document is related to machine learning.",
    "Here is a document about deep learning."
]

# 문서 임베딩 생성
document_embeddings = model.encode(list_of_documents)

# 임베딩의 차원 (문서 임베딩의 크기)
embedding_dim = document_embeddings.shape[1]

# FAISS 인덱스 생성 (L2 거리 기반)
index = faiss.IndexFlatL2(embedding_dim)

# 문서 임베딩 추가
index.add(np.array(document_embeddings))

# 추가된 벡터 개수 확인
print(f"Number of embeddings in the index: {index.ntotal}")

# 사용자의 질문
user_question = "What is the difference between AI and machine learning?"

# 질문을 임베딩
question_embedding = model.encode([user_question])

# FAISS에서 유사한 문서 검색 (상위 2개의 유사 문서 찾기)
k = 2  # 찾고 싶은 유사 문서 개수
distances, indices = index.search(np.array(question_embedding), k)

# 검색 결과 출력
print(f"Top {k} similar documents:")
for i, idx in enumerate(indices[0]):
    print(f"Document {idx + 1}: {list_of_documents[idx]} (distance: {distances[0][i]})")
    
    
>>>
Number of embeddings in the index: 3
Top 2 similar documents:
Document 1: This is the first document about AI. (distance: 1.069716215133667)
Document 2: This document is related to machine learning. (distance: 1.1714024543762207)

먼저 문서 임베딩의 준비,

FAISS 인덱스의 생성, (L2 거리 기반 IndexFlatL2 사용)

질문에 대한 유사 문서 검색

결과 출력,

 

만약 여러 문서에 대한 정보를 입력하고, 질문에 해당하는 문서를 탐색할 경우

문서 요약 모델을 돌린 후, 

생성된 요약 정보의 임베딩에 대해 매칭하는 방식으로 사용해볼 수 있겠다.