【AI】chromaDBを用いた類似検索

pip install chromadb

import chromadb
chroma_client = chromadb.Client()

collection = chroma_client.create_collection(name="my_collection")

collection.add(
    documents=[
        "This is a document about pineapple",
        "This is a document about oranges"
    ],
    ids=["id1", "id2"]
)

results = collection.query(
    query_texts=["This is a query document about hawaii"],
    n_results=2 
)
print(results)

import os
import fitz  # PyMuPDF

def extract_pdf_info(pdf_folder):
    pdf_info_list = []
    
    # 指定されたフォルダ内の全てのPDFファイルをリストアップ
    pdf_files = [f for f in os.listdir(pdf_folder) if f.endswith('.pdf')]
    
    for pdf_file in pdf_files:
        pdf_path = os.path.join(pdf_folder, pdf_file)
        pdf_document = fitz.open(pdf_path)
        
        # 各ページから情報を抽出
        for page_num in range(pdf_document.page_count):
            page = pdf_document.load_page(page_num)
            text = page.get_text()
            text = text.replace('\n', ' ').replace('\r', '')
            pdf_info = {
                'file_name': pdf_file,
                'page_number': page_num + 1,
                'text_content': text
            }
            pdf_info_list.append(pdf_info)
    
    return pdf_info_list

# PDFファイルが含まれているフォルダを指定
pdf_folder = 'pdf'

# PDF情報を抽出
pdf_info_list = extract_pdf_info(pdf_folder)

documents = [info['text_content'] for info in pdf_info_list]
metadatas = [{'file_nam}e': info['file_name'], 'page_number': info['page_number']} for info in pdf_info_list]
ids = [str(i) for i in range(1, len(pdf_info_list) + 1)]

from openai import OpenAI
from tqdm import tqdm

client = OpenAI(
    api_key = 自身のOPENAI_API_KEYを記載する
)
def get_embedding(text):
    response = client.embeddings.create(
        model="text-embedding-3-large",
        input=text
    )
    return response.data[0].embedding

embeddings = []
for doc in tqdm(documents, desc="Getting embeddings"):
    embeddings.append(get_embedding(doc))

import chromadb
import chromadb.utils.embedding_functions as embedding_functions

#chromaDBでembeddingに使用するモデルを指定
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
                api_key=os.getenv("OPENAI_API_KEY"),
                model_name="text-embedding-3-large"
            )

#chromaDBに接続
chroma_client = chromadb.Client()
collection = chroma_client.create_collection(name="my_collection", 
                                             metadata={"hnsw:space": "cosine"}, 
                                             embedding_function=openai_ef)

#chromaDBにデータの追加
collection.add(
    documents=documents,
    embeddings=embeddings,
    metadatas=metadatas,
    ids=ids
)

import time

# 問い合わせクエリーの用意
query_text = "職域における資産形成に関する取組について教えてください"

# 問い合わせクエリーのembedding
query_embedding = get_embedding(query_text)

# 検索実行（上位10件を表示）
start_time = time.time()
results = collection.query(
    query_embeddings=[query_embedding],
    n_results=10
)
end_time = time.time()
execution_time = end_time - start_time

print(f'検索時間:{execution_time}')