📗

学術論文をインデックス化し、AIエージェント向けにメタデータを抽出する

に公開

paper indexing

このブログでは、研究論文のインデックス作成において、さまざまなメタデータを抽出する包括的な例を紹介します。全文のチャンク化や埋め込みだけでなく、インデックス作成や検索のためのセマンティック埋め込みの構築までを解説します。

このチュートリアルが役に立った場合は、ぜひ CocoIndex on GitHub にスターをお願いします

ユースケース

  • 学術論文の検索・取得や、研究ベースのAIエージェント
  • 論文推薦システム
  • 研究知識グラフ
  • 科学文献のセマンティック分析

本記事で達成すること

例として、この PDF を見てみましょう。

本記事で目指すことは以下の通りです:

  1. 論文のメタデータ(ファイル名、タイトル、著者情報、アブストラクト、ページ数など)を抽出します。

  2. タイトルやアブストラクトなどのメタデータに対してベクトル埋め込みを作成し、セマンティック検索を可能にします。これにより、タイトルやアブストラクトに対してテキストクエリをマッチさせるなど、より高度なメタデータ駆動型の検索が実現できます。

  3. 著者ごとに関連するファイル名をすべてインデックス化し、「Jeff Deanの論文をすべて教えて」などの質問に答えられるようにします。

  4. 論文全体のPDF埋め込みを行いたい場合は、こちらの記事もご参照ください。

完全なコードはこちらでご覧いただけます。

コアコンポーネント

  1. PDF前処理

    • pypdfを使ってPDFを読み込み、以下を抽出します:
      • 総ページ数
      • 1ページ目の内容(メタデータが多く含まれるため)
  2. Markdown変換

    • 1ページ目を Marker を使ってMarkdownに変換します。
  3. LLMによるメタデータ抽出

    • 1ページ目のMarkdownをCocoIndexのExtractByLlm関数を使ってGPT-4oに送信します。
    • 抽出されるメタデータ例:
      • title(文字列)
      • authors(名前、メール、所属を含む)
      • abstract(文字列)
  4. セマンティック埋め込み

    • タイトルはSentenceTransformerのall-MiniLM-L6-v2モデルで直接埋め込みます。
    • アブストラクトは意味的な句読点やトークン数でチャンク化し、それぞれ個別に埋め込みます。
  5. リレーショナルデータ収集

    • 著者情報を展開し、author_papersリレーションに格納することで、
      • Xの論文をすべて表示
      • Yと共著した著者を表示
        などのクエリが可能になります。

前提条件

  • PostgreSQLのインストール
    CocoIndexはインクリメンタル処理のために内部的にPostgreSQLを使用します。
  • OpenAI APIキーの設定
    もしくは、Gemini、Ollama、LiteLLMなどのネイティブサポートもあります。ガイドもご参照ください。
    お好きなLLMプロバイダーを選択でき、完全にオンプレミスで動作させることも可能です。

インデックス化フローの定義

このプロジェクトは、より実際的なユースケースに近いメタデータ理解の包括的な例を示します。
CocoIndexを使えば、100行以内のインデックスロジックでこの設計を簡単に実現できることが分かります。
コードはこちら

これから説明する内容を分かりやすくするため、フローダイアグラムを用意しました。

  1. PDF形式の論文リストをインポートします。
  2. 各ファイルごとに以下を実施します:
    • 論文の1ページ目を抽出
    • 1ページ目をMarkdownに変換
    • 1ページ目からメタデータ(タイトル、著者、アブストラクト)を抽出
    • アブストラクトをチャンクに分割し、それぞれ埋め込みを計算
  3. Postgres(PGVector)に以下のテーブルとしてエクスポートします:
    • 各論文のメタデータ(タイトル、著者、アブストラクト)
    • 著者と論文のマッピング(著者ベースのクエリ用)
    • タイトルやアブストラクトチャンクの埋め込み(セマンティック検索用)

それでは、各ステップを詳しく見ていきましょう。

論文のインポート

@cocoindex.flow_def(name="PaperMetadata")
def paper_metadata_flow(
    flow_builder: cocoindex.FlowBuilder, data_scope: cocoindex.DataScope
) -> None:
    data_scope["documents"] = flow_builder.add_source(
        cocoindex.sources.LocalFile(path="papers", binary=True),
        refresh_interval=datetime.timedelta(seconds=10),
    )

