[NLP, embedding] FAISS 사용하기

FAISS는 Chroma_db와 함께 로컬에서 사용하기 벡터데이터 저장 및 검색 도구

chroma_db는 지난번에 다뤘기에 오늘은 FAISS를 사용해보고자 한다. FAISS는 메타에서 만든것

 

벡터데이터를 생성하고, 저장하고, 검색하는 순서를 정리해보면

기존 자료 임베딩 > 벡터db에 저장 > 벡터db에서 가장 유사한것 검색하여 찾기

이런 순서로 진행된다

 

#1. 텍스트 데이터를 적절한 모델을 사용하여 임베딩(예: 안녕하세요 > [1,8,10,100,.....])

내가 가지고 있는 자료를 임베딩하는 과정이고, 여러번 하면 임베딩모델 사용료가 나오니, 전략을 잘 구상해서, 한방에 잘 임베딩 하고, 그 임베딩값을 잘 저장해두자.

 

용량이 많이 않아서, 책 5권을 적절하게 분할해서, openai의 새로나온 임베딩 모델로 텍스트 임베딩 진행.

embedding.csv파일에 저장하였다

# openai embedding함수
def get_embedding(text, model="text-embedding-3-small"):
    text = text.replace("\n", " ")
    return client.embeddings.create(input=[text], model=model).data[0].embedding

folder_path = './database'
embedding_file = 'embedding.csv'
embedding_file_path = os.path.join(folder_path, embedding_file)

# embedding.csv가 존재하면, 데이터프레임 df로 로드
if os.path.exists(embedding_file_path):
    print(f"{embedding_file} is exist")
    df = pd.read_csv(embedding_file_path)
    # string으로 저장된 embedding을 list로 변환
    df['embedding'] = df['embedding'].apply(ast.literal_eval)
# pdf파일을 읽어서 embedding하여 csv파일로 저장
else:
    dataset_path = './bookdata'
    pdf_files = [file for file in os.listdir(dataset_path) if file.endswith('.pdf')]

    data = []
    # 모든 PDF 파일을 순회하며 embedding
    for file in pdf_files:
        pdf_file_path = os.path.join(dataset_path, file)
        reader = PdfReader(pdf_file_path)

        # 각 pdf의 페이지별로 embedding
        for i, page in enumerate(reader.pages):
            text = page.extract_text() if page.extract_text() else '' #있을 경우 추출, 없으면 빈문자열
            data.append([file, i, text]) # 파일명, 페이지번호, 텍스트

    # 데이터프레임으로 생성
    df = pd.DataFrame(data, columns=['filename', 'page', 'text'])

    # embedding
    df['embedding'] = df['text'].apply(lambda x: get_embedding(x, model="text-embedding-3-small"))

    # csv파일로 저장
    df.to_csv(embedding_file_path, index=False, encoding='utf-8')

나는 PdfReader를 사용했는데, 다른 라이브러리를 사용하는 것도 하나의 전략이다.

또한 split도 다양한 전략이 있으니, 성능좋은 것을 잘 찾아보자!

위 코드를 통해서, 책 5권을 임베딩하여, csv파일에 저장하였습니다.

 

#2. 임베딩된 N차원의 데이터를 vector db에 저장(chroma, faiss)

임베딩된 텍스트데이터를 벡터디비에 저장하는 직관적인 묘사는 아래와 같습니다.

임베딩된 문단(문장)들이 N차원의 데이터가 되고, 이는 N차원의 공간에 하나의 점으로 표현이 가능합니다.

당연히 관련있는 문장들은 비슷한 위치에 분포하게 됩니다.

[1,8,10,100,...] >

점 하나하나가 텍스트가 임베딩 된것

 

#1에서 임베딩한 데이터들을 벡터db에 저장하는 코드입니다.

def create_and_save_faiss_index(embeddings, index_path):
    # 첫 번째 임베딩을 기반으로 인덱스의 차원을 설정
    d = len(embeddings[0])
    index = faiss.IndexFlatL2(d)

    # 임베딩을 순회하면서 인덱스에 추가
    for embedding in embeddings:
        # 각 임베딩을 넘파이 배열로 변환하고 차원을 (1, d)로 재조정
        np_embedding = np.array(embedding, dtype='float32').reshape(1, -1)
        index.add(np_embedding)

    # 인덱스를 파일로 저장
    faiss.write_index(index, index_path)
    print(f"FAISS 인덱스가 {index_path}에 성공적으로 저장되었습니다.")


