|
|
|
引言今天带来一篇BM25变种的论文笔记，不要低估BM25，在RAG中检索中通常都会引入BM25检索，然后配合嵌入模型进行混合检索。BM25S:Ordersofmagnitudefasterlexicalsearchviaeagersparsescoring，题目翻译过来是：通过快速稀疏评分实现数量级速度提升的词汇检索。BM25S是一种高效的基于Python的BM25实现，仅依赖于Numpy和Scipy。与最流行的基于Python的框架(Rank-BM25)相比，BM25S在索引期间通过即时(eager)计算BM25分数并将其存储到稀疏矩阵中，实现了高达500倍的速度提升。代码开源在：https://github.com/xhluca/bm25s1.背景稀疏词汇搜索算法(Sparselexicalsearchalgorithms)，比如BM25，仍然得到了广泛地应用，因为不需要训练，可应用于多种语言，并且速度较快。尤其是基于Lucene的实现，通常比现有的基于Python的实现(如Rank-BM25)更快。本篇工作通过引入两项改进，能带来比现有基于Python实现显著的速度提升：在索引语料库时即时计算所有可能分配给任何未来查询标记的分数，并将这些计算存储在稀疏矩阵中，以便实现更快速的切片和求和。稀疏矩阵的思想在BM25-PT中被探索，它使用PyTorch预计算BM25分数，并通过稀疏矩阵乘法将其与查询的词袋编码相乘。而本篇工作的BM25S不依赖于Pytorch，使用Scipy的稀疏矩阵实现，BM25-PT通过词袋与文档矩阵相乘，而BM25S则切片相关索引并在标记维度上进行求和，消除了矩阵乘法的需求。实现时BM25S还引入了一种简单快速的基于Python的分词器。2.实现BM25的计算我们使用Lucene提出的评分方法，对于一个给定查询QQQ(分词为q1,⋯ ,q∣Q∣)q_1,\cdots,q_{|Q|})q1,⋯,q∣Q∣)和来自集合CCC中的文档DDD，计算下面的分数：B(Q,D)=∑i=1∣Q∣S(qi,D)=∑i=1∣Q∣IDF(qi,C)TF(qi,D)DB(Q,D)=∑i=1|Q|S(qi,D)=∑i=1|Q|IDF(qi,C)TF(qi,D)DB(Q,D)=∑i=1|Q|S(qi,D)=∑i=1|Q|IDF(qi,C)TF(qi,D)DB(Q,D)=i=1∑∣Q∣S(qi,D)=i=1∑∣Q∣IDF(qi,C)DTF(qi,D)其中D=TF(t,D)+k1(1−b+b∣D∣Lavg)\mathcalD=\text{TF}(t,D)+k_1\left(1-b+b\frac{|D|}{L_{\text{avg}}}\right)D=TF(t,D)+k1(1−b+bLavg∣D∣)，LavgL_{\text{avg}}Lavg是语料库CCC中的平均文档长度(通过标记数计算)，TF(qi,D)\text{TF}(q_i,D)TF(qi,D)是标记qiq_iqi在文档DDD中的词频，IDF\text{IDF}IDF是逆文档频率，计算为：IDF(qi,C)=ln(∣C∣−DF(qi,C)+0.5DF(qi,C)+0.5+1)\text{IDF}(q_i,C)=\ln\left(\frac{|C|-\text{DF}(q_i,C)+0.5}{\text{DF}(q_i,C)+0.5}+1\right)IDF(qi,C)=ln(DF(qi,C)+0.5∣C∣−DF(qi,C)+0.5+1)其中文档频率DF(qi,C)\text{DF}(q_i,C)DF(qi,C)是CCC中包含qiq_iqi的文档数。尽管B(Q,D)B(Q,D)B(Q,D)依赖于查询(query)，而查询仅在检索期间提供，下面会展示如何重构等式，以便在索引时即时计算TF\text{TF}TF和IDF\text{IDF}IDF。即时索引评分现在考虑词表VVV中的所有标记，记作t∈Vt\inVt∈V。我们可以重构S(t,D)S(t,D)S(t,D)为：S(t,D)=TF(t,D)⋅IDF(t,C)1DS(t,D)=\text{TF}(t,D)\cdot\text{IDF}(t,C)\frac{1}{\mathcalD}S(t,D)=TF(t,D)⋅IDF(t,C)D1当ttt是文档DDD中不存在的标记时，TF(t,D)=0\text{TF}(t,D)=0TF(t,D)=0，这也导致S(t,D)=0S(t,D)=0S(t,D)=0。这意味着，对于词表VVV中的大多数标记，我们可以简单地将相关性评分设置为0，仅计算实际出现在文档DDD中的ttt的值。这个计算可以在索引过程中完成，从而避免在查询时计算S(qi,D)S(q_i,D)S(qi,D)。分配查询评分考虑形状为∣V∣×∣C∣|V|×|C|∣V∣×∣C∣的稀疏矩阵，可以使用查询标记选择相关的行，从而得到一个形状为∣Q∣×∣C∣|Q|×|C|∣Q∣×∣C∣的矩阵，然后可以在列维度上进行求和，最终得到一个单一的∣C∣|C|∣C∣维向量(表示每个文档对查询的评分)。高效矩阵稀疏采用压缩稀疏列(CompressedSparseColumn,CSC)格式实现稀疏矩阵(使用scipy.sparse.csc_matrix)，这在坐标格式与CSC格式之间提供了高效的转换。由于我们在列维度上进行切片和求和，这一实现是稀疏矩阵实现中的最佳选择。在实践中，我们直接使用Numpy数组来复制稀疏操作。分词为了拆分文本，使用与Scikit-Learn为其自己的标记器所用的相同正则表达式模式，模式为r"(?u)\b\w\w+\b"。该模式便捷地解析UTF-8中的单词(涵盖多种语言)，其中\b处理单词边界。然后，如果需要词干提取，我们可以对词汇表中的所有单词进行词干提取，以便查找集合中每个单词的词干版本。最后，我们构建一个字典，将每个唯一的(词干)单词映射到一个整数索引，以便将标记转换为相应的索引，从而显著减少内存使用并使其能够用于切片Scipy矩阵和Numpy数组。Top-k选择在计算完集合中所有文档的评分后，我们可以通过选择前k个最相关的元素来完成搜索过程。一种简单的方法是对评分向量进行排序，并选择最后k个元素；而作者采用的是对数组的划分，仅选择最后k个文档(无序)。使用像Quickselect这样的算法，可以在平均时间复杂度为O(n)O(n)O(n)的情况下完成这一操作，其中n是集合中的文档数量，而排序需要O(nlogn)O(n\logn)O(nlogn)。如果用户希望按顺序接收前k个结果，排序划分后的文档将额外消耗O(klogk)O(klogk)O(klogk)的时间复杂度，这在假设k≪nk≪nk≪n的情况下是可以忽略不计的。在实践中，BM25S允许使用两种实现：一种基于numpy，利用np.argpartition，另一种基于jax，依赖于XLA的top-k实现。Numpy的argpartition使用了内省选择算法，该算法修改了quickselect算法，以确保最坏情况下的性能保持在O(n)O(n)O(n)内。实际上观察到JAX的实现上表现更佳。多线程通过池化执行器实现可选的多线程功能，以在检索过程中进一步加速。BM25的替代实现上述内容描述了如何为BM25的一种变体(Lucene)实现BM25S。然而，可以很容易地将BM25S方法扩展到许多BM25的变体。2.1基于未出现调整扩展稀疏性对于BM25L、BM25+，当TF(t,D)=0\text{TF}(t,D)=0TF(t,D)=0时，S(t,D)S(t,D)S(t,D)的值不会为零。将这个值记作标量Sθ(t)S^θ(t)Sθ(t)，表示当ttt不出现在文档DDD中时的得分。显然，构建一个∣V∣×∣C∣|V|×|C|∣V∣×∣C∣的稠密矩阵会消耗过多内存。相反，我们仍然可以通过从得分矩阵中每个标记ttt和文档DDD中减去Sθ(t)S^θ(t)Sθ(t)来实现稀疏性(因为词汇中的大多数标记ttt不会出现在任何给定文档DDD中，它们在得分矩阵中的值将为0)。然后，在检索过程中，我们只需计算每个查询qi∈Qq_i∈Qqi∈Q的Sθ(qi)S^θ(q_i)Sθ(qi)，并对其进行求和，以获取一个可以加到最终得分上的标量。更正式地，对于一个空文档∅∅∅，定义Sθ(t)=S(t,∅)S^θ(t)=S(t,∅)Sθ(t)=S(t,∅)为标记ttt的非出现得分。然后，差分S∆(t,D)S^∆(t,D)S∆(t,D)定义为：S∆(t,D)=S(t,D)−Sθ(t)S^∆(t,D)=S(t,D)-S^\theta(t)S∆(t,D)=S(t,D)−Sθ(t)因此，我们重构BM25(B)评分为：B(Q,D)=∑i=1∣Q∣S(qi,D)=∑i=1∣Q∣(S(qi,D)−Sθ(qi)+Sθ(qi))=∑i=1∣Q∣(SΔ(qi,D)+Sθ(qi))=∑i=1∣Q∣SΔ(qi,D)+∑i=1∣Q∣Sθ(qi)B(Q,D)=∑i=1|Q|S(qi,D)=∑i=1|Q|(S(qi,D)−Sθ(qi)+Sθ(qi))=∑i=1|Q|(SΔ(qi,D)+Sθ(qi))=∑i=1|Q|SΔ(qi,D)+∑i=1|Q|Sθ(qi)B(Q,D)=∑i=1|Q|S(qi,D)=∑i=1|Q|(S(qi,D)−Sθ(qi)+Sθ(qi))=∑i=1|Q|(SΔ(qi,D)+Sθ(qi))=∑i=1|Q|SΔ(qi,D)+∑i=1|Q|Sθ(qi)B(Q,D)=i=1∑∣Q∣S(qi,D)=i=1∑∣Q∣(S(qi,D)−Sθ(qi)+Sθ(qi))=i=1∑∣Q∣(SΔ(qi,D)+Sθ(qi))=i=1∑∣Q∣SΔ(qi,D)+i=1∑∣Q∣Sθ(qi)其中∑i=1∣Q∣SΔ(qi,D)\sum_{i=1}^{|Q|}S^\Delta(q_i,D)∑i=1∣Q∣SΔ(qi,D)可以使用差分稀疏得分矩阵高效计算在scipy中实现。此外，∑i=1∣Q∣Sθ(qi)\sum_{i=1}^{|Q|}S^\theta(q_i)∑i=1∣Q∣Sθ(qi)只需要对查询Q计算一次，随后可以应用于每个检索到的文档以获得确切的得分。3.基准测试吞吐量为了进行基准测试，使用BEIR基准中公开可用的数据集。表1中的结果显示，BM25S的速度显著快于Rank-BM25。分词的影响进一步通过比较BM25SLucene(k1=1.5,b=0.75k_1=1.5,b=0.75k1=1.5,b=0.75)的表现，检查分词对每个模型的影响，分别为(1)不进行词干提取，(2)不去除停用词，(3)两者都不去除，以及(4)都去除。总体而言，添加词干提取器平均上会提高得分，而停用词的影响较小。然而，在个别情况下，停用词可能会产生更大的影响。比较模型变体在表3中，比较了多个实现变体。大多数实现的平均得分在39.7到40之间，只有Elastic实现了稍高的得分。这种差异可以归因于分词方案的不同。4.结论作者提供了一种新颖的计算BM25分数的方法，称为BM25S，该方法不仅提供开箱即用的快速分词和高效的前K个选择，还最小化了依赖关系，使其能够直接在Python中使用。通过最小化依赖关系，BM25S成为在存储可能受到限制的场景(例如边缘部署)中的良好选择。实战首先安装相相关库：pipinstallbm25s1该项目默认只支持英文系列的语言，但我们可以简单修改，加入结巴分词让它支持中文：frombm25s.tokenizationimportTokenizedimportjiebafromtypingimportList,Unionfromtqdm.autoimporttqdmdeftokenize(texts,return_ids:bool=True,show_progress:bool=True,leave:bool=False,)->Union[List[List[str]],Tokenized]:ifisinstance(texts,str):texts=[texts]corpus_ids=[]token_to_index={}fortextintqdm(texts,desc="Splitstrings",leave=leave,disable=notshow_progress):splitted=jieba.lcut(text)doc_ids=[]fortokeninsplitted:iftokennotintoken_to_index:token_to_index[token]=len(token_to_index)token_id=token_to_index[token]doc_ids.append(token_id)corpus_ids.append(doc_ids)#Createalistofuniquetokensthatwewillusetocreatethevocabularyunique_tokens=list(token_to_index.keys())vocab_dict=token_to_index#ReturnthetokenizedIDsandthevocabdictionaryorthetokenizedstringsifreturn_ids:returnTokenized(ids=corpus_ids,vocab=vocab_dict)else:#WeneedareversedictionarytoconvertthetokenIDsbacktotokensreverse_dict=unique_tokens#WeconvertthetokenIDsbacktotokensin-placefori,token_idsinenumerate(tqdm(corpus_ids,desc="Reconstructingtokenstrings",leave=leave,disable=notshow_progress,)):corpus_ids[i]=[reverse_dict[token_id]fortoken_idintoken_ids]returncorpus_ids12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758然后通过bm25s.tokenize=tokenize替换默认，下面参考官网给出一个简单的例子：importbm25sbm25s.tokenize=tokenize#Createyourcorpusherecorpus=["今天天气晴朗，我的心情美美哒","小明和小红一起上学","我们来试一试吧","我们一起学猫叫","我和Faker五五开","明天预计下雨，不能出去玩了",]#corpus=load_corpus("data")#Tokenizethecorpusandindexitcorpus_tokens=bm25s.tokenize(corpus)print(corpus_tokens)retriever=bm25s.BM25(corpus=corpus)retriever.index(corpus_tokens)query="明天天气怎么样"query_tokens=bm25s.tokenize(query)docs,scores=retriever.retrieve(query_tokens,k=3)print(f"Bestresult(score:{scores[0,0]:.2f}):{docs[0,0]}")print(docs,scores)#Happywithyourindex?Saveitforlater...retriever.save("bm25s_index_animals")#...andloaditwhenneededret_loaded=bm25s.BM25.load("bm25s_index_animals",load_corpus=True)12345678910111213141516171819202122232425262728293031323334Tokenized(ids=[[0,1,2,3,4,5,6,7],[8,9,10,11],[12,13,14,15],[12,10,16,17],[3,18,19,20],[21,22,23,2,24,25,26]],vocab={'今天':0,'天气晴朗':1,'，':2,'我':3,'的':4,'心情':5,'美美':6,'哒':7,'小明':8,'和小红':9,'一起':10,'上学':11,'我们':12,'来':13,'试一试':14,'吧':15,'学':16,'猫叫':17,'和':18,'Faker':19,'五五开':20,'明天':21,'预计':22,'下雨':23,'不能':24,'出去玩':25,'了':26})Bestresult(score:0.53):明天预计下雨，不能出去玩了[['明天预计下雨，不能出去玩了''今天天气晴朗，我的心情美美哒''小明和小红一起上学']][[0.53133570.0.]]1234
|
|
|手机版|心飞设计-版权所有：微度网络信息技术服务中心
( 鲁ICP备17032091号-12 )|网站地图
发表于 2024-9-13 11:24:43
