BM25Retriever под капотом

В современных rag системах центральным инструментом являются ретриверы - объекты, которые отвечают за поиск близкой к запросу информации (контекста). Одним из них является BM25Retriever, основанный на частоте встречаемости. В отличие от аналогов, использующих векторные представления, он полагается на точные совпадение единиц, на которые разбит текст (токенов).

Для демонстрационных целей возмем набор текстов с описанием компьютерных угроз с сайта MITRE ATT&CK (тут и тут):

import numpy as np

train_sents = ["Contagious Interview has utilized open-source indicator of compromise repositories to determine their exposure to include VirusTotal, and MalTrail",
"Kimsuky has used LLMs to identify think tanks, government organizations, etc. that have information",
'''Sandworm Team researched Ukraine's unique legal entity identifier (called an "EDRPOU" number), including running queries on the EDRPOU website, in preparation for the NotPetya attack. Sandworm Team has also researched third-party websites to help it craft credible spearphishing emails''',
"During the 2015 Ukraine Electric Power Attack, Sandworm Team moved their tools laterally within the corporate network and between the ICS and corporate network",
"During the 2022 Ukraine Electric Power Attack, Sandworm Team used a Group Policy Object (GPO) to copy CaddyWiper's executable msserver.exe from a staging server to a local hard drive before deployment"]


query_sent="Mustang Panda has used open-source research to identify information about victims to use in targeting to include creating weaponized phishing lures and attachments"

BM25Retriever

Для корректной работы BM25Retriever важен способ разбиения текста на единицы, для этого используется параметр preprocess_func. Зададим функцию, осуществляющую деление по словам и их стемминг:

import nltk
from nltk.stem import SnowballStemmer
nltk.download('punkt_tab')

def preprocess_func(text):
    text = text.lower()
    words_l = nltk.tokenize.word_tokenize(text)
    stemmer = SnowballStemmer('english')
    words_l = [stemmer.stem(w) for w in words_l]
    return words_l

Верхнеуровнево для работы надо уметь создать BM25Retriever, используя, например, метод from_texts и найти ближайшие тексты с invoke:

from langchain_community.retrievers import BM25Retriever

# Возвращается объект pydantic BaseModel (в __init__ Serializable)
bm_retriever = BM25Retriever.from_texts(train_sents, k=2, preprocess_func=preprocess_func,
                                        metadatas=[{'len':len(it)} for it in train_sents],
                                        ids=range(len(train_sents)))

bm_retriever.invoke(query_sent)

from_texts

from_texts получает аргументы:

texts - список текстов, формирующих базу для поиска;
k - количество ближайших текстов в ответ на запрос поиска;
preprocess_func - функция для разбиения текстов на токены;
metadatas - список словарей с метаданными для каждого текста в texts;
ids - список идентификаторов для каждого текста в texts.

Следует отметить, что k можно поменять (например, когда вы загружаете настроенный дамп retriever-а) так: bm_retriever.k = 1.

from_texts выполняет следующий код:

texts_processed = [preprocess_func(t) for t in texts]
vectorizer = BM25Okapi(texts_processed...)

invoke

В invoke выполняются:

invoke из BaseRetriever вызывает _get_relevant_documents;
_get_relevant_documents осуществляет препроцессинг и возврат наиболее подходящих через get_top_n:

processed_query = self.preprocess_func(query)
self.vectorizer.get_top_n(processed_query, self.docs, n=self.k)

Подытоживая, from_texts инициирует препроцессинг текстов и создает класс BM25Okapi, а invoke - препроцессинг запроса и вызывает метод get_top_n объекта класса BM25Okapi. Теперь разберемся с BM25Okapi и его особенностями.

okapi

в конструкторе подсчитываются:

в get_top_n:

каждый документ (объект с текстом и метаданными) получает скор близости к запросу get scores in BM25Okapi
документы сортируются по убыванию скора и выбирается заданное число
соответствующий код:

scores = self.get_scores(query)
top_n = np.argsort(scores)[::-1][:n]
[documents[i] for i in top_n]

Создадим вручную объект класса BM25Okapi и набор документов (без метаданных для простоты):

from rank_bm25 import BM25Okapi
from langchain_core.documents import Document


texts_processed = [preprocess_func(t) for t in train_sents]
vectorizer = BM25Okapi(texts_processed)

docs = [Document(page_content=t) for t in train_sents]

processed_query = preprocess_func(query_sent)

свойства

idf - содержит idf токенов
doc_freqs - содержит список словарей с частотами токенов для каждого документа
get_scores - возвращает список скоров каждого документа относительно query
get_top_n - возвращает топ k самых близких документов
doc_len - содержит список количества токенов в каждом документе
avgdl - содержит среднее количество токенов в документах

формула скора

Посчитаем вручную скор для нулевого документа, который возвращает get_scores:

vectorizer.get_scores(processed_query)

idf для каждого токена считается в BM25Okapi конструкторе по следующей формуле:

Добавление константы (0.5) не дает знаменателю или числителю стать равным нулю и сглаживает рост idf для слов, которые встречаются только в единичных документах.

Посчитаем idf для токена "open-sourc", который встречается в 1 из 5 текстов:

idf = np.log((5-1+0.5)/(1+0.5))
vectorizer.idf['open-sourc'], idf

tf для "open-sourc" в нулевом документе получим так:

N = 0
tf = vectorizer.doc_freqs[N]['open-sourc']
tf

Для извлечения скора документа надо итерировать по всем токенам query и посчитать сумму для каждого по формуле:

Из формулы следует, что для документов с длиной меньше средней слагаемое будет увеличено (короткие получают бонус), а для больше средней - уменьшено (штраф за размытость). Посчитаем добавку для токена "open-sourc":

idf*(tf * (1.5 + 1) /(tf + 1.5 * (1 - 0.75 + 0.75 * vectorizer.doc_len[N] / vectorizer.avgdl)))

А теперь выведем скор для всего нулевого документа:

score = 0
# для нулевого документа
for q in processed_query:
    q_freq = vectorizer.doc_freqs[N].get(q,0)
    adding = vectorizer.idf.get(q, 0) * (q_freq * (1.5 + 1) /
                                        (q_freq + 1.5 * (1 - 0.75 + 0.75 * vectorizer.doc_len[N] / vectorizer.avgdl)))
    if q=='open-sourc':
      print(f'принт добавка от "open-sourc" - {adding}')
    score += adding
score

По скору наш (нулевой) документ второй, в такой очередности он и выводится в get_top_n:

k = 2

vectorizer.get_top_n(processed_query, docs, n=k)

#bm25retriever #retrievers #llm #rag #tfidf #nlp #ml