# 디렉토리 설정
CUR_DIR = os.path.dirname(os.path.abspath(__file__))
FAISS_INDEX_DIR = os.path.join(CUR_DIR, 'faiss_index')
FAISS_INDEX_NAME = "ai-kwon-book.index"

# FAISS 인덱스를 저장할 디렉토리가 없다면 생성
if not os.path.exists(FAISS_INDEX_DIR):
    os.makedirs(FAISS_INDEX_DIR)

# 인덱스 파일 경로
index_path = os.path.join(FAISS_INDEX_DIR, FAISS_INDEX_NAME)

# embedding.csv 파일 경로
folder_path = './database'
embedding_file = 'embedding.csv'
embedding_file_path = os.path.join(folder_path, embedding_file)
# embedding.csv가 존재하면, 데이터프레임 df로 로드
if os.path.exists(embedding_file_path):
    print(f"{embedding_file} is exist")
    df = pd.read_csv(embedding_file_path)
    # string으로 저장된 embedding을 list로 변환
    df['embedding'] = df['embedding'].apply(ast.literal_eval)

    # embedding을 numpy array로 변환
    embeddings = df['embedding'].to_list()
    # faiss 인덱스 생성
    create_and_save_faiss_index(embeddings, index_path)

else:
    print(f"{embedding_file} is not exist")

- 임베딩데이터 1개가 들어오면, 해당 데이터를 faiss에 저장합니다.

- 데이터의 차원을 확인(임베딩데이터중 첫번째 데이터로 차원확인 = len(embedding[0])하여 index로 만들고,

- 임베딩데이터를 np.array로 2차원 행렬로 만들어 저장합니다(faiss가 그렇게 해야함)

- .index파일을 만들어 faiss디비를 생성합니다.

- embedding.csv파일에 들어있는 임베딩 데이터를 한개씩 함수에 넣어서 faiss데이터를 저장합니다.

 

#3. 사용자에게 문장을 입력받아서, 임베딩하고, 벡터db에서 입력된 문장과 가장 가까운 점 K개 가져오기

 이제 벡터db에 저장 되었으니, 사용자의 입력을 받아, 입력과 유사한 문장을 가져오면 된다

# 쿼리에 대한 유사 문서 검색 함수
def query_faiss(query: str, use_retriever: bool = False):
    # 디렉토리 설정
    CUR_DIR = os.path.dirname(os.path.abspath(__file__))
    FAISS_INDEX_DIR = os.path.join(CUR_DIR, 'database/faiss_index')
    FAISS_INDEX_NAME = "ai-kwon-book.index"
    # 인덱스 파일 경로
    index_path = os.path.join(FAISS_INDEX_DIR, FAISS_INDEX_NAME)
    print(index_path)

    # 인덱스 파일에서 FAISS 인덱스 로드
    index = faiss.read_index(index_path)

    # 입력된 문장 임베딩
    query_embedding = get_embedding(
        query,
        model="text-embedding-3-small"
    )

    np_query_embedding = np.array(query_embedding, dtype='float32').reshape(1, -1)

    # 쿼리 벡터에 대한 가장 유사한 k개의 항목 검색
    D, I = index.search(np_query_embedding, 3)  # D는 거리, I는 인덱스

    # csv파일을 읽어서 임베딩값과 가장 가까운 3개 문장을 반환
    df = pd.read_csv('./database/embedding.csv')
    top_docs = df.iloc[I[0]]
    from pprint import pprint
    pprint(top_docs)
    top_docs = top_docs['text'].to_list()

    return top_docs

- query는 사용자의 입력 문장

- 해당 문장을 get_embedding함수로 동일하게 임베딩 진행 > [1,8,100,192,....]

- 해당 임베딩을 행렬화 시키고,

- faiss에 검색 실행 D, I = index.search(np_query_embedding, 3)  # D는 거리, I는 인덱스

- 가장 가까운 3개점(임베딩데이터)을 가지고 원본 데이터에서 검색

- 원본 문장을 받아서, top_docs로 리턴

 

 

이렇게 벡터db를 통한, 맥락검색이 가능하게 된다