用100行代码将PDF库转换为可搜索的研究数据库

太长不看版

本教程通过完整示例演示如何索引研究论文，提取多种元数据（不仅限于全文分块和嵌入），并构建用于索引和查询的语义嵌入。

如果您觉得本教程有帮助，请在GitHub上为CocoIndex点个⭐星标支持我们。

应用场景

学术搜索与检索，以及基于研究的AI智能体
论文推荐系统
研究知识图谱
科学文献的语义分析

实现目标

以这篇PDF为例，我们想要实现：

提取论文元数据，包括文件名、标题、作者信息、摘要和页数
为标题和摘要等元数据构建向量嵌入，实现语义搜索
建立作者索引和每个作者关联的所有文件名，回答诸如"给我Jeff Dean的所有论文"这类问题

如需对论文进行完整PDF嵌入，可参考这篇文章。

完整代码可在此处找到。

核心组件

PDF预处理

使用pypdf读取PDF并提取：

总页数
首页内容（用作富含元数据信息的代理）

Markdown转换

使用Marker将首页转换为Markdown格式。

LLM驱动的元数据提取

使用CocoIndex的ExtractByLlm函数将首页Markdown发送到GPT-4o。提取的元数据包括：

标题（字符串）
作者（包含姓名、邮箱和所属机构）
摘要（字符串）

语义嵌入

标题直接使用SentenceTransformer的all-MiniLM-L6-v2模型进行嵌入
摘要根据语义标点和词元数量进行分块，然后对每个分块单独嵌入

关系数据收集

将作者展开并收集到author_papers关系中，支持以下查询：

显示X的所有论文
Y与哪些合著者合作过？

前置要求

安装PostgreSQL（CocoIndex内部使用PostgreSQL进行增量处理）
配置OpenAI API密钥
或者，我们原生支持Gemini、Ollama、LiteLLM，请查看指南

您可以选择喜欢的LLM提供商，并完全在本地运行。

定义索引流程

本项目展示了更接近真实使用场景的元数据理解示例。

您将看到通过CocoIndex仅用100行索引逻辑代码就能轻松实现这一设计。

为了更好地引导您了解我们将要演示的内容，以下是流程图：

导入PDF论文列表
对每个文件：
- 提取论文首页
- 将首页转换为Markdown
- 从首页提取元数据（标题、作者、摘要）
- 将摘要分块，并计算每个分块的嵌入
使用PGVector导出到Postgres的以下表：
- 每篇论文的元数据（标题、作者、摘要）
- 作者到论文的映射，用于基于作者的查询
- 标题和摘要分块的嵌入，用于语义搜索

详细步骤

导入论文

1
2
3
4
5
6
7
8


@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将创建包含子字段（文件名、内容）的表，详情请参考文档。

提取和收集元数据

提取首页基本信息

定义自定义函数提取PDF的首页和页数：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16


@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
2


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

完成此步骤后，您应该拥有每篇论文的基本信息。

解析基本信息

我们将使用Marker将首页转换为Markdown。或者，您可以轻松插入喜欢的PDF解析器，如Docling。

定义marker转换器函数并进行缓存，因为其初始化很耗资源。这确保相同的转换器实例对不同输入文件重复使用。

1
2
3
4
5
6


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

插入到自定义函数中：

1
2
3
4
5
6
7
8
9


@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

传递给transform：

1
2
3
4


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

完成此步骤后，您应该拥有每篇论文首页的Markdown格式。

使用LLM提取基本信息

为LLM提取定义模式。CocoIndex原生支持具有复杂嵌套模式的LLM结构化提取。

1
2
3
4
5
6
7
8
9


@dataclasses.dataclass
class PaperMetadata:
    """
    Metadata for a paper.
    """
    
    title: str
    authors: list[Author]
    abstract: str

插入到ExtractByLlm函数中。定义了数据类后，CocoIndex会自动将LLM响应解析到数据类中。

1
2
3
4
5
6
7
8
9


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.",
    )
)

完成此步骤后，您应该拥有每篇论文的元数据。

收集论文元数据

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11


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"],
    )

只需收集您需要的任何内容 :)

收集作者到文件名的信息

我们已经提取了作者列表。这里我们希望在单独的表中收集作者→论文信息以构建查找功能。

简单地按作者收集：

1
2
3
4
5
6
7
8


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"],
        )

计算和收集嵌入

标题

1
2
3
4
5


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

摘要

将摘要分块，嵌入每个分块并收集它们的嵌入。有时摘要可能很长。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


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,
)

完成此步骤后，您应该拥有每篇论文的摘要分块。

嵌入每个分块并收集它们的嵌入：

1
2
3
4
5
6


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

完成此步骤后，您应该拥有每篇论文摘要分块的嵌入。

收集嵌入

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22


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：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23


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，详情请参阅此指南。

我们的目标是标准化接口，使其像搭建乐高一样简单。

在CocoInsight中逐步查看

您可以在CocoInsight中逐步查看项目，了解每个字段是如何构建的以及幕后发生了什么。

查询索引

您可以参考文本嵌入部分的指南来构建针对嵌入的查询。

目前CocoIndex不提供额外的查询接口。我们可以编写SQL或依赖目标存储的查询引擎。

许多数据库已经具有优化的查询实现和最佳实践。查询空间有优秀的解决方案用于查询、重新排名和其他搜索相关功能。

如需查询编写帮助，请随时在Discord上联系我们。

支持我们

我们正在不断改进，更多功能和示例即将推出。

如果本文对您有帮助，请在GitHub上为我们点个⭐星标，帮助我们成长。

感谢阅读！