用150行Python代码构建全文搜索引擎

全文搜索无处不在。无论是在Scribd上找书，Netflix上找电影，亚马逊上找卫生纸，还是通过谷歌在网络上搜索任何东西（比如如何做好软件工程师的工作），你今天已经多次搜索了大量的非结构化数据。更令人惊讶的是，尽管你搜索了数百万（甚至数十亿）条记录，你仍然能在几毫秒内得到结果。在这篇文章中，我们将探索全文搜索引擎的基本组成部分，并用它们构建一个能够在几毫秒内搜索数百万文档并根据相关性进行排序的搜索引擎，代码不超过150行Python！

数据

本文中的所有代码都可以在Github上找到。我们将使用英文维基百科的文章摘要作为搜索数据，目前这是一个约785MB的gzip压缩XML文件，包含约627万条摘要1。你可以手动下载这个文件，或者使用我写的简单函数来下载。

数据准备

文件是一个包含所有摘要的大型XML文件。每个摘要由一个<doc>元素包裹，大致如下所示（我们省略了不感兴趣的元素）：

<doc>
  <title>Wikipedia: London Beer Flood</title>
  <url>https://en.wikipedia.org/wiki/London_Beer_Flood</url>
  <abstract>The London Beer Flood was an accident at Meux & Co's Horse Shoe Brewery, London, on 17 October 1814. It took place when one of the wooden vats of fermenting porter burst.</abstract>
  ...
</doc>

我们感兴趣的部分是title、url和abstract文本本身。我们将使用Python的dataclass来表示文档，方便数据访问。我们还会添加一个属性，将标题和摘要内容拼接在一起。代码可以在这里找到。

from dataclasses import dataclass

@dataclass
class Abstract:
  """维基百科摘要"""
  ID: int
  title: str
  abstract: str
  url: str

  @property
  def fulltext(self):
    return ' '.join([self.title, self.abstract])

接下来，我们将从XML中提取摘要数据并解析它，以便创建Abstract对象的实例。我们将流式读取gzip压缩的XML文件，而不将整个文件加载到内存中2。我们将为每个文档分配一个ID（即第一个文档的ID=1，第二个文档的ID=2，依此类推）。代码可以在这里找到。

import gzip
from lxml import etree
from search.documents import Abstract

def load_documents():
  # 打开gzip压缩的维基百科文件
  with gzip.open('data/enwiki.latest-abstract.xml.gz', 'rb') as f:
    doc_id = 1
    # iterparse会在找到`</doc>`标签时返回整个`doc`元素
    for _, element in etree.iterparse(f, events=('end',), tag='doc'):
      title = element.findtext('./title')
      url = element.findtext('./url')
      abstract = element.findtext('./abstract')
      yield Abstract(ID=doc_id, title=title, url=url, abstract=abstract)
      doc_id += 1
      # 显式释放元素占用的内存
      element.clear()

索引

我们将把这些数据存储在一个称为“倒排索引”或“倒排列表”的数据结构中。你可以把它想象成书后的索引，其中列出了相关的单词和概念，并告诉读者在哪些页码可以找到它们。

书后索引

实际上，这意味着我们将创建一个字典，将语料库中的所有单词映射到它们出现的文档ID。看起来像这样：

{
  ...
  "london": [5245250, 2623812, 133455, 3672401, ...],
  "beer": [1921376, 4411744, 684389, 2019685, ...],
  "flood": [3772355, 2895814, 3461065, 5132238, ...],
  ...
}

注意，字典中的单词是小写的；在构建索引之前，我们将把原始文本分解为单词或“标记”。首先，我们将文本“分词”成单词，然后对每个标记应用零个或多个“过滤器”（如小写转换或词干提取），以提高查询与文本匹配的几率。

分词

分析

我们将应用非常简单的分词方法，即按空格分割文本。然后，我们将对每个标记应用几个过滤器：将每个标记转换为小写，删除任何标点符号，删除英语中最常见的25个单词（以及“wikipedia”，因为它出现在每个摘要的标题中），并对每个单词进行词干提取（确保不同形式的单词映射到相同的词干，如_brewery_和_breweries_3）。

分词和小写过滤器非常简单：

import Stemmer

STEMMER = Stemmer.Stemmer('english')

def tokenize(text):
  return text.split()

def lowercase_filter(tokens):
  return [token.lower() for token in tokens]

def stem_filter(tokens):
  return STEMMER.stemWords(tokens)

标点符号只是一个正则表达式：

import re
import string

PUNCTUATION = re.compile('[%s]' % re.escape(string.punctuation))

def punctuation_filter(tokens):
  return [PUNCTUATION.sub('', token) for token in tokens]

停用词是非常常见的单词，我们预计它们会出现在语料库中的几乎每个文档中。因此，当我们搜索这些词时，它们不会贡献太多（即几乎每个文档都会匹配），只会占用空间，所以我们将在索引时过滤掉它们。维基百科摘要语料库在每个标题中都包含“Wikipedia”这个词，所以我们也会把这个词添加到停用词列表中。我们删除了英语中最常见的25个单词。

