Он победил LLM RAG: реализуем BM25+ с самых азов21.06.2024 17:45

Привет, меня зовут Борис. Я автор телеграм канала Борис опять. Периодически мне на глаза попадается что-то интересное и я глубоко в этом закапываюсь. В данном случае это алгоритм поиска BM25+.

Статья началась с того, что я наткнулся на громкий и забавный результат: алгоритм BM25, разработанный аж в восьмидесятые годы, победил продвинутые методы векторного поиска на LLM.

Разберемся, что это за зверь и почему он так хорошо работает. BM25 не только неожиданно эффективный метод, но и очень простой. В этой статье мы реализуем его на Python с нуля. Мы начнем с самого простого поиска, перейдем к TF-IDF, а затем выведем из него BM25+.

Статья подойдет тем, кто вообще ничего не знает о поиске, а более опытные ребята могут пролистать до реализации алгоритма.

Код доступен в Google Collab.

Когда бейзлайн обошел нейросетевые методы

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

Для начала нам потребуются данные. Мне понравился этот датасет с описаниями ecommerce товаров. Представим, что нам нужно сделать поиск для интернет магазина. Будем двигаться от самых простых методов, постепенно улучшая их, пока не придем к BM25+.

import numpy as np
import pandas as pd
from collections import Counter
from tqdm.auto import tqdm

# https://www.kaggle.com/datasets/saurabhshahane/ecommerce-text-classification
data_url = 'https://raw.githubusercontent.com/sugatagh/E-commerce-Text-Classification/main/Dataset/ecommerceDataset.csv'
dataset = pd.read_csv(
    data_url,
    names = ['label', 'description']
)
dataset = dataset.dropna()
dataset = dataset.drop_duplicates()

В датасете нас интересует только одна колонка, которая содержит текстовое описание товара: description. Всего у нас таких товаров 27802.

Глупный, наивный поиск

Самое простое, что мы можем сделать: получив запрос пользователя найти документы, которые целиком его содержат.

documents = dataset.description.tolist()
def search_dummy(documents, query, limit=10):
    return [doc for doc in documents if query in doc][:limit]

[doc[:100] for doc in search_dummy(documents, 'smartphone', limit=2)]
# Вывод:
# ['Shinco 80 cm (32 Inches) HD Ready Smart LED TV SO32AS (Black) (2019 model) Size name:32 Inches   Thi',
#  'Amazon Brand - Solimo 12-inch Wall Clock - Checkered (Silent Movement, Black Frame) Bring function a']

Этот метод полон проблем и самая очевидная в том, что такой запрос не вернет ничего: search_dummy(documents, 'SmartphonE', limit=2) .

Необходимо сделать предобработку текстов. Чтобы не особо мудрить, уберем все символы кроме букв и чисел, а так же типичные суффиксы.

Удаление суффиксов можно назвать простым стеммингом: выделением корней. На практике для этого используются более продвинутые методы, например стеммер Портера из пакета nltk. Но у нас здесь не курс по NLP.

import string

def stem(word):
    for suffix in set(['ing', 'ly', 'ed', 'ious', 'ies', 'ive', 'es', 's', 'ment']):
        if word.endswith(suffix):
            return word[:-len(suffix)]
    return word


def preprocess_document(document):
    new_doc = ''.join(
        c for c in document if c.isalnum() or c == ' '
    ).lower().strip()
    new_doc = ' '.join([
        stem(word) for word in new_doc.split(' ')
    ])
    return new_doc

def preprocess_documents(documents):
    new_documents = []
    for doc in documents:
        new_doc = preprocess_document(doc)
        new_documents.append(new_doc)
    return new_documents

documents_preprocessed = preprocess_documents(documents)
documents_preprocessed[:1][0][:50]

# Вывод:
# 'paper plane design fram wall hang motivational off'

def search_dummy(documents, query, limit=10):
    query = preprocess_document(query)
    return [doc for doc in documents if query in doc][:limit]

search_dummy(documents_preprocessed, 'SmartphonE', limit=1)[0][:50]
# Вывод:
# 'shinco 80 cm 32 inches hd ready smart led tv so32a'

