Векторные БД vs Точность — часть 1

Решил я собрать «по-быстрому» локальный RAG(retrieval augmentation generation), который будет находить термины из словаря Ожегова. Изучив просторы интернета, понял. Все сводится к рецепту (упрощенная интерпретация):

1. Берем нужный нам текст и нарезаем на куски

2. С помощью «эмбеддера» превращаем в эмбеддинги

3. Грузим в векторную БД

4. Цепляем ChatGPT или альтернативу с помощью Langchain

5. Пишем промпт

6. Радуемся

В случае с толковым словарем Ожегова, может быть 2 типа вопросов:

1. По содержимому или значению найти термин

2. Найти термин, что он обозначает

Радовался я не долго. Через время начал понимать, что он путается в терминах во втором типе вопросов. Прошу один термин, а получаю совсем другой. Давайте разбираться почему.

На схеме выше видно. Для ответа используется «результат поиска + вопрос» — нужно проверить, запрос и ответ из БД. Так как мы сами грузим датасет и знаем какой ответ должен быть, то мы можем автоматически проверить, правильно ли нам возвращает векторная БД ответ по нашему запросу.

Готовим данные

Берем нужный нам текст и нарезаем на куски.

import re import pandas as pd  with open("ozhegov.txt", mode="r", encoding="UTF-8") as file:     text_lines = file.readlines()  def return_first_match(pattern, text):     result = re.findall(pattern,text)     result = result[0] if result else ""     return result  data = []  for line in text_lines:      title = return_first_match(r"^[а-яА-Я]{2,}(?=,)", line)     text = return_first_match(r"\.\s([А-Я]+.*)\n", line)      if(len(title) > 3):         data.append(             {                 "title": title,                 "text" : text             })  dataset = pd.DataFrame(data) dataset.to_csv('ozhegov_dataset.csv') dataset.head()

Для нашего эксперимента не будем использовать весь датасет, а возьмем только 1000 случайных штук.

#бывают пустые в моем датасете dataset = dataset[dataset['title'].str.len() > 0] dataset = dataset[dataset['text'].str.len() > 0] dataset = dataset.sample(n=1000)  dataset.astype({"text": str, "title": str}) dataset.info(show_counts=True)  dataset.head()

<class 'pandas.core.frame.DataFrame'> Index: 1000 entries, 28934 to 16721 Data columns (total 3 columns):  #   Column      Non-Null Count  Dtype  ---  ------      --------------  -----   0   Unnamed: 0  1000 non-null   int64   1   title       1000 non-null   object  2   text        1000 non-null   object dtypes: int64(1), object(2) memory usage: 31.2+ KB

Превращаем в эмбеддинги

«Эмбеддер» долго не выбирал, воспользовался Рейтинг русскоязычных энкодеров предложений

Про сами векторные БД и как они работают, какие есть — рассказывать не буду. По этому поводу уже написаны статьи. Буду использовать Chroma (топ-1 из этой статьи).

Для чистоты эксперимента, решил взять топ-5 моделей и 3 разные функции расстояния доступные в Chroma

models = [     "intfloat/multilingual-e5-large",     "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",     "symanto/sn-xlm-roberta-base-snli-mnli-anli-xnli",     "cointegrated/LaBSE-en-ru",     "sentence-transformers/LaBSE" ]  distances = [     "l2",     "ip",     "cosine" ]

Готовим Chroma

Ставим нужные pip пакеты

%pip install -U sentence-transformers ipywidgets chromadb chardet charset-normalizer

Бывает ошибка с установкой, в самой ошибке есть решение

HINT: This error might have occurred since this system does not have Windows Long Path support enabled. You can find information on how to enable this at https://pip.pypa.io/warnings/enable-long-paths

https://learn.microsoft.com/en-us/windows/win32/fileio/maximum-file-path-limitation?tabs=powershell#enable-long-paths-in-windows-10-version-1607-and-later

Запустим Chroma в Docker

docker pull chromadb/chroma docker run -p 8000:8000 chromadb/chroma

Функции для работы с Chroma

