最精简的搜索系统实现：倒排索引

众所周知，搜推广业务在机器学习中一直都属于是应用面最广的部分，也是最直接为公司带来收益的数据科学业务。毕业之后我的第一个实习就是做简单的推荐系统，因此有过一些浅显的了解。如今尽管已经走上了另外一个方向，我仍然会偶尔看一下相关的内容。

今天带来一个简单的搜索引擎实现，只有短短几行代码。其中用到了jieba这个经典的中文python库，核心的思想则是倒排索引。

什么是倒排索引？

要记住倒排索引，关键就是文档–>ID。

搜索引擎中，每个文件都对应一个文件ID，而文件内容又被表示为一系列关键词的集合。例如“文档1”经过分词，提取了20个关键词，每个关键词都会记录它在文档中的出现次数和出现位置。

倒排索引的基本原理就是将文档中的词汇作为索引项，记录每个词汇所在的文档集合以及该词汇在每个文档中的位置信息。

过程如下：

  1. 预处理：对每个文档进行分词处理，将文档分割成一个个独立的词汇。
  2. 创建词汇表：将所有文档中的词汇汇总到一个词汇表中，去除重复的词汇。
  3. 构建索引：对于词汇表中的每个词汇，记录其所在的文档集合以及在每个文档中的位置信息。
  4. 存储索引：存储构建好的倒排索引以便后续的搜索查询。

假设商品A包含了以下特征：AAA，蓝色，M码，猴子。那么这几个特征就是商品A的索引项（Value），当用户检索了其中任何一个特征，就会返回这个商品的ID（Key）。这样用户通过搜“AAA”，“蓝色”，“M码”，“猴子”，均可找到该商品。

代码实现

下面是具体的代码实现。首先我们先构建一系列的文档：

1
2
3
4
5
6
7
8
9
10
from collections import defaultdict
import jieba

forward_index = [
    "苹果发布新款iPhone",
    "iPhone6s发布",
    "苹果发布iPhone6s",
    "苹果公布2023年财报",
    "苹果发布全新M3芯片的MacBook Pro",
]

随后构建一个生成索引表的字典：

1
2
3
4
5
6
7
def build_invert_index(docs: list[str]) -> dict[str, list[int]]:
    invert_index: dict[list[int]] = defaultdict(list)
    for doc in docs:
        words = jieba.cut_for_search(doc)
        for word in words:
            invert_index[word].append(docs.index(doc))
    return invert_index

最后就是检索的方法：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def retrive(query: str, inverted_index: dict[str, list[int]]) -> list[int]:
    words = jieba.cut_for_search(query)
    result = set()
    for word in words:
        if word in inverted_index:
            result.update(inverted_index[word])
    return list(result)


if __name__ == "__main__":
    invert_index = build_invert_index(forward_index)

    while True:
        query = input("请输入查询词：")
        if query == "q":
            break
        result = retrive(query, invert_index)
        print(f"查询词：{query}的结果为：")
        for doc in result:  
            print(forward_index[doc])

        print("=" * 20)

2024/3/10 于苏州