ТЕХНИЧЕСКОЕ ЗАДАНИЕ: Реализовать BM25 с нуля

1. Цель задачи

Разработать собственную реализацию ранжирующей функции BM25 (Okapi BM25) на основе инвертированного индекса, вычисления TF-IDF и нормализации длины документа. Реализация должна быть выполнена с использованием только стандартной библиотеки Python и базовых математических операций без готовых IR-фреймворков. Финальное качество retrieval (по метрикам nDCG@10, MAP, Recall@10) должно быть не ниже, чем у эталонной реализации Elasticsearch на том же корпусе и наборе запросов.

Ключевой результат Рабочий модуль bm25.py, который принимает коллекцию документов и запрос, возвращает ранжированный список документов с баллами BM25, с качеством, эквивалентным Elasticsearch (разница метрик ≤ 2%).

2. Исходные данные

Если нет реального Elasticsearch — симулируем:

1. Установить библиотеку rank_bm25 (pip install rank_bm25) и реализовать эталонный scorer на ней.
2. Использовать ту же коллекцию и запросы — сравнить свои баллы с эталонными (допустимое расхождение из-за параметров k1, b — фиксируем их одинаковыми).

3. Технологический стек

4. Этапы выполнения

Этап 1: Подготовка данных и разметка (1 час)

Действия

1. Выбрать корпус
   Рекомендуется fetch_20newsgroups(subset='train') из sklearn.datasets — первые 500 документов (или произвольная выборка). Сохранить список documents (строки) и doc_ids.

2. Сформировать запросы и релевантные документы

   • Для каждого запроса указать список релевантных doc_id (можно вручную: 5–10 запросов, 3–5 релевантных каждый).
   • Формат: queries.json (Query ID → текст) и qrels.json (Query ID → список doc_id).
   • Пример queries.json:

     {
       "q1": "space exploration rockets nasa",
       "q2": "machine learning neural networks"
     }
     

   • Пример qrels.json:

     {
       "q1": [0, 12, 45],
       "q2": [3, 17, 88, 99]
     }
     

3. Выполнить предобработку текста

   • Нижний регистр
   • Удалить пунктуацию (кроме апострофов)
   • Токенизация (nltk.word_tokenize)
   • Стемминг (PorterStemmer)
   • Удалить стоп-слова (nltk.corpus.stopwords.words('english'))
   • Итог: список токенов для каждого документа и запроса

4. Разделить данные на train/validation (если потребуется настройка k1, b).
   Можно использовать 80% корпуса для построения индекса, 20% для тестирования или весь корпус.

Ожидаемый результат этапа

• Файлы corpus.json (doc_id → токенизированный список слов), queries.json, qrels.json
• Функция preprocess(text: str) -> list[str] в utils.py

Этап 2: Построение инвертированного индекса (1.5 часа)

Действия

1. Определить структуру
   Словарь inverted_index: dict[str, dict[int, int]] — токен → {doc_id: term_frequency}.
   Дополнительно хранить doc_lengths: list[int] — длина каждого документа в словах.

2. Реализовать класс InvertedIndex

   class InvertedIndex:
       def __init__(self):
           self.index: dict[str, dict[int, int]] = {}
           self.doc_lengths: list[int] = []
           self.num_docs: int = 0
           self.avg_doc_length: float = 0.0
   
       def build(self, corpus: dict[int, list[str]]):
           # проход по документам, подсчёт tf и заполнение index
           # после построения вычислить avg_doc_length
           pass
   

3. Построить индекс

   • Для каждого документа увеличить num_docs, записать его длину.
   • Для каждого токена добавить запись {doc_id: tf}. Использовать defaultdict.
   • После всех документов вычислить среднюю длину.

4. Добавить метод doc_frequency(token) -> int — количество документов, содержащих токен (длина self.index[token]).

5. Сериализовать индекс в файл inverted_index.pkl (pickle) для повторного использования.

Ожидаемый результат этапа

• Класс InvertedIndex с построителем и атрибутами.
• Файл inverted_index.pkl
• Тестовый вывод: количество уникальных токенов, средняя длина документа.

Этап 3: Реализация BM25 scoring (2 часа)

Действия

1. Реализовать функцию bm25_score(query_tokens, doc_id, index, k1=1.2, b=0.75):

   • Формула Okapi BM25:

     score(q, d) = sum_{t in q} IDF(t) * (tf(t,d) * (k1+1)) / (tf(t,d) + k1 * (1 - b + b * (|d| / avgdl)))
     

   • IDF(t) = log((N - df(t) + 0.5) / (df(t) + 0.5)) (вариант с +0.5, чтобы избежать деления на ноль)
   • N — общее число документов
   • df(t) — docs frequency (число документов с токеном t)
   • tf(t,d) — term frequency токена t в документе d
   • |d| — длина документа d
   • avgdl — средняя длина документа

2. Обойти все документы, которые содержат хотя бы один токен из запроса (используя index.index). Если запрос – несколько токенов, пересечение не обязательно; можно брать объединение doc_id из всех токенов.

3. Оптимизация

   • Вычислить IDF для каждого токена заранее (словарь idf_cache), обновлять при изменении k1, b не нужно, т.к. IDF не зависит от параметров.
   • Для каждого документа, содержащего токен, вычислить частичные баллы и суммировать.