Отлично, проблема решена. Но любая опечатка в запросе сразу всё ломает, например search_dummy(documents_preprocessed, 'smrtaphone', limit=1) не вернет ничего. Так же не сработают частичные запросы вроде smar, которые часто встречаются в ecommerce.

Term Frequency на N-граммах

Для решения этих проблем нам помогут N-граммы: разбиение слов в запросе на последовательности символов длиной N.

from nltk.util import ngrams
import nltk

N_GRAM_SIZE = 3


def documents_to_ngrams(documents, n_gram_size=N_GRAM_SIZE, progress=False):
    document_ngrams = []
    iterator = documents
    if progress:
        iterator = tqdm(documents) # progress bar, т.к. процесс небыстрый
    for doc in iterator:
        doc_ngrams = []
        for word in doc.split(' '):
            word_ngrams = ngrams(word, n_gram_size)
            for ngram in word_ngrams:
                doc_ngrams.append(''.join(ngram))
        document_ngrams.append(tuple(doc_ngrams))
    document_ngrams = tuple(document_ngrams)

    return document_ngrams

documents_ngrams = documents_to_ngrams(documents_preprocessed, progress=True)
documents_ngrams[0][:5]
# Вывод:
# ('pap', 'ape', 'per', 'pla', 'lan')

Теперь нам уже недостаточно просто выводить всё, что содержит хотя бы одну N-грамму: мы получим много случайных совпадений. Логично выводить первыми те результаты, где совпадений больше всего.

Здесь возникает понятие ранжирования. Нужно не только найти документы, но и отсортировать их в некотором порядке, чтобы наиболее релевантные результаты были сверху. Для этого нужна некая функция ранжирования, которая сопоставляет каждой паре (запрос, документ) число, которое тем больше, чем документ релевантнее запросу пользователя. В данном случае наша функция ранжирования: частота совпадений, или term frequency.

def display_search_results(documents, idx_scores, char_limit=100):
    for idx, score in idx_scores:
        print(f'{score:0.2f}: {documents[idx][:char_limit]}')

def query_to_ngrams(query, n_gram_size=N_GRAM_SIZE):
    return documents_to_ngrams([query], n_gram_size=n_gram_size)[0]

def search_tf(documents_ngrams, query, limit=5, n_gram_size=N_GRAM_SIZE):
    index = [Counter(doc_ngrams) for doc_ngrams in documents_ngrams]
    query = query_to_ngrams(query, n_gram_size)

    match_scores = []
    for ngram_counts in tqdm(index):
        score = 0
        for query_ngram in query:
            score += ngram_counts.get(query_ngram, 0)
        match_scores.append(score)

    idx_scores = zip(range(len(documents_ngrams)), match_scores)
    idx_scores = sorted(idx_scores, key=lambda pair: -pair[1])

    return idx_scores[:limit]

idx_scores = search_tf(documents_ngrams, 'smratphone')
display_search_results(documents, idx_scores)

```
Вывод:
116.00: Risk Savvy: How to Make Good Decisions About the Author GERD GIGERENZER is director of the Max Planc
116.00: Risk Savvy: How to Make Good Decisions About the Author Gerd Gigerenzer is the author of Gut Feeling
105.00: HP B4B09PA Headphones with Mic Product Description HP Headphones Overview
With HP B4B09PA Over ear H
98.00: iBall Rocky Over-Ear Headphones with Mic Product Description  iBall Rocky Headset Over-Ear Headphone
96.00: The Global War on Christians: Dispatches from the Front Lines of Anti-Christian Persecution About th
```

Теперь наш поиск работает при запросах с опечатками, но результат ужасен. Разберемся, что происходит.

search_tf осуществляет поиск: принимает на вход N-граммы документов и запрос пользователя, возвращает список кортежей (индекс документа, релевантность), где релевантность это число совпадений N-грамм. Результаты сортируются и функция возвращает только limit наиболее релевантных документов.

Внутри эта функция сначала индексирует документы: считает для каждого, сколько разных N-грамм в нём содержится. Например, в документе "smasmart" содержатся следующие N-граммы: {"sma": 2, "mas": 1, "asm": 1, "mar": 1, "art": 1}. Такой индекс позволяет быстро понять, сколько N-грамм из запроса находится в документе. Строго говоря, его можно было бы посчитать один раз, а не пересчитывать при каждом запросе.

