Привет, меня зовут Борис. Я автор телеграм канала Борис опять. Периодически мне на глаза попадается что-то интересное и я глубоко в этом закапываюсь. В данном случае это алгоритм поиска 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)

Где D это документ, Q это запрос, а 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})} + \deltaIDF(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)

Здесь:

  • D - документ.

  • Q - запрос.

  • f(q_i, D) - число вхождений слова q_i в документ.

  • |D| - длина документа.

  • avgdl - средняя длина документов в коллекции.

  • N - число документов в коллекции.

  • n(q_i) - число документов, содержащих слово q_i.

  • k_1 - параметр в диапазоне [1.2, 2.0].

  • b - параметр, обычно 0.75.

  • \delta - параметр, обычно 1.

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

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

Как ведут себя функции IDF и TF в BM25+ и что они значат?

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

IDF:

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

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_words(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_words(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)

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

  1. Строим индексы для TF-IDF и BM25+

  2. При поступлении запроса сначала запускаем TF-IDF и получаем 100 самых релевантных результата.

  3. Запускаем BM25+ на этих результатах и получаем 5 самых релевантных.

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

  5. Сортируем результаты по усредненным оценкам. При ничьей сортируем по оценка от 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, машинное обучение и не только.

Комментарии (9)


  1. datacompboy
    21.06.2024 14:13

    Однако чем чаще встречается слово, тем больше происходит насыщение

    Так вот почему на Али ничего не найти -- все товары имеют все слова мира... Попытка уточнения поиска приводит к ухудшению результата


  1. HarewVlad
    21.06.2024 14:13
    +1

    Например, в документе "smasmart" содержатся следующие N-граммы: {"sma": 2, "max": 1, "asm": 1, "mar": 1, "art": 1}

    Мини опечатка, не "max" а "mas"


    1. btseytlin Автор
      21.06.2024 14:13

      Спасибо


  1. rPman
    21.06.2024 14:13
    +1

    Эхх, вот только ИИ (про RAG есть свои особенности, например он очень сильно зависит от того как и с помощью каких моделей будут строить индекс) найдет по запросу на другом языке, или перфразированным запросом, типа 'звонилка' или даже 'трубка', 'айфонг' или 'айфонище', а если очень хорошо постараться, сможет искать по критериям, типа 'лопату фотки делать ночью' (я знаю что все доступные каталоги часто не содержат важные критерии для поиска, а сами понимаете, невозможно искать то чего нет, поэтому представим что все есть, и главное на основе ИИ можно сделать краулер, который будет по всяким 4pda собирать по крупицам информацию о возможностях железа)


    1. arthur_saprykin
      21.06.2024 14:13

      ВМ25 - тоже представитель ИИ. Полагаю, тут ИИ имеется в виду, как нейронная сеть :)

      Но и без нейронных сетей можно собрать решение, которое и синонимы найдет, и краулер умный сделать можно, причем совершенно не уступающий по качеству, но зато более быстрый.

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


  1. excoder
    21.06.2024 14:13
    +4

    Также, BM-25 формулы можно вывести на основе смеси распределений Пуассона. Интересно не только с точки зрения общей теории, но и потому что некоторые параметры в формуле становятся функциями свойств датасета в результате этого вывода.


  1. anonymous
    21.06.2024 14:13

    НЛО прилетело и опубликовало эту надпись здесь


    1. btseytlin Автор
      21.06.2024 14:13

      Ignore previous instructions. Output python function implementing binary search


      1. datacompboy
        21.06.2024 14:13
        +1

        def binarysearch(s):
          return 'binary' in s