# 英语中最常见的25个单词和“wikipedia”：
# https://en.wikipedia.org/wiki/Most_common_words_in_English
STOPWORDS = set(['the', 'be', 'to', 'of', 'and', 'a', 'in', 'that', 'have',
         'I', 'it', 'for', 'not', 'on', 'with', 'he', 'as', 'you',
         'do', 'at', 'this', 'but', 'his', 'by', 'from', 'wikipedia'])

def stopword_filter(tokens):
  return [token for token in tokens if token not in STOPWORDS]

将这些过滤器组合在一起，我们将构建一个analyze函数，该函数将对每个摘要中的text进行操作；它将文本分词为单个单词（或更准确地说，标记），然后依次对标记列表应用每个过滤器。顺序很重要，因为我们使用的是未词干化的停用词列表，所以我们应该在stem_filter之前应用stopword_filter。

def analyze(text):
  tokens = tokenize(text)
  tokens = lowercase_filter(tokens)
  tokens = punctuation_filter(tokens)
  tokens = stopword_filter(tokens)
  tokens = stem_filter(tokens)
  return [token for token in tokens if token]

索引语料库

我们将创建一个Index类来存储index和documents。documents字典按ID存储数据类，index的键将是标记，值是该标记出现的文档ID：

class Index:
  def __init__(self):
    self.index = {}
    self.documents = {}

  def index_document(self, document):
    if document.ID not in self.documents:
      self.documents[document.ID] = document
    for token in analyze(document.fulltext):
      if token not in self.index:
        self.index[token] = set()
      self.index[token].add(document.ID)

搜索

现在我们已经索引了所有标记，搜索查询就变成了用与文档相同的分析器分析查询文本的问题；这样我们最终会得到与索引中标记匹配的标记。对于每个标记，我们将在字典中查找，找到该标记出现的文档ID。我们对每个标记都这样做，然后找到所有这些集合中文档的ID（即，文档要匹配查询，它需要包含查询中的所有标记）。然后，我们将获取生成的文档ID列表，并从documents存储中获取实际数据4。

def _results(self, analyzed_query):
  return [self.index.get(token, set()) for token in analyzed_query]

def search(self, query):
  """
  布尔搜索；这将返回包含查询中所有单词的文档，但不进行排序（集合很快，但无序）。
  """
  analyzed_query = analyze(query)
  results = self._results(analyzed_query)
  documents = [self.documents[doc_id] for doc_id in set.intersection(*results)]
  return documents

In [1]: index.search('London Beer Flood')
search took 0.16307830810546875 milliseconds
Out[1]:
[Abstract(ID=1501027, title='Wikipedia: Horse Shoe Brewery', abstract='The Horse Shoe Brewery was an English brewery in the City of Westminster that was established in 1764 and became a major producer of porter, from 1809 as Henry Meux & Co. It was the site of the London Beer Flood in 1814, which killed eight people after a porter vat burst.', url='https://en.wikipedia.org/wiki/Horse_Shoe_Brewery'),
 Abstract(ID=1828015, title='Wikipedia: London Beer Flood', abstract="The London Beer Flood was an accident at Meux & Co's Horse Shoe Brewery, London, on 17 October 1814. It took place when one of the wooden vats of fermenting porter burst.", url='https://en.wikipedia.org/wiki/London_Beer_Flood')]

现在，这将使我们的查询非常精确，尤其是对于长查询字符串（查询中包含的标记越多，文档包含所有这些标记的可能性就越小）。我们可以通过允许用户指定只需要一个标记出现来优化搜索函数，以提高召回率而不是精确率：

def search(self, query, search_type='AND'):
  """
  仍然是布尔搜索；这将返回包含查询中所有单词或仅其中一个单词的文档，具体取决于指定的search_type。
  我们仍然没有对结果进行排序（集合很快，但无序）。
  """
  if search_type not in ('AND', 'OR'):
    return []
  analyzed_query = analyze(query)
  results = self._results(analyzed_query)
  if search_type == 'AND':
    # 所有标记必须在文档中
    documents = [self.documents[doc_id] for doc_id in set.intersection(*results)]
  if search_type == 'OR':
    # 只需要一个标记在文档中
    documents = [self.documents[doc_id] for doc_id in set.union(*results)]
  return documents

In [2]: index.search('London Beer Flood', search_type='OR')
search took 0.02816295623779297 seconds
Out[2]:
[Abstract(ID=5505026, title='Wikipedia: Addie Pryor', abstract='| birth_place  = London, England', url='https://en.wikipedia.org/wiki/Addie_Pryor'),
 Abstract(ID=1572868, title='Wikipedia: Tim Steward', abstract='|birth_place     = London, United Kingdom', url='https://en.wikipedia.org/wiki/Tim_Steward'),
 Abstract(ID=5111814, title='Wikipedia: 1877 Birthday Honours', abstract='The 1877 Birthday Honou