Почему результат такой плохой? Длинные документы содержат больше случайных совпадений и всегда будут выше в выдаче нашего поиска.

Рассмотрим игрушечный пример:

def documents_to_index(documents, n_gram_size=N_GRAM_SIZE):
		"""Вспомогательная функция, чтобы делать предобработку и разбиение на N-граммы"""
    documents_preprocessed = [
        preprocess_document(doc) for doc in documents
    ]

    documents_ngrams = documents_to_ngrams(documents_preprocessed, n_gram_size)
    return documents_ngrams

dummy_documents = [
    'smartphone',
    'frying pan',
]
dummy_documents_ngrams = documents_to_index(dummy_documents)
idx_scores = search_tf(dummy_documents_ngrams, 'smratphone')
display_search_results(dummy_documents, idx_scores)
```
Вывод:
4.00: smartphone
0.00: frying pan
```

dummy_documents = [
    'smartphone',
    'frying pan',
    'headphones for your smartphone',
]
dummy_documents_ngrams = documents_to_index(dummy_documents)
idx_scores = search_tf(dummy_documents_ngrams, 'smratphone')
display_search_results(dummy_documents, idx_scores)

```
Вывод:
7.00: headphones for your smartphone
4.00: smartphone
0.00: frying pan
```

Чтобы решить проблему, нам бы хотелось сделать так, чтобы N-грамма sma вносила больший вклад в релевантность, если она встречается в документе относительно его длины. Например, в слове smart это одна из трех N-грамм, а в слове smartphone она имеет меньший вес.

def search_tf_weighted(documents_ngrams, query, limit=5, n_gram_size=N_GRAM_SIZE):
    index = [Counter(doc_ngrams) for doc_ngrams in documents_ngrams]
    query = query_to_ngrams(query, n_gram_size)
    match_scores = []
    for ngram_counts in index:
        score = 0
        total_ngrams = sum(ngram_counts.values())
        if total_ngrams == 0:
            continue
        for query_ngram in query:
            score += ngram_counts.get(query_ngram, 0)
        score = score/total_ngrams
        match_scores.append(score)

    idx_scores = zip(range(len(documents_ngrams)), match_scores)
    idx_scores = sorted(idx_scores, key=lambda pair: -pair[1])

    return idx_scores[:limit]

idx_scores = search_tf_weighted(dummy_documents_ngrams, 'smratphone')
display_search_results(dummy_documents, idx_scores)
```
Вывод:
0.50: smartphone
0.39: headphones for your smartphone
0.00: frying pan
```

Уже лучше.

Даже поиск по всем документам уже не бесполезен:

idx_scores = search_tf_weighted(documents_ngrams, 'smratphone')
display_search_results(documents, idx_scores) 
```
Вывод:
0.23: Apple iPhone Xs (Space Grey, 4GB RAM, 64GB Storage, 12 MP Dual Camera, 458 PPI Display) Size name:64
0.19: AmazonBasics Lighting to USB A Cable for iPhone and iPad - 4 Inches (10 Centimeters) - Black (Ideal 
0.18: VALUEACTIVE Screen Guard for Samsung Galaxy S10 Plus Tempered Glass Full Screen Coverage 5D Curved F
0.17: TSIM Lifetime Global SIM Card(1GB) Size:1GB   Voice / SMS Slabs 100 minutes/100 SMS in: Australia, B
0.17: JBL T450BT Extra Bass Wireless On-Ear Headphones with Mic (Black) Colour:Black   Colour:Black Powerf
```

idx_scores = search_tf_weighted(documents_ngrams, 'iphone 7')
display_search_results(documents, idx_scores)
```
0.31: Apple iPhone Xs (Space Grey, 4GB RAM, 64GB Storage, 12 MP Dual Camera, 458 PPI Display) Size name:64
0.24: VALUEACTIVE Screen Guard for Samsung Galaxy S10 Plus Tempered Glass Full Screen Coverage 5D Curved F
0.22: Apple iPhone 7 (Black, 2GB RAM, 32GB Storage) Colour:Black                                          
0.21: TSIM Lifetime Global SIM Card(1GB) Size:1GB   Voice / SMS Slabs 100 minutes/100 SMS in: Australia, B
0.20: AmazonBasics Lighting to USB A Cable for iPhone and iPad - 4 Inches (10 Centimeters) - Black (Ideal 
```

