chroma database 로컬 저장 및 불러오기 ( csv 파일 로드 )

안녕하세요

 

오늘은 개인적으로 간단하게 테스트했던 코드를 공유합니다.

 

요즘에 핫한 LLM (ChatGPT, Gemini) 를 활용한 RAG 어플리케이션 개발시 중요한 부분중에 하나인 Vector database 샘플 코드 입니다.

 

간단히 Chroma 에 저장하고 이를 다시 로드하는 코드 입니다.

 

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain_community.document_loaders import CSVLoader

def init_db():
    embedding_function = OpenAIEmbeddings()
    if os.path.exists("chroma.db"):
        db = Chroma(persist_directory="./chroma.db", embedding_function=embedding_function)
    else:
        loader = CSVLoader("csv_filename.csv", encoding="UTF-8")
        documents = loader.load()
        db = Chroma.from_documents(documents, embedding_function, persist_directory="./chroma.db")
    return db

 

다음은 코드 설명입니다.

1. 현재 저장소 경로에 chroma.db 라는 경로가 있는지 확인합니다.

2. 저장소 경로에 있다면 데이터를 로드합니다. 이때 embedding 은 Openai를 사용했습니다.
   개발의 선택에 따라 다른 오픈소스도 가능합니다.

3. 저장소 경로에 chroma.db 가 없다면 csv 파일을 읽어서 Chroma Database를 생성합니다.

    생성된 데이터베이스는 로컬에 ./chroma.db 라는 이름으로 저장합니다.

 

참고로, csv 파일은 csvLoader를 이용하여 row 별로 데이터를 읽어서 vector database에 저장하는 구조를 사용했습니다.

일반적인 pdf 파일 또는 text 파일에는 적합하지 않습니다. 

 

구현하고자 하는 어플리케이션의 유형 및 데이터의 유형에 따라 파일을 읽어들이는 구조의 변경은 필요합니다.