Библиотека NumPy или библиотека SciKit-Learn могут удовлетворить все ваши потребности в извлечении данных.
Делиться
В настоящее время, благодаря технологии Retrieval Augmented Generation (RAG), векторные базы данных привлекают к себе большое внимание в мире искусственного интеллекта.
Многие утверждают, что для создания системы RAG и управления эмбеддингами необходимы такие инструменты, как Pinecone, Weaviate, Milvus или Qdrant. Если вы работаете над корпоративными приложениями, содержащими сотни миллионов векторов, то подобные инструменты просто необходимы. Они позволяют выполнять операции CRUD, фильтровать данные по метаданным и использовать индексирование на диске, выходящее за рамки оперативной памяти компьютера.
Однако для большинства внутренних инструментов, ботов для создания документации или агентов MVP добавление выделенной базы данных векторов может оказаться излишним. Это увеличивает сложность, задержки в сети, добавляет затраты на сериализацию и усложняет управление.
На самом деле «векторный поиск» (то есть часть RAG, отвечающая за извлечение данных) — это всего лишь умножение матриц. И в Python уже есть одни из лучших в мире инструментов для этого.
В этой статье мы покажем, как создать готовый к использованию компонент поиска в конвейере RAG для небольших и средних объемов данных, используя только NumPy и SciKit-Learn. Вы увидите, что поиск по миллионам текстовых строк возможен за миллисекунды, при этом все данные хранятся в памяти и не требуют каких-либо внешних зависимостей.
Понимание процесса извлечения информации как матричной математики
Как правило, RAG включает четыре основных этапа:
- Встраивание: Преобразуйте текст исходных данных в векторы (списки чисел с плавающей запятой).
- Хранение: Сохраните эти векторные данные в базе данных.
- Retrieve: Найти векторы, которые математически «близки» к искомому вектору.
- Генерация: Передайте соответствующий текст в программу LLM и получите окончательный ответ.
Шаги 1 и 4 основаны на больших языковых моделях. Шаги 2 и 3 относятся к области применения векторных баз данных. Мы сосредоточимся на частях 2 и 3 и на том, как мы можем полностью избежать использования векторных баз данных.
Но когда мы ищем в нашей векторной базе данных, что же на самом деле означает «близость»? Обычно это косинусное сходство . Если ваши два вектора нормализованы так, что их величина равна 1, то косинусное сходство — это просто скалярное произведение этих двух векторов.
Если у вас есть одномерный вектор запроса размером N, Q (1xN), и база данных векторов документов размером M × N, D (MxN), то поиск наилучших совпадений не является запросом к базе данных; это операция умножения матриц, скалярное произведение матрицы D на транспонированную матрицу Q.
Баллы = DQ^T
NumPy разработан для эффективного выполнения подобных операций с использованием подпрограмм, задействующих современные возможности процессора, такие как векторизация.
Реализация
Мы создадим класс под названием SimpleVectorStore для обработки ввода, индексирования и извлечения данных. Наши входные данные будут состоять из одного или нескольких файлов, содержащих текст, по которому мы хотим выполнить поиск. Использование Sentence Transformers для локальных встраиваний позволит всему работать в автономном режиме.
Предварительные требования
Создайте новую среду разработки, установите необходимые библиотеки и запустите блокнот Jupyter.
Введите следующие команды в командную оболочку. Я использую UV в качестве менеджера пакетов; измените настройки в соответствии с тем, какой инструмент вы используете.
$ uv init ragdb $ cd ragdb $ uv venv ragdb $ source ragdb/bin/activate $ uv pip install numpy scikit-learn sentence-transformers jupyter $ jupyter notebook
Встроенное векторное хранилище
Нам не нужен сложный сервер. Всё, что нам нужно, это функция для загрузки текстовых данных из входных файлов и разбиения их на фрагменты размером в байты, а также класс с двумя списками: один для исходных фрагментов текста, а другой для матрицы встраивания. Вот код.
import numpy as np import os from sentence_transformers import SentenceTransformer from sklearn.metrics.pairwise import cosine_similarity from typing import List, Dict, Any from pathlib import Path class SimpleVectorStore: def __init__(self, model_name: str = 'all-MiniLM-L6-v2'): print(f»Загрузка модели встраивания: {model_name}…») self.encoder = SentenceTransformer(model_name) self.documents = [] # Хранит необработанный текст и метаданные self.embeddings = None # Станет массивом numpy def add_documents(self, docs: List[Dict[str, Any]]): «»» Загружает документы. docs format: [{'text': '…', 'metadata': {…}}, …] «»» texts = [d['text'] for d in [документы] # 1. Генерация эмбеддингов print(f»Эмбеддинг {len(texts)} документов…») new_embeddings = self.encoder.encode(texts) # 2. Нормализация эмбеддингов # (Критическая оптимизация: позволяет скалярному произведению аппроксимировать косинусное сходство) norm = np.linalg.norm(new_embeddings, axis=1, keepdims=True) new_embeddings = new_embeddings / norm # 3. Обновление хранилища, если self.embeddings равно None: self.embeddings = new_embeddings else: self.embeddings = np.vstack([self.embeddings, new_embeddings]) self.documents.extend(docs) print(f»Хранилище теперь содержит {len(self.documents)} документов.») def search(self, query: str, k: int = 5): «»» Извлекает k наиболее похожих документов. «»» if self.embeddings is None or len(self.documents) == 0: print(«Предупреждение: хранилище векторов пусто. Нет документов для поиска.») return [] # 1. Встраивание и нормализация запроса query_vec = self.encoder.encode([query]) norm = np.linalg.norm(query_vec, axis=1, keepdims=True) query_vec = query_vec / norm # 2. Векторизованный поиск (умножение матриц) # Форма результата: (1, N_docs) scores = np.dot(self.embeddings, query_vec.T).flatten() # 3. Получение K лучших индексов # argsort сортирует по возрастанию, поэтому мы берем последние k и переворачиваем их # Убедитесь, что k не превышает количество документов k = min(k, len(self.documents)) top_k_indices = np.argsort(scores)[-k:][::-1] results = [] for idx in top_k_indices: results.append({ «score»: float(scores[idx]), «text»: self.documents[idx]['text'], «metadata»: self.documents[idx].get('metadata', {}) }) return results def load_from_directory(directory_path: str, chunk_size: int = 1000, overlap: int = 200): «»» Читает файлы .txt и разбивает их на перекрывающиеся фрагменты. «»» docs = [] # Используйте pathlib для надежной обработки и разрешения путей path = Path(directory_path).resolve() if not path.exists(): print(f»Ошибка: Каталог '{path}' не найден.») print(f»Текущий рабочий каталог: {os.getcwd()}») return docs print(f»Загрузка документов из: {путь}») for file_path in path.glob(«*.txt»): try: with open(file_path, «r», encoding=»utf-8″) as f: text = f.read() # Простая разбивка текста на фрагменты с помощью скользящего окна # Мы перебираем текст с шагом меньше размера фрагмента # чтобы создать перекрытие (сохраняя контекст между фрагментами). шаг = размер фрагмента — перекрытие для i в диапазоне (0, len(text), шаг): фрагмент = text[i : i + размер фрагмента] # Пропускаем фрагменты, которые слишком малы (например, оставшиеся пробелы) если len(chunk) < 50: продолжить docs.append({ "text": chunk, "metadata": { "source": file_path.name, "chunk_index": i } }) except Exception as e: print(f"Предупреждение: Не удалось прочитать файл {file_path.name}: {e}") print(f"Успешно загружено {len(docs)} фрагментов из {len(list(path.glob('*.txt')))} файлов." return docs
Используемая модель встраивания
В коде используется модель MiniLM-L6-v2 из библиотеки Sentence Transformers . Она была выбрана по следующим причинам:
- Он быстрый и лёгкий.
- Она создает 384-мерные векторы, которые используют меньше памяти, чем более крупные модели.
- Он хорошо справляется с широким спектром задач, связанных с английским языком, и не требует специальной тонкой настройки.
Эта модель — всего лишь предложение. Вы можете использовать любую другую модель встраивания, если у вас есть предпочтительная.
Зачем нужна нормализация?
Вы можете заметить этапы нормализации в коде. Мы упоминали об этом ранее, но для ясности, для двух векторов X и Y косинусное сходство определяется следующим образом:
Сходство = (X · Y) / (||X|| * ||Y||)
Где:
- X · Y — скалярное произведение векторов X и Y.
- ||X|| — это величина (длина) вектора X.
- ||Y|| — это величина вектора Y.
Поскольку деление требует дополнительных вычислений, если все наши векторы имеют единичную величину, знаменатель равен 1, поэтому формула сводится к скалярному произведению X и Y, что значительно ускоряет поиск.
Тестирование производительности
Первое, что нам нужно сделать, это получить входные данные для работы. Для этого можно использовать любой текстовый файл. Для предыдущих экспериментов с RAG я использовал книгу, которую скачал с Project Gutenberg. Неизменно захватывающая:
« Болезни крупного рогатого скота, овец, коз и свиней» Джона А. В. Доллара и Г. Муссу.
Обратите внимание, что вы можете просмотреть страницу «Разрешения, лицензирование и другие часто задаваемые вопросы» проекта Gutenberg, перейдя по следующей ссылке.
https://www.gutenberg.org/policy/permission.html
Вкратце, подавляющее большинство электронных книг проекта Gutenberg находятся в общественном достоянии в США и других странах мира. Это означает, что никто не может предоставить или отказать в разрешении распоряжаться этим материалом по своему усмотрению.
«…по вашему усмотрению» включает в себя любое коммерческое использование, переиздание в любом формате, создание производных работ или исполнение произведений.
Я скачал текст книги с сайта Project Gutenberg на свой локальный компьютер, используя эту ссылку.
https://www.gutenberg.org/ebooks/73019.txt.utf-8
Эта книга содержала приблизительно 36 000 строк текста. Для запроса к книге требуется всего шесть строк кода. В моем примере, в строке 2315 книги обсуждается заболевание, называемое кондилома. Вот отрывок:
ВОСПАЛЕНИЕ МЕЖЦИФРОВОГО ПРОСТРАНСТВА.
(КОНДИЛОМАТЫ.)
Кондиломы возникают в результате хронического воспаления кожи, покрывающей
межпальцевая связка. Любое повреждение этой области, вызывающее даже
Поверхностные повреждения могут привести к хроническому воспалению кожи и
гипертрофия сосочков, первая стадия образования
кондиломы.Травмы, вызванные смещением шнуров в межпальцевое пространство, приводили к…
При подковывании рабочих волов также полезно поднимать ноги, что помогает достичь желаемых результатов.
причины.
Итак, мы зададим вопрос: «Что такое кондиломы?» Обратите внимание, что мы не получим точного ответа, поскольку не передаём результаты поиска в LLM, но мы должны увидеть, что наш поиск возвращает фрагмент текста, который предоставил бы LLM всю необходимую информацию для формулирования ответа, если бы мы это сделали.
%%time # 1. Инициализация хранилища = SimpleVectorStore() # 2. Загрузка документов real_docs = load_from_directory(«/mnt/d/book») # 3. Добавление в хранилище, если real_docs: store.add_documents(real_docs) # 4. Поиск results = store.search(«Что такое кондиломы?», k=1) results
И вот результат.
Загрузка модели встраивания: all-MiniLM-L6-v2… Загрузка документов из: /mnt/d/book Успешно загружено 2205 фрагментов из 1 файла. Встраивание 2205 документов… Теперь хранилище содержит 2205 документов. Время работы ЦП: пользовательское 3,27 с, системное: 377 мс, общее: 3,65 с. Время выполнения: 3,82 с. [{'score': 0.44883957505226135, 'text': 'две последниеnфаланги, причем последняя операция проще первой иnобеспечивает лоскуты более правильной формы и лучше приспособлены дляnобразования удовлетворительного культя.nnn ВОСПАЛЕНИЕ МЕЖПАЛЬЦЕВОГО ПРОСТРАНСТВА.nn(КОНДИЛОМАТЫ.)nn Кондиломы возникают в результате хронического воспаления кожи, покрывающейnмежпальцевую связку. Любое повреждение этой области, вызывающее даже поверхностное поражение, может привести к хроническому воспалению кожи и гипертрофии сосочков, первой стадии образования кондилом. Побочными причинами также являются травмы, полученные от шнуров, попавших в межпальцевое пространство при подковывании рабочих волов. Воспаление межпальцевого пространства также является распространенным осложнением афтозных высыпаний вокруг когтей и в пространстве между ними. Постоянный контакт с подстилкой, навозом и мочой способствует инфицированию поверхностных или глубоких ран и, вызывая чрезмерную грануляцию, приводит к гипертрофии сосочкового слоя.
Прочитать, разбить на фрагменты, сохранить и корректно запросить текстовый документ объемом 36000 строк менее чем за 4 секунды — это довольно хороший результат.
SciKit-Learn: Путь обновления
NumPy хорошо подходит для поиска методом перебора. Но что делать, если у вас десятки или сотни документов, и перебор слишком медленный? Прежде чем переходить к векторной базе данных, можно попробовать алгоритм NearestNeighbors из библиотеки SciKit-Learn. Он использует древовидные структуры, такие как KD-Tree и Ball-Tree, чтобы ускорить поиск до O(log N) вместо O(N).
Чтобы это проверить, я скачал с Gutenberg множество других книг, в том числе:
- Рождественская песнь Чарльза Диккенса
- Жизнь и приключения Санта-Клауса, автор Л. Фрэнк Баум.
- «Война и мир» Толстого
- «Прощание с оружием» Хемингуэя
В общей сложности эти книги содержат около 120 000 строк текста. Я скопировал и вставил все пять входных файлов с книгами десять раз, в результате чего получилось пятьдесят файлов и 1,2 миллиона строк текста. Это примерно 12 миллионов слов, если предположить, что в среднем в строке 10 слов. Для сравнения, эта статья содержит приблизительно 2800 слов, поэтому объем данных, с которым мы проводим тестирование, эквивалентен более чем 4000-кратному объему этого текста.
$ dir achristmascarol — Copy (2).txt cattle_disease — Copy (9).txt santa — Copy (6).txt achristmascarol — Copy (3).txt cattle_disease — Copy.txt santa — Copy (7).txt achristmascarol — Copy (4).txt cattle_disease.txt santa — Copy (8).txt achristmascarol — Copy (5).txt farewelltoarms — Copy (2).txt santa — Copy (9).txt achristmascarol — Copy (6).txt farewelltoarms — Copy (3).txt santa — Copy.txt achristmascarol — Copy (7).txt farewelltoarms — Copy (4).txt santa.txt achristmascarol — Copy (8).txt farewelltoarms — Copy (5).txt warandpeace — Copy (2).txt achristmascarol — Copy (9).txt farewelltoarms — Copy (6).txt warandpeace — Copy (3).txt achristmascarol — Copy.txt farewelltoarms — Copy (7).txt warandpeace — Copy (4).txt achristmascarol.txt farewelltoarms — Copy (8).txt warandpeace — Copy (5).txt cattle_disease — Copy (2).txt farewelltoarms — Copy (9).txt warandpeace — Copy (6).txt cattle_disease — Copy (3).txt farewelltoarms — Copy.txt warandpeace — Copy (7).txt cattle_disease — Copy (4).txt farewelltoarms.txt warandpeace — Copy (8).txt cattle_disease — Copy (5).txt santa — Copy (2).txt warandpeace — Copy (9).txt cattle_disease — Copy (6).txt santa — Copy (3).txt warandpeace — Copy.txt cattle_disease — Copy (7).txt santa — Copy (4).txt warandpeace.txt cattle_disease — Copy (8).txt santa — Copy (5).txtДопустим, мы ушли
Допустим, в конечном итоге мы ищем ответ на следующий вопрос:
К кому же Николай признался своей матери в любви после рождественских каникул?
Если вы не знали, это взято из романа «Война и мир».
Давайте посмотрим, как наш новый поиск покажет себя на этом большом объеме информации.
Вот код, использующий SciKit-Learn.
Во-первых, у нас появился новый класс, реализующий алгоритм ближайшего соседа из библиотеки SciKit-Learn.
from sklearn.neighbors import NearestNeighbors class ScikitVectorStore(SimpleVectorStore): def __init__(self, model_name='all-MiniLM-L6-v2'): super().__init__(model_name) # Метод перебора часто быстрее деревьев для многомерных данных, # если только N не очень велико, но 'ball_tree' может помочь в определенных случаях. self.knn = NearestNeighbors(n_neighbors=5, metric='cosine', algorithm='brute') self.is_fit = False def build_index(self): print(«Building Scikit-Learn Index…») self.knn.fit(self.embeddings) self.is_fit = True def search(self, query: str, k: int = 5): if not self.is_fit: self.build_index() query_vec = self.encoder.encode([query]) # Примечание: Scikit-learn обрабатывает нормализацию внутри себя для косинусной метрики, # если она настроена, но явное указание лучше. distances, indices = self.knn.kneighbors(query_vec, n_neighbors=k) results = [] for i in range(k): idx = indices[0][i] # Преобразуем расстояние обратно в оценку сходства (1 — dist) score = 1 — distances[0][i] results.append({ «score»: score, «text»: self.documents[idx]['text'] }) return results
А наш код поиска так же прост, как и для версии на NumPy.
%%time # 1. Инициализация хранилища = ScikitVectorStore() # 2. Загрузка документов real_docs = load_from_directory(«/mnt/d/book») # 3. Добавление в хранилище, если real_docs: store.add_documents(real_docs) # 4. Поиск results = store.search(«Кто, после рождественских каникул, рассказал своей матери о своей любви», k=1) results
И результаты нашей работы.
Загрузка модели встраивания: all-MiniLM-L6-v2… Загрузка документов из: /mnt/d/book Успешно загружено 73060 фрагментов из 50 файлов. Встраивание 73060 документов… Теперь хранилище содержит 73060 документов. Создание индекса Scikit-Learn… Время работы ЦП: пользовательское 1 мин 46 с, системное: 18,3 с, общее: 2 мин 4 с. Общее время: 1 мин 13 с [{'score': 0.6972659826278687, 'text': 'nГЛАВА XIIInnВскоре после рождественских каникул Николас рассказал своей матери о своей любвиnк Соне и о своем твердом намерении жениться на ней. Графиня, которая давно заметила происходящее между ними и ожидала этого заявления, выслушала его молча, а затем сказала сыну, что он может жениться на ком пожелает, но ни она, ни его отец не дадут своего благословения на такой брак. Николай впервые почувствовал, что мать недовольна им и что, несмотря на свою любовь к нему, она не уступит. Холодно, не глядя на сына, она послала за мужем и, когда он пришел, попыталась кратко и холодно сообщить ему о случившемся в присутствии сына, но, не в силах сдержаться, разрыдалась от досады и вышла из комнаты. Старый граф начал нерешительно увещевать Николая и умолять его отказаться от своего намерения. Николай ответил, что не может отказаться от своего слова, и его отец, вздыхая и явно смущенный, очень скоро замолчал.
Почти все 1 минуту 13 секунд, затраченные на описанную выше обработку, ушли на загрузку и разбиение входных данных на блоки. Сам поиск, когда я запускал его отдельно, занял менее одной десятой секунды!
Совсем неплохо.
Краткое содержание
Я не утверждаю, что векторные базы данных не нужны. Они решают конкретные задачи, с которыми не справляются NumPy и SciKit-Learn. Вам следует перейти с таких библиотек, как SimpleVectorStore или ScikitVectorStore , на Weaviate/Pinecone/pgvector и т. д., если выполняется любое из следующих условий.
Сохранение данных: Вам необходимо, чтобы данные сохранялись после перезапуска сервера без необходимости перестраивать индекс из исходных файлов каждый раз. Хотя np.save или сериализация (picking) подходят для простого сохранения данных. Разработка всегда предполагает компромиссы. Использование векторной базы данных усложняет вашу конфигурацию в обмен на масштабируемость, которая вам, возможно, сейчас не нужна. Если вы начнете с более простой конфигурации RAG, используя NumPy и/или SciKit-Learn для процесса извлечения данных, вы получите:
Оперативная память является узким местом: ваша матрица встраивания превышает объем памяти вашего сервера. Примечание: 1 миллион векторов 384-мерного типа [float32] занимает всего около 1,5 ГБ оперативной памяти, поэтому в память можно поместить очень много данных.
Частота операций CRUD: Вам необходимо постоянно обновлять или удалять отдельные векторы во время чтения. Массивы NumPy, например, являются неизменяемыми, и добавление требует копирования всего массива, что медленно.
Фильтрация метаданных: Вам понадобятся сложные запросы, например, «Найти векторы, близкие к X, где user_id=10 И дата > 2023». В NumPy для этого требуются булевы маски, что может привести к путанице.
В инженерной практике всегда приходится идти на компромиссы. Использование векторной базы данных усложняет вашу систему в обмен на масштабируемость, которая вам, возможно, сейчас не нужна. Если начать с более простой конфигурации RAG с использованием NumPy и/или SciKit-Learn для процесса поиска, вы получите:
- Низкая задержка. Отсутствие сетевых переходов.
- Снижение затрат. Никаких подписок на SaaS или дополнительных экземпляров.
- Простота. Это всего лишь скрипт на Python.
Так же, как вам не нужен спортивный автомобиль, чтобы поехать в продуктовый магазин. Во многих случаях NumPy или SciKit-Learn могут быть всем, что вам нужно для поиска с использованием алгоритма RAG.
Источник: towardsdatascience.com
