用于并行处理的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
可能遇到的错误
- 操作系统兼容性问题:如前文所述,Windows和Linux/mac在处理并行任务时表现不同,在Windows上可能会得到较小的性能提升。
- 文件路径问题:确保输入目录路径正确并且目录中包含需要处理的文件。
- 资源限制:并行处理需要更多的系统资源。在系统资源不足的情况下,可能会出现性能下降或系统崩溃的情况。
如果你觉得这篇文章对你有帮助,请点赞,关注我的博客,谢谢!
参考资料: