KeyBERT优化：支持SentenceTransformer的batch_size参数传递

背景介绍

KeyBERT是一个基于BERT的关键词提取工具，它利用预训练语言模型来识别文档中最相关的关键词。在实际应用中，当处理大规模文档时，GPU资源的有效利用变得尤为重要。其中，SentenceTransformer作为KeyBERT的常用后端之一，其encode方法支持batch_size参数调整，这对GPU内存利用率和处理速度有显著影响。

问题分析

当前KeyBERT的设计存在一个局限性：虽然SentenceTransformer的encode方法允许通过batch_size参数控制批量处理大小，但这一参数无法通过KeyBERT的extract_sentences方法直接设置。这导致用户无法灵活调整批量大小以优化GPU资源利用，特别是在处理大规模文档时，这种限制会显著影响处理效率。

技术解决方案

经过社区讨论，确定了两种可能的改进方案：

1. 通用kwargs方案：通过**kwargs将参数传递给底层encode方法
2. 专用后端方案：在SentenceTransformerBackend中增加专门的encode_kwargs参数

最终选择了第二种方案，原因如下：

• 保持API的明确性和可读性
• 避免为特定后端引入通用参数
• 更符合软件设计原则

实现细节

在SentenceTransformerBackend中新增encode_kwargs参数，允许用户传递包括batch_size在内的各种编码参数。这一改动使得高级用户能够：

• 根据GPU内存容量调整批量大小
• 优化处理速度与内存占用的平衡
• 利用SentenceTransformer的其他编码参数

性能影响

适当增大batch_size可以带来显著的性能提升：

• 减少GPU内核启动开销
• 提高内存带宽利用率
• 降低数据传输频率
• 充分利用并行计算能力

使用建议

对于不同规模的文档处理，建议：

• 小规模文档：使用默认batch_size
• 中等规模文档：尝试256-512的batch_size
• 大规模文档：根据GPU内存测试最优batch_size

总结

这一改进使KeyBERT在处理大规模文档时更加灵活高效，特别是对于需要频繁处理大量文本的生产环境。通过允许用户调整batch_size参数，KeyBERT能够更好地适应不同硬件配置和处理需求，为用户提供更优的性能体验。