flow_builder.add_source は、サブフィールド(filenamecontent)を持つテーブルを作成します。詳細はドキュメントをご覧ください。

メタデータの抽出と収集

基本情報の抽出

PDFの1ページ目とページ数を抽出するカスタム関数を定義します。

@dataclasses.dataclass
class PaperBasicInfo:
    num_pages: int
    first_page: bytes
@cocoindex.op.function()
def extract_basic_info(content: bytes) -> PaperBasicInfo:
    """Extract the first pages of a PDF."""
    reader = PdfReader(io.BytesIO(content))

    output = io.BytesIO()
    writer = PdfWriter()
    writer.add_page(reader.pages[0])
    writer.write(output)

    return PaperBasicInfo(num_pages=len(reader.pages), first_page=output.getvalue())

これをフローにプラグインします。
PDF全体が非常に大きいため、メタデータの処理コストを最小化するために、1ページ目からメタデータを抽出します。

with data_scope["documents"].row() as doc:
    doc["basic_info"] = doc["content"].transform(extract_basic_info)

このステップ後、各論文の基本情報が取得できます。

基本情報のパース

1ページ目をMarkerを使ってMarkdownに変換します。

Markerのコンバーター関数を定義し、キャッシュします。
初期化にはリソースが必要なため、キャッシュを使用します。
また、初期化にはリソースが必要なため、キャッシュを使用します。

@cache
def get_marker_converter() -> PdfConverter:
    config_parser = ConfigParser({})
    return PdfConverter(
        create_model_dict(), config=config_parser.generate_config_dict()
    )

カスタム関数にプラグインします。

@cocoindex.op.function(gpu=True, cache=True, behavior_version=1)
def pdf_to_markdown(content: bytes) -> str:
    """Convert to Markdown."""

    with tempfile.NamedTemporaryFile(delete=True, suffix=".pdf") as temp_file:
        temp_file.write(content)
        temp_file.flush()
        text, _, _ = text_from_rendered(get_marker_converter()(temp_file.name))
        return text

換関数にプラグインします。

with data_scope["documents"].row() as doc:      
    doc["first_page_md"] = doc["basic_info"]["first_page"].transform(
            pdf_to_markdown
        )

このステップ後、各論文の1ページ目がMarkdown形式で取得できます。

LLMによる基本情報の抽出

LLMによるメタデータ抽出のためのスキーマを定義します。
CocoIndexは、複雑でネストされたスキーマを持つLLM構造化抽出をネイティブにサポートしています。
ネストされたスキーマについて詳しく学びたい場合は、こちらの記事をご覧ください。

@dataclasses.dataclass
class PaperMetadata:
    """
    Metadata for a paper.
    """

    title: str
    authors: list[Author]
    abstract: str

ExtractByLlm関数にプラグインします。
データクラスを定義すると、CocoIndexは自動的にLLMのレスポンスをデータクラスにパースします。

doc["metadata"] = doc["first_page_md"].transform(
    cocoindex.functions.ExtractByLlm(
        llm_spec=cocoindex.LlmSpec(
            api_type=cocoindex.LlmApiType.OPENAI, model="gpt-4o"
        ),
        output_type=PaperMetadata,
        instruction="Please extract the metadata from the first page of the paper.",
    )
)

このステップ後、各論文のメタデータが取得できます。

論文メタデータの収集

  paper_metadata = data_scope.add_collector()
  with data_scope["documents"].row() as doc:
    # ... process
    # Collect metadata
    paper_metadata.collect(
        filename=doc["filename"],
        title=doc["metadata"]["title"],
        authors=doc["metadata"]["authors"],
        abstract=doc["metadata"]["abstract"],
        num_pages=doc["basic_info"]["num_pages"],
    )

必要なメタデータを収集します。

著者と論文のマッピングの収集

著者リストを抽出しました。ここでは、著者と論文のマッピングを別のテーブルに収集して、検索機能を構築します。
著者ごとに収集します。

author_papers = data_scope.add_collector()

