用于并行处理的SimpleDirectoryReader在LLamaIndex中的应用-CSDN博客

本文链接：https://blog.csdn.net/ppoojjj/article/details/140193812

用于并行处理的SimpleDirectoryReader在LLamaIndex中的应用

在本文中，我们将演示如何在使用SimpleDirectoryReader加载数据时运用并行处理技术。并行处理在处理较重的工作负载时非常有用，例如从包含许多文件的目录中加载数据。（注意：如果使用Windows，使用并行处理加载数据时可能获得的收益较小。这与multiprocessing在Linux/mac和Windows上的工作方式差异有关，参考这里或这里)。

必要库导入

import cProfile
import pstats
from pstats import SortKey
#中转API
from llama_index.core import SimpleDirectoryReader

示例代码

在这个示例中，我们将使用来自llamahub的PatronusAIFinanceBenchDataset数据集。该数据集基于32个PDF文件，这些文件包含在从llamahub下载的内容中。

# 下载数据集
!llamaindex-cli download-llamadataset PatronusAIFinanceBenchDataset --download-dir ./data

# 定义包含32个PDF文件的目录
reader = SimpleDirectoryReader(input_dir="./data/source_files")

# 顺序加载数据
documents = reader.load_data()
len(documents)  # 输出文档的数量

# 使用cProfile进行性能分析
cProfile.run("reader.load_data()", "oldstats")
p = pstats.Stats("oldstats")
p.strip_dirs().sort_stats(SortKey.CUMULATIVE).print_stats(15)

使用并行处理加载数据

要使用并行处理加载数据，我们需要将num_workers设置为一个正整数值。

documents = reader.load_data(num_workers=10)

len(documents)

# 使用cProfile进行性能分析
cProfile.run("reader.load_data(num_workers=10)", "newstats")
p = pstats.Stats("newstats")
p.strip_dirs().sort_stats(SortKey.CUMULATIVE).print_stats(15)

结果分析

通过上述结果分析，可以看到使用并行处理从包含许多文件的目录中加载数据时，速度大约提高了13倍（或1200%的速度增加）。

# 顺序加载统计信息
Wed Jan 10 12:40:50 2024    oldstats

         1857432165 function calls (1853977584 primitive calls) in 391.159 seconds

# 并行加载统计信息
Wed Jan 10 13:05:13 2024    newstats

         12539 function calls in 31.319 seconds