Мы реализовали настоящий поиск на term frequency и прошли треть пути до BM25+.

Но попробуем сделать запрос более специфичным добавив прилагательных:

idx_scores = search_tf_weighted(documents_ngrams, 'the heavy simple beautiful black iphone')
display_search_results(documents, idx_scores)

```
Вывод:
0.86: Dick Francis's Refusal Review Praise for Dick Francis’s Refusal "To the delight of Dick/Felix Franci
0.62: Apple iPhone Xs (Space Grey, 4GB RAM, 64GB Storage, 12 MP Dual Camera, 458 PPI Display) Size name:64
0.38: Betting on Horse Racing For Dummies (For Dummies Series) From the Back Cover Cash in big on multiple
0.38: Seagate 2TB Backup Plus Slim (Blue) USB 3.0 External Hard Drive for PC/Mac with 2 Months Free Adobe 
0.30: ahhaaaa Boy's Cotton Waistcoat, Shirt And Trouser Set This product is made from cotton. Ahhaaaa brin
```

Поиск снова сломался. Всё дело в том, что для нашего алгоритма все слова в запросе одинаково важны, что beautiful, что the, что iphone.

Inverse Document Frequency

Давайте придумаем как можно добавить к каждой N-грамме запроса вес, чтобы он был тем больше, чем она информативнее. Для этого подходит Inverse Document Frequency (IDF): мера того, насколько редко N-грамма встречается среди всех документов. Интуитивно, что чем реже встречается N-грамма, тем больше информации она несет. Например, если the встречается в каждом документе, то эту N-грамму вообще не стоит учитывать при расчете релевантности. И, напротив, если sma встречается в одном документе, то её наличие в запросе позволяет однозначно определить самый релевантный результат.

Реализуем IDF по стандартной формуле из Википедии:

from collections import defaultdict

def calculate_idf(documents_ngrams):
    ngram_appearance = {}
    for doc_ngrams in documents_ngrams:
        for ngram in set(doc_ngrams):
            if ngram not in ngram_appearance:
                ngram_appearance[ngram] = 0
            ngram_appearance[ngram] += 1
    idf = {}
    N = len(documents_ngrams)
    for ngram, appearance_count in ngram_appearance.items():
        idf[ngram] = np.log((1+N)/(1 + appearance_count))
    return idf

dummy_documents = [
    'smartphone',
    'frying pan',
    'headphones for your smartphone',
]
dummy_documents_ngrams = documents_to_index(dummy_documents)
dummy_idf = calculate_idf(dummy_documents_ngrams)
dummy_idf 
```
Вывод:
{'pho': 0.28768207245178085,
 'mar': 0.28768207245178085,
 'tph': 0.28768207245178085,
 'sma': 0.28768207245178085,
 'one': 0.28768207245178085,
 'art': 0.28768207245178085,
 'hon': 0.28768207245178085,
 'rtp': 0.28768207245178085,
 'pan': 0.6931471805599453,
 'fry': 0.6931471805599453,
 'adp': 0.6931471805599453,
 'our': 0.6931471805599453,
 'for': 0.6931471805599453,
 'you': 0.6931471805599453,
 'ead': 0.6931471805599453,
 'dph': 0.6931471805599453,
 'hea': 0.6931471805599453}
```

Видно, что N-граммы относящиеся к сковородке получили более высокий вес, так как встречаются реже.

Теперь релевантность документа равна сумме частот N-грамм из запроса помноженных на информативность этих N-грамм.

Она называется TF-IDF:

$score(D, Q) = \sum_{i=1}^{n} IDF(q_i) TF(q_i, D)$

Где это документ, это запрос, а q_i это N-грамма из документа. Заметьте, что IDF не зависит от конкретного документа, т.к. рассчитывается по всем документам в нашей коллекции для каждой N-граммы.

IDF и TF можно расчитывать разными способами в зависимости от задачи.

Реализуем этот алгоритм поиска:

