Улучшение точности поиска с Rerank

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

Платформа Compressa предоставляет Rerank модель для реализации такого подхода.

Процесс работы может выглядеть следующим образом:

Сначала система берет вопрос пользователя и извлекает топ-100 документов из датасета, используя либо лексический, либо семантический поиск.
Затем эти 100 документов вместе с исходным вопросом передаются в Compressa Rerank.
Для каждого документа вычисляется оценка релевантности.
На основе полученных оценок документы переранжируются, предоставляя пользователю наиболее подходящие результаты.

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

В этом руководстве мы рассмотрим, как использовать Compressa Rerank для улучшения результатов поиска, и продемонстрируем её применение на практических примерах.

Весь представленный ниже код можно получить в виде .ipynb ноутбука здесь.

# Установка необходимых библиотек
!pip install langchain-compressa
!pip install langchain


# Импорт необходимых библиотек
import os
import requests
import numpy as np
from time import time
from typing import List
from pprint import pprint
from langchain_compressa import CompressaRerank
from langchain.schema import Document

Для ячейки ниже вам понадобится API ключ Compressa. Вы можете получить его после регистрации.

os.environ["COMPRESSA_API_KEY"] = "ваш_ключ"

# Инициализация Compressa Rerank
reranker = CompressaRerank(api_key=COMPRESSA_API_KEY, base_url="https://compressa-api.mil-team.ru/v1", top_n=3)

# Пример запроса и "документов"
query = "Кто является автором картины 'Черный квадрат'?"
docs = [
    "Казимир Малевич создал 'Черный квадрат' в 1915 году, что стало революционным событием в мире искусства.",
    "Илья Репин известен своими реалистичными картинами, такими как 'Бурлаки на Волге' и 'Иван Грозный и сын его Иван'.",
    "'Черный квадрат' Малевича считается одним из самых влиятельных произведений абстрактного искусства XX века.",
    "Василий Кандинский, современник Малевича, также работал в жанре абстрактного искусства.",
    "Малевич написал несколько версий 'Черного квадрата' между 1915 и 1930 годами.",
    "Картина 'Черный квадрат' впервые была выставлена на футуристической выставке '0,10' в Петрограде.",
    "Помимо 'Черного квадрата', Малевич известен своими работами в стиле супрематизма.",
    "Марк Шагал, еще один известный русский художник, работал в стиле модернизма и сюрреализма.",
    "В Третьяковской галерее в Москве хранится одна из версий 'Черного квадрата' Малевича.",
    "Картина 'Черный квадрат' вызвала множество дискуссий и интерпретаций в мире искусства."
]

# Создаем объекты типа Langchain Document из наших вопросов
doc_objects = [
    Document(page_content=doc, metadata={"index": idx})
    for idx, doc in enumerate(docs)
]

Используем Compressa Rerank

В следующей ячейке мы вызовем Compressa Rerank для ранжирования docs на основе их релевантности запросу query

results = reranker.compress_documents(query=query, documents=doc_objects)

top_n = 3 # Измените top_n, чтобы изменить количество возвращаемых результатов.
sorted_results = sorted(results, key=lambda x: x.metadata['relevance_score'], reverse=True)
selected_results = sorted_results[:top_n]

print(query)

for idx, r in enumerate(selected_results):
    doc_index = r.metadata["index"]
    doc_text = r.page_content if doc_index < len(docs) else "Текст не найден"
    print(f"Ранг документа: {idx + 1}, Индекс документа: {doc_index}")
    print(f"Документ: {doc_text}")
    print(f"Оценка релевантности: {r.metadata['relevance_score']:.2f}")
    print("\n")

Поиск по Википедии - End2end демо

В следующем примере мы покажем, как использовать Rerank модель для улучшения поиска по статьям русскоязычной Wikipedia. Для создания удобного датасета мы уменьшили mrtydi-v1.1-russian до порядка 500к отрывков из статей, сохранив "правильные" ответы для всех вопросов, которые будем использовать ниже, и добавили другие рандомные статьи.

Мы будем использовать лексический поиск BM25 для извлечения топ-100 отрывков, соответствующих запросу, а затем отправим эти 100 отрывков и запрос в Compressa Rerank, чтобы получить отранжированные результаты.

Мы выведем топ-3 результата согласно BM25 (как это используется, например, в Elasticsearch) и улучшенный результат от нашей модели Compressa Rerank.

Для начала установим и импортируем дополнительные библиотеки.

!pip install -U  rank_bm25
!pip install gdown

import gdown
import json
import gzip
import os
from rank_bm25 import BM25Okapi
from sklearn.feature_extraction import _stop_words
import string
from tqdm.autonotebook import tqdm

#Скачиваем датасет c google диска Сompressa, он уже нарезан на отрывки
file_id = '1u-3EQ4yoYys1wSSK_EifIPPRyh_6oaR6'
url = f'https://drive.google.com/uc?id={file_id}'
gdown.download(url, 'mrtydi-russian-subset.jsonl', quiet=False)

with open('mrtydi-russian-subset.jsonl', 'r', encoding='utf-8') as f:
    passages = [json.loads(line) for line in tqdm(f)]