Определяем необходимые функции для создания и удаления коллекций. А также функцию для поиска, в которой мы забираем запись минимальной дистанцией + добавляем данных по индексу из датасета. Индексы в БД равны индексам в датасете.

from chromadb.utils import embedding_functions import chromadb chroma_client = chromadb.HttpClient(host="localhost", port=8000)  def create_collection(model_name, distance):          chroma_client = chromadb.HttpClient(host="localhost", port=8000)      sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=model_name)          #в этом эксперименте не будем использовать, нам нужно найти термин     #text_collection = chroma_client.create_collection(name='text', embedding_function=sentence_transformer_ef)          title_collection = chroma_client.create_collection(name="title", embedding_function=sentence_transformer_ef, metadata={"hnsw:space": distance})      ids = list(map(str, dataset.index.values.tolist()))     #text_collection.add(ids = ids, documents=dataset["text"].tolist())     title_collection.add(ids = ids, documents=dataset["title"].tolist())      return title_collection  def delete_collection():     chroma_client.delete_collection("title")   def query_collection(collection, query, max_results, dataframe, model_name, distance):     results = collection.query(query_texts=query, n_results=max_results, include=['distances'])      #print(results)     df = pd.DataFrame({                 'id':results['ids'][0],                  'score':list(map(float,results['distances'][0])),                 'query': query,                 'title': dataframe[dataframe.index.isin(list(map(int,results['ids'][0])))]['title'],                 'content': dataframe[dataframe.index.isin(list(map(int,results['ids'][0])))]['text'],                 'model_name': model_name,                 'distance': distance                 })          # Забираем с минимальной дистанцией, значит он ближе и больше похож     df = df[df.score == df.score.min()]     df['is_found'] = df.apply(lambda row: row.query == row.title, axis=1)          return df 

Формируем тестовый датасет и стартуем

Формируем тестовый датасет из случайных 100 штук из датасета загруженного в Chroma.

test_dataset = dataset.sample(n=100) test_dataset.head() test_results = pd.DataFrame()

Запускаем и собираем результаты для каждой модели с разной функцией расстояния.

for model in models:     for distance in distances:         print(f"{model} - {distance}")         try:             delete_collection()         except Exception as ex:             print(f"delete_collection error: {ex}")          collection = create_collection(model, distance)          for title in test_dataset["title"].tolist():             test_results = test_results._append(query_collection(             collection=collection,             query=title,             max_results=5,             dataframe=dataset,             model_name=model,             distance=distance))              print(f"{len(test_results)}")            test_results.to_csv("results_ozhegov2.csv") test_results.head()

Смотрим результаты

Теперь посчитаем количество найденных (правильных результатов)

finally_result = pd.DataFrame() for model in models:     for distance in distances:         df = test_results.loc[test_results['model_name'].str.contains(model) == True]         df = df.loc[df['distance'].str.contains(distance) == True]          finally_result = finally_result._append(pd.DataFrame({                 'founded': [len(df[df['is_found'] == True])],                 'model_name': [model],                 'distance': [distance]                 }))          finally_result.head(15)

Заключение

Собрать «по-быстрому» локальный RAG для работы с терминами пока не удалось и выдать готовый рецепт.

Текущие ~25% — сложно назвать хорошей точностью.

Какие я вижу варианты решения проблемы с точностью:

1. Использовать гибридный поиск c BM25, не все векторные БД его поддерживают

   1. Chroma — в процессе реализации

   2. Pinecone — public preview — нету в Docker — запрос фичи

   3. Milvus — в процессе реализации

   4. Weaviate — поддерживает

   5. Qdrant — поддерживает

   6. Elasticsearch — поддерживает

   7. Opensearch — поддерживает

2. Прикрутить в качестве второго ретривера Postgres + BM25 и искать сразу в двух — звучит так себе + дублировать информацию…

3. Тюнить модель, но это уже далеко не «по-быстрому»

4. Поработать с текстом (лемматизация, нормализация и т.п.) — сильно сомневаюсь, что поможет.

Исходный код

Во второй части попробую использовать гибридный поиск

P.S. Буду ждать в комментариях какой вариант еще попробовать, чтобы можно было «по-быстрому» и локально развернуть.