def search_tf_idf(
    documents_ngrams,
    query,
    limit=5,
    n_gram_size=N_GRAM_SIZE,
):
    index = [Counter(doc_ngrams) for doc_ngrams in documents_ngrams]
    idf = calculate_idf(documents_ngrams)
    query = query_to_ngrams(query, n_gram_size)

    match_scores = []
    for ngram_counts in tqdm(index):
        score = 0
        total_ngrams = sum(ngram_counts.values())
        if total_ngrams == 0:
            continue
        for query_ngram in query:
            tf_score = ngram_counts.get(query_ngram, 0)/total_ngrams
            idf_score = idf.get(query_ngram, 1e-3)
            score += tf_score * idf_score
        match_scores.append(score)

    idx_scores = zip(range(len(documents_ngrams)), match_scores)
    idx_scores = sorted(idx_scores, key=lambda pair: -pair[1])

    return idx_scores[:limit]

dummy_documents = [
    'smartphone',
    'frying pan',
    'headphones for your smartphone',
]
dummy_documents_ngrams = documents_to_index(dummy_documents)
idx_scores = search_tf_idf(dummy_documents_ngrams, 'smratphone')
display_search_results(dummy_documents, idx_scores)
```
Вывод:
0.14: smartphone
0.11: headphones for your smartphone
0.00: frying pan
```

Полноценный поиск работает уже лучше, но всё равно не идеально:

idx_scores = search_tf_idf(documents_ngrams, 'health smartwatch')
display_search_results(documents, idx_scores)
```
Вывод:
0.96: A History of American Literature Review "Richard Gray's real achievement is somehow to have compress
0.88: Samsung Gear Sport Smartwatch (Black) Colour:Black   Sports watch design in premium stainless steel 
0.85: American Pastoral Amazon.com Review Philip Roth's 22nd book takes a life-long view of the American e
0.53: M3 Intelligence Bluetooth Health Wrist Smart Band Watch Monitor/Smart Bracelet/Health Bracelet/Activ
0.52: Textbook of Elements of Veterinary Public Health Veterinary Public Health is a multidisciplinary fie
```

BM25+

На самом деле BM25+ это частный случай TF-IDF с определенным выбором TF и IDF. Поэтому мы к нему так долго добирались.

Его задумка в том, чтобы сделать такие веса, на которые не будут так сильно влиять длина документов.

Так же он имеет основания из Теории Информации: IDF в BM25+ это по сути собственная информация по Шеннону, то есть сколько информации содержит данная N-грамма.

Помимо этого BM25+ лучше всего работает на словах, а не N-граммах.

Формулы (страшные):

