Полнотекстовый поиск в 150 строках Питона

Полнотекстовый поиск можно найти повсюду. Начиная c выбора книг на Scribd и заканчивая поиском фильмов на Netflix, туалетной бумаги на Amazon или вообще чего угодно через Google (например, как работать инженером-программистом), только за сегодня вы многократно искали что-то среди огромных объемов неструктурированных данных. Что еще более удивительно, так это то, что даже если вы искали среди миллионов (или миллиардов) записей, вы получили ответ за миллисекунды. В этом посте мы рассмотрим основные компоненты полнотекстового поиска и используем их для создания системы, которая может искать в миллионах документов и ранжировать их по релевантности за миллисекунды, менее чем за 150 строк кода на Python!

Данные

Весь код из этой статьи можно найти на Github. Я буду приводить ссылки на фрагменты кода, чтобы вы могли запускать его самостоятельно. Вы можете запустить пример полностью, установив зависимости (pip install -r requirements.txt) и запустив python run.py. Так вы загрузите все данные и выполните образец запроса с ранжированием и без.

Прежде чем мы перейдем к созданию поискового движка, нам сначала понадобятся текстовые неструктурированные данные, по которым мы будем искать. В качестве таких данных мы возьмем аннотации статей из английской Википедии, которые в настоящее время представляют собой сжатый gzip'ом XML-файл размером около 785 Мб и объемом в 6,27 миллионов аннотаций. Я написал простую функцию для загрузки этого XML-файла, но вы можете просто загрузить файл вручную.

Подготовка данных

