炼丹秘术：为了赢，我重新捡起了BM25

炼丹笔记

发布于 2021-12-29 08:50:02

2630

发布于 2021-12-29 08:50:02

在实践中，无论是搜索问题，还是文本问题，如何找到相似的文本都是一个常见的场景，但TFIDF文本相似度计算用多了，年轻人往往会不记得曾经的经典。

毕业快4年了，最近准备梳理一下《我毕业这4年》，在整理文档时看到了好久之前的一个比赛，想起了当时TFIDF、BERT的方案都没在指标上赢过BM25的情景，本文作为“炼丹秘术”系列的第7期，我们来聊一聊相似文本搜索的相关知识点。

什么是BM25

BM25是信息索引领域用来计算Query与文档相似度得分的经典算法，不同于TFIDF，BM25的公式主要由三个部分组成：

对Query进行语素解析，生成语素qi；
对于每个搜索结果D，计算每个语素qi与D的相关性得分；
将qi相对于D的相关性得分进行加权求和，从而得到Query与D的相关性得分。

公式如下：

其中，Q表示Query，qi表示Q解析之后的一个语素，d表示一个搜索结果文档，Wi表示语速qi的权重，R(qi,d)表示语素qi与文档d的相关性得分。

BM25的实现

class BM25:
    def __init__(self, corpus, tokenizer=None):
        self.corpus_size = len(corpus)
        self.avgdl = 0
        self.doc_freqs = []
        self.idf = {}
        self.doc_len = []
        self.tokenizer = tokenizer

        if tokenizer:
            corpus = self._tokenize_corpus(corpus)

        nd = self._initialize(corpus)
        self._calc_idf(nd)

    def _initialize(self, corpus):
        nd = {}  # word -> number of documents with word
        num_doc = 0
        for document in corpus:
            self.doc_len.append(len(document))
            num_doc += len(document)

            frequencies = {}
            for word in document:
                if word not in frequencies:
                    frequencies[word] = 0
                frequencies[word] += 1
            self.doc_freqs.append(frequencies)

            for word, freq in frequencies.items():
                try:
                    nd[word]+=1
                except KeyError:
                    nd[word] = 1

        self.avgdl = num_doc / self.corpus_size
        return nd

    def _tokenize_corpus(self, corpus):
        pool = Pool(cpu_count())
        tokenized_corpus = pool.map(self.tokenizer, corpus)
        return tokenized_corpus

    def _calc_idf(self, nd):
        raise NotImplementedError()

    def get_scores(self, query):
        raise NotImplementedError()

    def get_batch_scores(self, query, doc_ids):
        raise NotImplementedError()

    def get_top_n(self, query, documents, n=5):

        assert self.corpus_size == len(documents), "The documents given don't match the index corpus!"

        scores = self.get_scores(query)
        top_n = np.argsort(scores)[::-1][:n]
        return [documents[i] for i in top_n]

数据竞赛实战案例

在这里，用自己之前写过的一个真实案例，来介绍一下。任务提供一个论文库（约含20万篇论文），同时提供对论文的描述段落，来自论文中对同类研究的介绍。任务需要为描述段落匹配三篇最相关的论文。

例子：
描述段落：
An efficient implementation based on BERT [1] and graph neural network (GNN) [2] is introduced.
相关论文：
[1] BERT: Pre-training of deep bidirectional transformers for language understanding.
[2] Relational inductive biases, deep learning, and graph networks.

从任务描述中我们可以看到，该任务需要对描述段落匹配三篇最相关的论文。单从形式上可以理解为这是一个“完形填空”任务。但相较于在本文的相应位置上填上相应的词语不同的是，这里需要填充的是一个Sentence，也就是论文的题目。但是如果你按照这个思路去寻求解决方案，你会发现在这个量级的文本数据上，一般算力是满足不了的。

既然如此，那我们不如换一个思路来思考这个问题，“对描述段落匹配三篇最相关的论文”，其实最简单的实现方式是计算描述段落和论文库里所有论文的相似度，找出最相似的即可。但这同样会存在一个问题，通过对数据进行探查你会发现“An efficient implementation based on BERT [1] and graph neural network (GNN) [2] is introduced.”这一描述段落，同时引用了两篇文章，那么在计算相似度时，到底哪个位置该是哪篇文章呢？

01建模核心思路

在解决本问题时，使用了两种方式，其一是利用Wrod2Vec方法，将描述段落利用Word2Vec得到每个词的词向量，同时对句子中的词使用IDF为权重进行加权得到Sentence Embedding，同时为了得到更好的效果，这里做了一个改进，即使用Smooth Inverse Frequency代替IDF作为每个词的权重；其二是利用BM25得到Sentence Embedding。两种方法各自计算余弦相似度得到3篇论文，去重后召回集中每个段落有3-6篇不等的召回论文。

02SIF Sentence Embedding