4. Написать метод score_query(query_tokens, index, k1, b) -> list[(doc_id, score)]:

   • Создать doc_scores = defaultdict(float)
   • Для каждого токена в запросе:
     • Если токен не в индексе – пропустить
     • Вытащить словарь {doc_id: tf}
     • Для каждого (doc_id, tf) вычислить компонент балла и добавить к doc_scores[doc_id]
   • Отсортировать по убыванию балла.

5. Параметры по умолчанию k1=1.2, b=0.75. Можно сделать возможность их настройки через аргументы.

Тестирование на маленьком примере (2 документа, короткий запрос) вручную:

• doc1: "the cat sat on the mat", doc2: "the dog sat on the mat", query: "cat dog"
• Рассчитать вручную, сверить с выводом функции.

Ожидаемый результат этапа

• Файл bm25.py с функцией bm25_score и методом score_query.
• Возможность запустить bm25.score("query text") -> list[(doc_id, float)]

Этап 4: Оценка качества и сравнение с Elasticsearch (1.5 часа)

Действия

1. Подготовить эталонный ранжирование с помощью Elasticsearch:

   • Создать индекс bm25_test с mapping: content – text, analyzer – english (стандартный, со стеммингом).
   • Загрузить документы (использовать helpers.bulk).
   • Для каждого запроса выполнить поиск с bm25 (по умолчанию) и получить top-100 результатов.
   • Сохранить ранги в формате qid Q0 docid rank score tag.

2. Запустить свою реализацию на тех же запросах:

   • Те же параметры k1, b (1.2, 0.75) – чтобы сравнение было корректным.
   • Получить top-100 для каждого запроса, записать в том же формате.

3. Рассчитать метрики nDCG@10, MAP, Recall@10 с помощью pytrec_eval:

   • Сначала на своих результатах, затем на эталонных.
   • Убедиться, что разница не превышает 2%.

4. Проанализировать разницу

   • Если есть расхождения >2%, найти причину: разная обработка стоп-слов, стемминг, параметры IDF, округление.
   • Исправить (например, синхронизировать стемминг).

5. Написать отчёт в формате

   • Метрики своей реализации
   • Метрики эталона
   • Разница в процентах
   • Примеры запросов, где есть отличие (если есть)

Ожидаемый результат этапа

• Файл retrieval_results/own_run.txt и elastic_run.txt
• Файл metrics_report.txt с цифрами

Этап 5: Оптимизация и документирование (1 час)

Действия

1. Оптимизировать скорость

   • Закешировать idf для всех токенов при инициализации.
   • Вместо вычисления doc_length каждый раз, хранить norm_factor = k1 * (1 - b + b * doc_length / avgdl) предвычисленным для каждого документа (вычисляется один раз после построения индекса, меняется только при изменении k1, b).

2. Добавить поддержку нескольких запросов сразу (батчинг).

3. Написать README.md

   • Описание алгоритма
   • Как запустить
   • Пример использования
   • Результаты сравнения с Elasticsearch

4. Оформить код

   • Разделить на модули: preprocessing.py, index.py, bm25.py, evaluate.py.
   • Добавить докстринги, типизацию.

Ожидаемый результат этапа

• Финальный репозиторий (папка проекта) с кодом, данными и отчётом.
• requirements.txt (только сторонние библиотеки для сравнения: nltk, elasticsearch-py, pytrec_eval).

5. Критерии приемки (Definition of Done)

• Инвертированный индекс построен и сериализован, avg_doc_length корректен (разница с прямым расчётом <1%).
• BM25 реализован без использования готовых IR-библиотек (только Python stdlib + math).
• На тестовом наборе из 2 документов и запроса результат совпадает с ручным расчётом (допустимо расхождение 1e-9 из-за float).
• Сравнение с эталоном (Elasticsearch или rank_bm25) на всем корпусе: разница nDCG@10 ≤ 0.02, MAP ≤ 0.02, Recall@10 ≤ 0.02.
• Параметры k1, b можно менять через аргументы функции (по умолчанию 1.2, 0.75).
• Код оформлен в модульной структуре, есть докстринги, README с примерами.
• Метрики собраны и представлены в текстовом файле metrics_report.txt.
• Воспроизводимость: достаточно выполнить python main.py (или аналогичный скрипт) для получения результатов.

6. Ожидаемый результат

Основной артефакт

• bm25.py — класс/функция для ранжирования по BM25.
• index.py — класс InvertedIndex.
• preprocessing.py — функции нормализации текста.
• evaluate.py — скрипт расчёта метрик и сравнения с эталоном.
• metrics_report.txt — отчёт с метриками.
• README.md — документация.

Опционально

• График сравнения Precision–Recall.
• Тесты (pytest) на маленьких данных.
• Ноутбук Jupyter с визуализацией процесса.

7. Возможные сложности и их решение

8. Бюджет времени (оценка)

Примечание: Для первого раза с изучением теории BM25 может потребоваться до 10 часов. Рекомендуется выполнять в спринте 1-2 дня.

9. Связанные вопросы из базы знаний

10. Чек-лист самопроверки

• Я реализовал инвертированный индекс с поддержкой term frequency и document frequency.
• Я реализовал BM25 в строгом соответствии с формулой Okapi BM25 (k1, b, IDF).
• Я протестировал на маленьком примере вручную и получил совпадающий балл.
• Я сравнил свои метрики с эталоном (Elasticsearch или rank_bm25) и убедился, что разница ≤2%.
• Я оформил код в модули, добавил docstring и README с инструкцией по запуску.