Файл представляет собой один большой 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. Мы обернем документы в dataclass'ы Python для удобной работы с ними. К ним мы добавим свойство, которое объединяет заголовок и содержимое аннотации. Код можно найти здесь.

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. Мы будем просматривать XML в формате gzip, не загружая весь файл в память. Мы присвоим каждому документу 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()` явно освободит память, 
            # выделенную под элемент
            element.clear()

Индексирование

Мы будем хранить все это в структуре данных, известной как "инвертированный индекс" ("inverted index" или "postings list"). Его можно сравнить с алфавитным указателем в конце книги, в котором есть список соответствующих слов и понятий, а также номера страниц, на которых читатель может их найти.

На практике это означает, что мы создадим справочник, в котором все слова нашего текстового корпуса будут сопоставлены с ID документов, в которых они встречаются. Это будет выглядеть примерно так:

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

Обратите внимание, что в приведенном примере слова написаны строчными буквами; перед построением индекса мы собираемся разбить (или проанализировать) необработанный текст на список слов (или лексем, токенов). Идея заключается в том, что сначала мы разбиваем текст на токены, а затем применяем к каждому токену ноль или более фильтров (например, перевод в нижний регистр или стемминг) для повышения вероятности совпадения запросов с текстом.

Анализ

Мы токенизируем текст очень простым методом, просто разбив его на пробелы. Затем мы применим несколько фильтров к каждой лексеме: мы переведем ее в нижний регистр, удалим все знаки препинания, удалим 25 самых распространенных слов в английском языке (плюс слово "wikipedia", потому что оно встречается в каждом заголовке в каждом реферате) и применим стемминг к слову (чтобы убедиться, что различные формы слова относятся к одной и той же основе, например, brewery и breweries).

Токенизация и перевод в нижний регистр очень просты:

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, которая будет обрабатывать текст каждой аннотации; она будет разбивать текст на отдельные слова (точнее, лексемы), а затем последовательно применять каждый фильтр к списку лексем. Порядок важен, потому что мы используем список стоп-слов, не прошедших стемминг, поэтому фильтр stopword_filter должен применяться перед фильтром stem_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, которые есть во всех этих наборах (т.е. чтобы документ соответствовал запросу, он должен содержать все токены в запросе). Затем мы возьмем полученный список идентификаторов и извлечем фактические данные из нашего хранилища documents.

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

def search(self, query):
    """
    Булевый поиск; вернет документы, которые содержат все слова 
    из запроса, но не применит ранжирование к ним 
    (множества (set) быстрые, но не упорядоченные).
    """
    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')]

Это алгоритм сделает наши запросы очень точными, особенно длинные запросы (чем больше токенов содержит наш запрос, тем меньше вероятность того, что мы найдем документ со всеми этими токенами). Мы можем оптимизировать нашу функцию поиска не для точности (precision), а для чувствительности (recall), позволяя пользователям указывать, что для соответствия запросу достаточно только одного вхождения лексемы:

def search(self, query, search_type='AND'):
    """
    Все еще булевый поиск; вернет документы, которые либо содержат все слова из запроса,
    либо хотя бы одно слово из запроса, в зависимости от search_type.

    Мы все еще не ранжируем результаты (множества (set) быстрые, но не упорядоченные).
    """
    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 Honours were appointments by Queen Victoria to various orders and honours to reward and highlight good works by citizens of the British Empire. The appointments were made to celebrate the official birthday of the Queen, and were published in The London Gazette on 30 May and 2 June 1877.', url='https://en.wikipedia.org/wiki/1877_Birthday_Honours'),
 ...
In [3]: len(index.search('London Beer Flood', search_type='OR'))
search took 0.029065370559692383 seconds
Out[3]: 49627

Релевантность

Мы реализовали довольно быстрый поисковый движок с помощью несложного кода на Python, но есть один аспект, который явно отсутствует в нашем маленьком движке, и это релевантность. Пока что мы просто возвращаем неупорядоченный список документов, и оставляем за пользователем выяснять, какие из них действительно его интересуют. Для больших наборов результатов это мучительно или просто-напросто невозможно (в нашем примере с OR почти 50.000 результатов).

Так возникает идея релевантности; что если бы мы могли оценить каждый документ по шкале, показывающей, насколько он соответствует запросу, и просто упорядочить их по этим значениям? Один из самых наивных и простых способов оценить релевантность документа запросу - просто подсчитать, как часто в документе упоминается конкретное слово. В конце концов, чем чаще в документе упоминается термин, тем больше вероятность того, что он относится к нашему запросу!

Частота вхождения

Давайте расширим наш класс Abstract, чтобы он вычислял и хранил частоту вхождения каждого слова при индексировании. Таким образом, у нас будет быстрый доступ к этим числам, когда мы захотим ранжировать наш список документов:

# в documents.py
from collections import Counter
from .analysis import analyze

@dataclass
class Abstract:
    # вырезано
    def analyze(self):
        # Counter создаст словарь, подсчитывающий уникальные значения в списке:
        # {'london': 12, 'beer': 3, ...}
        self.term_frequencies = Counter(analyze(self.fulltext))

    def term_frequency(self, term):
        return self.term_frequencies.get(term, 0)

Затем нам нужно подсчитывать эти частоты при индексировании данных:

# в index.py мы добавляем document.analyze()

def index_document(self, document):
    if document.ID not in self.documents:
        self.documents[document.ID] = document
        document.analyze()

Изменим нашу функцию поиска, чтобы можно было ранжировать документы в списке результатов. Мы получим документы с помощью того же самого булева запроса из индекса и хранилища документов, а затем для каждого документа в этом наборе результатов мы просто просуммируем, как часто каждый термин встречается в этом документе.

def search(self, query, search_type='AND', rank=True):
    # вырезано
    if rank:
        return self.rank(analyzed_query, documents)
    return documents


def rank(self, analyzed_query, documents):
    results = []
    if not documents:
        return results
    for document in documents:
        score = sum([document.term_frequency(token) for token in analyzed_query])
        results.append((document, score))
    return sorted(results, key=lambda doc: doc[1], reverse=True)

Обратная частота документа

Гораздо лучше, но все еще есть очевидные недостатки. Мы считаем, что все слова в запросе имеют одинаковую ценность при оценке релевантности запроса. Однако вполне вероятно, что некоторые термины практически не влияют на релевантность; например в библиотеке с большим количеством документов о пиве термин "пиво" будет часто встречаться почти в каждом документе (на самом деле, мы уже пытаемся решить эту проблему, исключая из индекса 25 наиболее распространенных английских слов). Поиск по слову "пиво" в таком случае будет, по сути, простой случайной сортировкой.

Чтобы решить эту проблему, мы добавим в алгоритм ранжирования еще один компонент, который уменьшит влияние на итоговую оценку документа тех слов, которые встречаются в индексе очень часто. Мы могли бы использовать частоту коллекции термина (т.е. как часто этот термин встречается во всех документах), но на практике вместо этого используется частота документа (т.е. сколько документов в индексе содержат этот термин). В конце концов, мы ранжируем документы, поэтому имеет смысл брать статистику на уровне документов.

Мы вычислим обратную частоту документов (idf) для термина j, разделив количество документов (n) в индексе на количество документов (df), содержащих этот термин, и возьмем логарифм этого значения.

При ранжировании мы будем просто умножать частоту термина на обратную частоту документа, так что совпадения с терминами, которые редко встречаются в корпусе, будут вносить больший вклад в оценку релевантности. Мы можем легко вычислить обратную частоту документа на основе данных, имеющихся в нашем индексе:

# index.py
import math

def document_frequency(self, token):
    return len(self.index.get(token, set()))

def inverse_document_frequency(self, token):
    # Мэннинг, Хинрих и Шютце используют log10, и мы будем, 
    # хотя разницы на самом деле нет
    # https://nlp.stanford.edu/IR-book/html/htmledition/inverse-document-frequency-1.html
    return math.log10(len(self.documents) / self.document_frequency(token))

def rank(self, analyzed_query, documents):
    results = []
    if not documents:
        return results
    for document in documents:
        score = 0.0
        for token in analyzed_query:
            tf = document.term_frequency(token)
            idf = self.inverse_document_frequency(token)
            score += tf * idf
        results.append((document, score))
    return sorted(results, key=lambda doc: doc[1], reverse=True)

Что дальше

И вот у нас готова базовая поисковая система всего в нескольких строках кода на Python! Вы можете найти весь код на Github, в том числе вспомогательную функцию, которая загрузит аннотации из Википедии и построит индекс. Установите зависимости, запустите код в вашей любимой консоли Python и наслаждайтесь работой со структурами данных и поиском.

Очевидно, что это проект для иллюстрации концепции поиска и демонстрации того, как он может быть настолько быстрым (я могу найти и даже ранжировать 6,27 млн. документов на своем ноутбуке с помощью такого "медленного" языка, как Python), но это не программа для применения в продакшн. Она работает полностью в памяти на моем ноутбуке, в то время как библиотеки вроде Lucene используют гиперэффективные структуры данных и даже оптимизируют поиск на диске, а такое ПО, как Elasticsearch и Solr, масштабируют Lucene на сотни, а то и тысячи машин.

Однако это не означает, что мы не можем пофантазировать о расширении этой базовой функциональности; например, мы сейчас предполагаем, что каждое поле в документе имеет одинаковый вклад в релевантность, тогда как совпадение слова запроса в заголовке, вероятно, должно иметь больший вес, чем совпадение в описании. Еще одним интересным проектом может стать сложный парсинг запросов; нет причин, по которым должны совпадать все термины или только один термин. Почему бы не исключить определенные термины, или позволить AND и OR между определенными словами? Можем ли мы сохранить индекс на диске и сделать так, чтобы он не ограничивался оперативной памятью ноутбука?

Материал подготовлен с ❤️ редакцией Кухни IT.