インターネット上のPDFファイルを読み込んでVector Search

2023/12/17に公開

FAISSによるVector 検索

前提

langchain は下記を利用

langchain                     0.0.308

コードの説明

各種importをまずします。

from pathlib import Path
from langchain.document_loaders import UnstructuredHTMLLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import PromptTemplate
from langchain.embeddings import VertexAIEmbeddings
from langchain.vectorstores import FAISS

chunk_sizeを1,000で分割したスプリッタを作成。
chunk_size1は、Vector Search1に登録する際に重要になってきます。

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000,chunk_overlap=0)

PDFファイルをloaderにてロードします。
langchainの場合、Pdf、XML、HTMLファイルを読み込むためのローダーが用意されています。

loader = UnstructuredURLLoader(urls=['https://storage.googleapis.com/deepmind-media/gemini/gemini_1_report.pdf','https://storage.googleapis.com/deepmind-media/AlphaCode2/AlphaCode2_Tech_Report.pdf'])
document = loader.load()

HTML、XML用のloaderもありますが、ローカルファイルからの読み込みローダーのみしかみつけらえれていません。

# document = UnstructuredHTMLLoader('./text/myfile.html').load()
# document = UnstructuredXMLLoader('./text/myfile.xml').load()

ロードしたドキュメントを最初にRecursiveCharacterTextSplitterにて分割しておきます。

doc_splits = text_splitter.split_documents(document)
print(len(doc_splits))

分割した最初の9ブロックを出力します。

for i in range(9):
  print(f"[{len(doc_splits[i].page_content)}文字]{doc_splits[i].page_content}", "\n")

FIASSにてVertexAIEmbeddingsモデルを利用して、ドキュメントデータベース作成

db = FAISS.from_documents(doc_splits, VertexAIEmbeddings())

検索してみます。
GeminiとAlphaCode2のレポートPDFをサンプルで利用しているので、なにか検索しています。

query = "Please explain AlphaCode 2 feature ."
docs = db.similarity_search(query, 3)

for index, doc in enumerate(docs):
     print(f"No{index+1}. {doc.page_content}")

最後に

ベクトルの効率的な類似性検索とクラスタリングのためのライブラリを使って、検索までおこないましたが、検索まで簡単におこなえることがわかりました。
続編として、VertexAI Searchとの連携をしていきます

https://github.com/facebookresearch/faiss

Discussion