print(f"Всего отрывков в подмножестве: {len(passages)}")
print(f"Пример первого отрывка:")
print(f"ID: {passages[0]['id']}")
print(f"Содержание: {passages[0]['contents'][:300]}...")  # Показываем первые 200 символов содержания

# Мы будем сравнивать результаты с обычным лексическим поиском (поиском по ключевым словам). 
# Здесь мы используем алгоритм BM25, который реализован в пакете rank_bm25.
# Мы приводим наш текст к нижнему регистру и удаляем стоп-слова при индексации

# Для начала зададим список русских стоп-слов, чтобы улучшить процесс токенизации. Список можно расширить или взять готовый, он приведен для примера.
RUSSIAN_STOP_WORDS = set([
    'и', 'в', 'во', 'не', 'что', 'он', 'на', 'я', 'с', 'со', 'как', 'а', 'то', 'все', 'она', 'так',
    'его', 'но', 'да', 'ты', 'к', 'у', 'же', 'вы', 'за', 'бы', 'по', 'только', 'ее', 'мне', 'было',
    'вот', 'от', 'меня', 'еще', 'нет', 'о', 'из', 'ему', 'теперь', 'когда', 'даже', 'ну', 'вдруг',
    'ли', 'если', 'уже', 'или', 'ни', 'быть', 'был', 'него', 'до', 'вас', 'нибудь', 'опять', 'уж',
    'вам', 'ведь', 'там', 'потом', 'себя', 'ничего', 'ей', 'может', 'они', 'тут', 'где', 'есть',
    'надо', 'ней', 'для', 'мы', 'тебя', 'их', 'чем', 'была', 'сам', 'чтоб', 'без', 'будто', 'чего',
    'раз', 'тоже', 'себе', 'под', 'будет', 'ж', 'тогда', 'кто', 'этот', 'того', 'потому', 'этого',
    'какой', 'совсем', 'ним', 'здесь', 'этом', 'один', 'почти', 'мой', 'тем', 'чтобы', 'нее', 'сейчас',
    'были', 'куда', 'зачем', 'всех', 'никогда', 'можно', 'при', 'наконец', 'два', 'об', 'другой',
    'хоть', 'после', 'над', 'больше', 'тот', 'через', 'эти', 'нас', 'про', 'всего', 'них', 'какая',
    'много', 'разве', 'три', 'эту', 'моя', 'впрочем', 'хорошо', 'свою', 'этой', 'перед', 'иногда',
    'лучше', 'чуть', 'том', 'нельзя', 'такой', 'им', 'более', 'всегда', 'конечно', 'всю', 'между'
])

def bm25_tokenizer(text):
    tokenized_doc = []
    for token in text.lower().split():
        token = token.strip(string.punctuation)
        if len(token) > 0 and token not in RUSSIAN_STOP_WORDS:
            tokenized_doc.append(token)
    return tokenized_doc

tokenized_corpus = []
for passage in tqdm(passages):
    tokenized_corpus.append(bm25_tokenizer(passage['contents']))

bm25 = BM25Okapi(tokenized_corpus)

# Эта функция будет искать во всех отрывках из Википедии те, которые отвечают на запрос. 
# Затем мы выполняем переранжирование, используя Compressa Rerank

def search(query, top_k=3, num_candidates=100):
    print("Наш вопрос:", query)

    # Поиск BM25, то есть поиск по ключевым словам 
    bm25_scores = bm25.get_scores(bm25_tokenizer(query))
    top_n = np.argpartition(bm25_scores, -num_candidates)[-num_candidates:]
    bm25_hits = [{'corpus_id': idx, 'score': bm25_scores[idx]} for idx in top_n]
    bm25_hits = sorted(bm25_hits, key=lambda x: x['score'], reverse=True)
    
    print(f"Топ-3 результата поиска BM25 по ключевым словам")
    for hit in bm25_hits[0:top_k]:
        print("\t{:.3f}\t{}".format(hit['score'], passages[hit['corpus_id']]['contents'].replace("\n", " ")))

    # Превратим отрывки в документы для переранжирования
    docs = [
        Document(page_content=passages[hit['corpus_id']]['contents'], metadata={"corpus_id": hit['corpus_id']})
        for hit in bm25_hits
    ]
    
    # Реранжируем документы
    reranked_results = reranker.compress_documents(query=query, documents=docs)
    sorted_results = sorted(reranked_results, key=lambda x: x.metadata["relevance_score"], reverse=True)

    print(f"\nТоп-3 результата после переранжирования ({len(bm25_hits)} результатов BM25 переранжированы)")
    for hit in sorted_results[:top_k]:
        print("\t{:.3f}\t{}".format(hit.metadata['relevance_score'], hit.page_content.replace("\n", " ")))

search(query = "В каком году была создана компания Apple?")

search(query = "Где были найдены останки человека флоресского в 2003 году?")

search(query = "Когда началось Возрождение?")

search(query = "Сколько людей живет в Мехико в 2019 году?")

search(query="Какой язык является самым древним?")

search(query="Какой климат в Бурятии?")

search(query="На какой реке расположен Тольятти?")

Улучшение точности поиска с Rerank

Используем Compressa Rerank​

Поиск по Википедии - End2end демо​

Используем Compressa Rerank

Поиск по Википедии - End2end демо