with data_scope["documents"].row() as doc:
    with doc["metadata"]["authors"].row() as author:
        author_papers.collect(
            author_name=author["name"],
            filename=doc["filename"],
        )

このステップ後、著者と論文のマッピングが取得できます。
Author Papers

埋め込みの計算と収集

タイトルの埋め込み

doc["title_embedding"] = doc["metadata"]["title"].transform(
    cocoindex.functions.SentenceTransformerEmbed(
        model="sentence-transformers/all-MiniLM-L6-v2"
    )
)

アブストラクトのチャンク化と埋め込み

アブストラクトをチャンクに分割し、それぞれ埋め込みを計算します。
アブストラクトは非常に長い場合があります。

doc["abstract_chunks"] = doc["metadata"]["abstract"].transform(
    cocoindex.functions.SplitRecursively(
        custom_languages=[
            cocoindex.functions.CustomLanguageSpec(
                language_name="abstract",
                separators_regex=[r"[.?!]+\s+", r"[:;]\s+", r",\s+", r"\s+"],
            )
        ]
    ),
    language="abstract",
    chunk_size=500,
    min_chunk_size=200,
    chunk_overlap=150,
)

このステップ後、各論文のアブストラクトチャンクが取得できます。

各チャンクを埋め込み、それぞれの埋め込みを収集します。

with doc["abstract_chunks"].row() as chunk:
    chunk["embedding"] = chunk["text"].transform(
        cocoindex.functions.SentenceTransformerEmbed(
            model="sentence-transformers/all-MiniLM-L6-v2"
        )
    )

このステップ後、各論文のアブストラクトチャンクの埋め込みが取得できます。

埋め込みの収集

metadata_embeddings = data_scope.add_collector()

with data_scope["documents"].row() as doc:
    # ... process
    # collect title embedding
    metadata_embeddings.collect(
        id=cocoindex.GeneratedField.UUID,
        filename=doc["filename"],
        location="title",
        text=doc["metadata"]["title"],
        embedding=doc["title_embedding"],
    )
    with doc["abstract_chunks"].row() as chunk:
        # ... process
        # collect abstract chunks embeddings
        metadata_embeddings.collect(
            id=cocoindex.GeneratedField.UUID,
            filename=doc["filename"],
            location="abstract",
            text=chunk["text"],
            embedding=chunk["embedding"],
        )

データのエクスポート

最後に、データをPostgresにエクスポートします。

paper_metadata.export(
    "paper_metadata",
    cocoindex.targets.Postgres(),
    primary_key_fields=["filename"],
)
author_papers.export(
    "author_papers",
    cocoindex.targets.Postgres(),
    primary_key_fields=["author_name", "filename"],
)    
metadata_embeddings.export(
    "metadata_embeddings",
    cocoindex.targets.Postgres(),
    primary_key_fields=["id"],
    vector_indexes=[
        cocoindex.VectorIndexDef(
            field_name="embedding",
            metric=cocoindex.VectorSimilarityMetric.COSINE_SIMILARITY,
        )
    ],
)

この例では、PGVectorを埋め込みストアとして使用します。
CocoIndexでは、Qdrantなどの他のサポートされたベクトルデータベースに対して1行の切り替えが可能です。詳細はガイドをご覧ください。
私たちはインターフェースを標準化し、それをLEGOのように構築することを目指しています。

CocoInsightでステップバイステップで見る

CocoInsightでプロジェクトをステップバイステップで見ることができます。
各フィールドがどのように構築され、どのように動作しているかを確認できます。

埋め込みに対するクエリの構築

Text Embeddingsのセクションを参照して、埋め込みに対するクエリの構築方法を確認してください。
CocoIndexでは、追加のクエリインターフェースは提供していません。SQLを書くか、ターゲットストレージのクエリエンジンに依存します。

  • 多くのデータベースは、独自のベストプラクティスで最適化されたクエリ実装を持っています
  • クエリ空間には、クエリ、ランキング、その他の検索関連機能の優れたソリューションがあります。

クエリの作成についてサポートが必要な場合は、Discordまでお問い合わせください。

サポート

私たちは常に改善しています。また、機能や例がさらに追加されています。
この記事が役に立った場合は、GitHubで⭐スターをお願いします。

ありがとうございました!

Discussion