[LLM] - Pinecone을 활용한 벡터 검색과 유사도 비교

LLM 프로젝트에서 벡터 데이터베이스는 데이터 검색과 관리를 더 효율적으로 만들어주었다. 그중 Pinecone은 높은 성능과 간편한 설정으로 많은 주목을 받았다. 이번 글에서는 내가 Pinecone을 설치하고 설정하며 벡터 DB에 데이터를 임베딩하고 활용한 방법을 정리해본다.

Pinecone이란?

• Pinecone은 임베딩된 데이터의 검색과 관리를 위한 벡터 데이터베이스 서비스다.
• 대규모 데이터셋에서 유사성 검색을 빠르고 정확하게 수행할 수 있다.
• API 기반의 쉬운 설정과 사용법을 제공해 프로젝트에 바로 적용할 수 있다.

패키지 설치하기

%pip install pinecone-client
%pip install --upgrade pinecone pinecone-plugin-assistant
%pip install --upgrade pinecone-client langchain-pinecone
%pip install langchain
%pip install pydantic<2.0
%pip install pandas
%pip install openai
%pip install python-dotenv

이 중 langchain은 LLM과의 통합을 위해, pandas는 데이터 처리, openai는 OpenAI API 사용, python-dotenv는 환경 변수 관리를 위해 설치했다.

Pinecone 설정 및 연결

• Pinecone 클라이어트를 초기화하고 설정하는 과정을 진행

from dotenv import load_dotenv
import os
import pinecone

# 환경 변수 로드
load_dotenv()
PINECONE_API = os.getenv('PINECONE_API_KEY')

# Pinecone 설정 및 연결
pinecone_client = pinecone.Pinecone()
indexes = pinecone_client.list_indexes()
print(indexes)
index_name = indexes[0].name
print(index_name)

Pinecone API 키는 보통 Pinecone의 웹 대시보드에서 생성할 수 있다.

나는 .env 파일을 사용해 API 키와 같은 중요한 정보를 관리하여 코드 보안을 유지했다.

OpenAI 및 LangChain 설정

from openai import OpenAI
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
import os

# OpenAI API 설정
openai_api_key = os.getenv('OPENAI_API_KEY')

# OpenAI Embeddings 설정
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small"
    )

llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0.1,
    )

OpenAI API 키는 OpenAI 계정에서 생성했다. LangChain은 다양한 LLM과 통합할 수 있는 라이브러리로, OpenAI의 임베딩 모델과 쉽게 연동할 수 있다.

케글 데이터 파일명, 컬럼 추출

CSV 파일에서 데이터를 로드하고 메타데이터 설명을 생성했다. 데이터 파일은 케글(Kaggle) 같은 데이터셋 제공 플랫폼에서 구했다.

import pandas as pd

# file 경로
file_path = 'C:/Users/thdgh/Documents/langchain-kr/99-Projects/SYU_metadata_test/data/Amazon_prime_dataset.csv'
# 데이터 로드
data = pd.read_csv(file_path)
# 컬럼명 추출
columns = data.columns.tolist()
# 파일명 추출
file_name = os.path.basename(file_path)
print(columns)
print(file_name)

데이터 설명 생성 및 저장

LLM을 사용하여 데이터 설명을 생성하고 이를 벡터 DB에 저장했다.

# 데이터에 대한 설명을 생성하기 위한 프롬프트 작성
prompt = f"'파일 제목 {file_name}'과(와) 컬럼명 {columns}들(을) 적당한 길이의 설명형식의 한 문장인 메타데이터 만들어줘." \
         f"제목과 컬럼명을 종합적으로 생각해야해"

# LLM을 통해 데이터 설명 생성
data_description = llm(prompt)

print("데이터 설명 메타데이터:\n", data_description)
print(data_description.content)

# Pinecone 인덱스 객체 생성
index = pinecone_client.Index(index_name)

# 메타데이터 설명을 임베딩하여 벡터 DB에 저장
metadata_embedding = embeddings.embed_query(data_description.content)

# Pinecone에 메타데이터 업로드
if index.upsert([(file_name, metadata_embedding, {"description": data_description.content})]):
    print("메타데이터가 벡터 DB에 성공적으로 저장되었습니다.")

임베딩은 텍스트 데이터를 벡터 형태로 변환하는 과정이다. 나는 OpenAI의 임베딩 API를 통해 생성된 벡터를 Pinecone 벡터 DB에 저장하여 효율적인 유사도 검색이 가능하게 했다.

유사 데이터 검색 함수

def search_similar_data(user_query, similarity_threshold=0.5, fallback_threshold=0.3):
    user_embedding = embeddings.embed_query(user_query)
    results = index.query(vector=user_embedding, top_k=3, include_metadata=True)
    
    # 기본 임계값으로 필터링
    primary_recommendations = [match for match in results['matches'] if match['score'] >= similarity_threshold]
    
    # 임계값에 해당하는 결과가 없을 경우 낮은 임계값으로 필터링
    if primary_recommendations:
        recommendations = primary_recommendations
    else:
        recommendations = [match for match in results['matches'] if match['score'] >= fallback_threshold]
    
    # 결과 출력
    if recommendations:
        print("Recommendations:")
        for match in recommendations:
            print(f"File Title: {match['id']}, Score: {match['score']}")
            print(f"Metadata Description: {match['metadata']['description']}\n")
    else:
        print("No relevant data found based on your query.")

    return recommendations

이 함수는 내가 입력한 쿼리를 벡터로 변환한 후 Pinecone에서 검색해 높은 유사도 순으로 결과를 반환하도록 했다.

사용자 - 대화형 챗봇

LangChain을 이용하여 사용자와의 대화형 인터페이스를 만들었다.

from langchain.memory import ConversationBufferMemory
from langchain import PromptTemplate, LLMChain

# 메모리 설정
memory = ConversationBufferMemory()

# 프롬프트 템플릿 설정
prompt_template = PromptTemplate(
    input_variables=["query"],
    template="사용자의 요청에 따라 유사한 데이터를 추천해줘. 요청: {query}"
)

# LLM 체인 설정
llm_chain = LLMChain(llm=llm, prompt=prompt_template, memory=memory)

def chatbot_interaction(user_query):
    # 사용자 쿼리에 대한 답변 생성
    response = llm_chain.run(query=user_query)
    
    # 유사 데이터 검색
    search_results = search_similar_data(user_query)
    
    # 응답과 추천 결과 반환
    return response, search_results

챗봇은 내가 사용자의 요청을 처리하고 관련 데이터를 벡터 DB에서 검색하여 결과를 반환할 수 있도록 구현되었다. 이를 통해 LLM 프로젝트에서 데이터 검색 및 추천 기능을 효과적으로 사용할 수 있다.