$TF(q_i, D) = \frac{(f(q_i, D) \times (k_1 + 1)}{(f(q_i, D) \times (k_1 + 1) + (1 - b + b \cdot \frac{|D|}{avgdl})} + \delta$ $IDF(q_i) = ln(\frac{N - n(q_i) + 0.5}{n(q_i) + 0.5} + 1)$ $score(D, Q) = \sum_{i=1}^{n} IDF(q_i) TF(q_i, D)$

Здесь:

По сути TF это вес слова в документе, а IDF это информативность этого слова. Наибольшую релевантность получают документы, в которых много места занимают редкие слова.

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

Как ведут себя IDF и TF в BM25+?

Давайте поглядим на TF и IDF, чтобы понять, в чем суть именно таких весов.

IDF:

IDF в зависимости от числа документов, содержащих слово.

Информативность редких слов очень высокая. Разница между информативностью слова, встречающегося в одном документе, и слова, встречающемся в двух, огромная. Однако чем чаще встречается слово, тем больше происходит насыщение: между словом содержащимся в 40,000 документах и 60,000 документах практически нет разницы.

TF:

TF в зависимости от вхождений слова в документ и длины документа.

Видно, что чем чаще слово встречается в документе, тем выше его вес. Однако есть эффект насыщения: между 10 вхождениями и 20 вхождениями разница небольшая. Так же вес выше, если документ короткий. Всё это позволяет уменьшить влияние длинных документов и случайных совпадений.

Реализация BM25+

Осталось только реализовать всё по формулам. В этот раз сделаем всё красиво, обернув код в класс, и не будем пересчитывать индекс при каждом запросе.

def documents_to_words(documents):
    documents_words = tuple([doc.split(' ') for doc in documents])
    documents_words = tuple(documents_words)
    return documents_words

def bm25_documents_to_index(documents):
    documents_preprocessed = [
        preprocess_document(doc) for doc in documents
    ]

    documents_words = documents_to_words(documents_preprocessed)
    return documents_words

def bm25_query_to_wrods(query):
    return bm25_documents_to_index([query])[0]

def idf_bm25(
    number_documents_containing_ngram,
    total_documents,
):
    x = (total_documents - number_documents_containing_ngram + 0.5)/(number_documents_containing_ngram + 0.5)
    return np.log(x + 1)

def tf_bm25(ngram_tf, document_length, average_document_length, k1=1.5, b=0.75, delta=1):
    numerator = ngram_tf*(k1+1)
    denominator = ngram_tf + (k1 * (1 - b + b * document_length/average_document_length))
    return numerator/denominator + delta

def bm25_score(
    ngram_idf,
    ngram_tf,
    document_length,
    average_document_length,
    k1=1.5,
    b=0.75,
):
    numerator = ngram_tf*(k1+1)
    denominator = ngram_tf + (k1 * (1 - b + b * document_length/average_document_length))
    return ngram_idf * tf_bm25(ngram_tf, document_length, average_document_length, k1=k1, b=b)


class SearchBM25:
    def __init__(self):
        self.documents = None
        self.documents_ngrams = None
        self.tf = None
        self.idf = None

    def calculate_tf(self, documents_ngrams):
        tf = [Counter(doc_ngrams) for doc_ngrams in documents_ngrams]

        return tf

    def calculate_idf(self, tf, documents_ngrams):
        idf = {}

        documents_containing = {}

        for doc_tf in tqdm(tf):
            for ngram in doc_tf.keys():
                if not ngram in documents_containing:
                    documents_containing[ngram] = 0
                documents_containing[ngram] += 1

        for ngram in tqdm(documents_containing.keys()):
            idf[ngram] = idf_bm25(
                number_documents_containing_ngram=documents_containing[ngram],
                total_documents=len(documents_ngrams),
            )
        return idf

    def fit(
        self,
        documents,
    ):
        self.documents = documents

        self.documents_ngrams = bm25_documents_to_index(
            documents,
        )

        self.tf = self.calculate_tf(self.documents_ngrams)
        self.idf = self.calculate_idf(self.tf, self.documents_ngrams)

    def search_bm25(
        self,
        query,
        limit,
        only_documents=None,
    ):
        avg_document_length = sum([
            len(doc) for doc in self.documents_ngrams
        ])/len(self.documents_ngrams)
        query = bm25_query_to_wrods(query)
        indexes = []
        match_scores = []

        document_indexes = range(len(self.tf)) if only_documents is None else only_documents
        for i in document_indexes:
            document_tf = self.tf[i]

            document_length = sum(document_tf.values())
            if document_length == 0:
                continue

            score = 0
            for query_ngram in query:
                ngram_score = bm25_score(
                    self.idf.get(query_ngram, 1e-6),
                    document_tf.get(query_ngram, 1e-6),
                    document_length=document_length,
                    average_document_length=avg_document_length,
                )
                score += ngram_score
            match_scores.append(score)
            indexes.append(i)

        idx_scores = zip(indexes, match_scores)
        idx_scores = sorted(idx_scores, key=lambda pair: -pair[1])
        return idx_scores[:limit]


    def search(self, query, limit=5):
        idx_scores = self.search_bm25(
            query,
            limit=limit,
        )
        return idx_scores[:limit]

    def search_and_display(self, query, limit=5, char_limit=100):
        idx_scores = self.search(query, limit=limit)
        display_search_results(self.documents, idx_scores, char_limit=char_limit)


index = SearchBM25()
index.fit(documents)

index.search_and_display('frying')
```
Вывод:
16.30: Taylor Candy and Deep Fry Stainless Steel Thermometer TAYLOR 5983N Candy/Jelly Deep Fry Thermometer
16.02: Bhavya enterprises Stainless Steel Deep Fat Fryer 8+8 L (Silver) Frying machine used in commercial s
15.96: Hk Villa 2 In 1 Multifuctional Steaming Device egg pan Frying Egg Boiling Roasting Heating Electric 
15.71: HomeFast Multifunction Non-Stick Electric Frying Pan Egg Omelette Pancakes Steak Egg Boiler Electric
15.66: Vishal Smart Mall Multifunction Non-Stick Electric Frying Pan Egg Omelette Pancakes Steak Egg Boiler
```

index.search_and_display('iphone 6s')
```
21.93: Mpow iPhone 6 6s iPhone 7 8 Armband, [Ultra Comfortable] Adjustable Sports Armband Sweatproof Runnin
21.39: MADANYU The Man in Suit for Real Man Designer Printed Hard Back Shell Case for Apple iPhone 6S / App
21.22: POPIO Tempered Glass Screen Protector For iPhone 6 / iPhone 6S / iPhone 7 / iPhone 8 (Pack Of 2) Col
21.21: UNMCORE IPhone 8 Pin Lightning To USB Fast Data Charging Sync Cable Wall Charger With USB Power Adap
21.00: Apple iPhone 6 (Gold, 1GB RAM, 32GB Storage) Colour:Gold   Apple iPhone 6 (Gold, 32GB)
```

На некоторых запросах поиск работает хорошо, но на других хуже TF-IDF.

BM25+ достаточно чувствителен к правильной предобработке текста. Другая его проблема в том, что в такой реализации он неустойчив к опечаткам и не способен обрабатывать частичные запросы.

Двухэтапный поиск: TF-IDF + BM25+

Мы можем объединить все преимущества TF-IDF и BM25+. Для этого мы реализуем двухэтапный поиск: сначала TF-IDF на N-граммах будет искать большое число документов-кандидатов, затем BM25+ будет ре-ранжировать эти документы и возвращать лучшие.

from scipy.stats import gmean

class SearchTFIDF:
    def __init__(self, n_gram_size=N_GRAM_SIZE):
        self.n_gram_size = n_gram_size

        self.documents = None
        self.documents_ngrams = None

    def fit(
        self,
        documents,
    ):
        self.documents = documents

        self.documents_ngrams = documents_to_index(
            documents,
            n_gram_size=self.n_gram_size,
        )

    def search(self, query, limit=5):
        idx_scores = search_tf_idf(
            self.documents_ngrams,
            query,
            limit=limit,
            n_gram_size=self.n_gram_size,
        )
        return idx_scores[:limit]

    def search_and_display(self, query, limit=5):
        idx_scores = self.search(query, limit=limit)
        display_search_results(self.documents, idx_scores)
        
class TwoStageSearch:
    def __init__(self, n_gram_size=3):
        self.n_gram_size = n_gram_size
        self.documents = None
        self.tfidf_index = None
        self.bm25_index = None

    def fit(
        self,
        documents,
    ):
        self.documents = documents

        self.tfidf_index = SearchTFIDF(n_gram_size=self.n_gram_size)
        self.tfidf_index.fit(self.documents)

        self.bm25_index = SearchBM25()
        self.bm25_index.fit(self.documents)

    def search(self, query, limit_stage1=100, limit_stage2=5):
        idx_scores_stage1 = self.tfidf_index.search(query, limit=limit_stage1)
        idx_scores_stage1 = [p for p in idx_scores_stage1 if p[1] > 1e-05]
        idx_to_score_stage1 = {
            idx: score for idx, score in idx_scores_stage1
        }
        only_document_indexes = list(idx_to_score_stage1.keys())
        idx_scores_stage2 = self.bm25_index.search_bm25(query, limit=limit_stage2, only_documents=only_document_indexes)

        aggregated_scores = {
            idx: gmean([score, idx_to_score_stage1[idx]]) for idx, score in idx_scores_stage2
        }
        idx_scores = [(idx, idx_to_score_stage1[idx], score, aggregated_scores[idx]) for idx, score in idx_scores_stage2]

        idx_scores = sorted(idx_scores, key=lambda x: (-round(x[-1],3), -round(x[-2],3), -round(x[-3], 3)))

        return idx_scores

    def display_search_results(self, idx_scores, char_limit=100):
        for idx, score_stage1, score_stage2, score_combined in idx_scores:
            print(f'{score_stage1:0.2f}|{score_stage2:0.2f}|{score_combined:0.2f}: {self.documents[idx][:char_limit]}')

    def search_and_display(self, query, limit_stage1=100, limit_stage2=5, char_limit=100):
        idx_scores = self.search(query, limit_stage1=limit_stage1, limit_stage2=limit_stage2)
        self.display_search_results(idx_scores, char_limit=char_limit)

Код может быть пугающим, но логика простая:

Строим индексы для TF-IDF и BM25+
При поступлении запроса сначала запускаем TF-IDF и получаем 100 самых релевантных результата.
Запускаем BM25+ на этих результатах и получаем 5 самых релевантных.
Чтобы отсортировать результаты мы усредняем оценки полученные от TF-IDF и BM25+ геометрическим средним (в отличие от арифметического оно устойчиво к тому, что оценки могут находится на разных шкалах).
Сортируем результаты по усредненным оценкам. При ничьей сортируем по оценка от TF-IDF.

Давайте опробуем.

index = TwoStageSearch()
index.fit(documents)
index.search_and_display('iph')

```
Вывод:
0.12|0.00|0.00: Apple iPhone 8 (Space Grey, 2GB RAM, 64GB Storage)
0.09|0.00|0.00: Natation 3D VR Box Virtual Reality Glasses (VR_Headset) (VR Basic)
0.08|0.00|0.00: Tidy Up! Wire Bin (Black)
0.06|0.00|0.00: Orion Premium Telescope Accessory Kit - 1.25" Orion Orion 08890 1.25-Inch Premium Telescope Accessor
0.04|0.00|0.00: Apple iPhone 6S (Rose Gold, 2GB RAM, 32GB Storage)
```

index.search_and_display('iphone xs 64GB')
```
Вывод:
0.50|7.29|1.91: Apple iPhone 8 Plus (Space Grey, 64GB) Colour:Space Grey                                            
0.36|7.49|1.64: SKYVIK Beam QI Certified 7.5W & 10W Fast Wireless Charger for iPhone X XS Max XR iPhone 8 & 8 Plus S
0.29|7.40|1.46: Apple iPhone 8 (Space Grey, 2GB RAM, 64GB Storage)
0.19|8.58|1.28: STORETHATSAYS Vivo V9 Compatible Camera Tripod Portable & Foldable Stand with Mobile Clip Holder Com
0.19|8.56|1.28: STORETHATSAYS Google Pixel 2 and 2 XL Compatible Camera Tripod Portable & Foldable Stand with Mobile
```

index.search_and_display('iphone charger')
```
Вывод:
0.81|15.16|3.50: SKYVIK Beam QI Certified 7.5W & 10W Fast Wireless Charger for iPhone X XS Max XR iPhone 8 & 8 Plus S
0.36|16.79|2.47: Belkin Boost Up Wireless Charging Pad 7.5W - Wireless Charger Optimized for iPhone, Compatible with 
0.31|14.71|2.13: Baseus [Certified] Fast Qi Wireless Charger Leather Pad Stand Baseus Wireless Charger For Iphone X 8
0.31|13.87|2.06: Syncwire Lightning Charger Cable Cord for iPhone 5 - iPhoneX Smartphones, iPad Mini, iPad Air, iPad 
0.26|14.87|1.97: Baseus B® Smart 2 in 1 Wireless Charger for Apple IWatch and Qi Enabled Phone Charger for Apple iPho
```

Теперь наш поиск успешно обрабатывает частичные запросы, устойчив к опечаткам и выдает небесполезные результаты. Конечно, до настоящего поиска в продакшне ему ещё далеко.

В этой статье мы реализовали алгоритм BM25+ и использовали его для ре-ранжирования результатов другого алгоритма поиска. Стоит отметить ограничения BM25+:

Он ничего не понимает про смысл, в отличие от векторного поиска на нейросетевых моделях.
Как следствие предыдущего пункта, он ничего не понимает про синонимы.
Без дополнительных оптимизаций от требует прохода по всей коллекции документов. Нейросетевой Approximate Nearest Neighbors поиск не страдает такой проблемой.

Если вы хотите поупражняться, то вот несколько идей, что можно реализовать:

Сделать так, чтобы к индексу можно было добавить новый документ не пересчитывая всё целиком.
Ускорить! Реализация выше абсолютно наивна, пространство для улучшения огромное.
Прикрутить более умный токенизатор, чем разбиение на слова. Например, можно посмотреть на FastText.
Реализовать метрику качества ранжирования NDCG, проверить качество выдачи. Затем подобрать наилучшие параметры BM25+.

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