Использование векторной базы данных для поиска речей Белого дома
Эта статья была первоначально опубликована на The New Stack.
Приближается сезон кампании для президентской кампании в США. Это хорошее время, чтобы оглянуться на некоторые речи, произнесенные администрацией Байдена в течение его первых двух лет на посту. Разве не было бы здорово выполнить поиск по расшифровкам некоторых речей, чтобы больше узнать о том, как Белый дом до сих пор освещал определенные темы?
Допустим, мы хотим искать по содержанию речи. Как бы мы это сделали? Мы могли бы использовать семантический поиск. Семантический поиск сейчас является одной из самых горячих тем в искусственном интеллекте (ИИ). Он стал важнее по мере того, как мы наблюдали рост популярности приложений обработки естественного языка (NLP), таких как ChatGPT. Вместо того чтобы многократно обращаться к GPT, что дорого как экономически, так и экологически, мы можем использовать векторную базу данных для кэширования результатов (например, с помощью GPTCache).
В этом руководстве мы локально запустим векторную базу данных, чтобы искать речи Байдена с 2021 по 2022 год по содержанию. Набор данных, который мы используем, — это “The White House (Speeches and Remarks) 12/10/2022”, который мы нашли на Kaggle и сделали доступным для загрузки через Google Drive для этого примера. Пошаговый блокнот для этого руководства доступен на GitHub.
Прежде чем мы погрузимся в код, убедитесь, что вы скачали необходимые предварительные компоненты. Нам нужны четыре библиотеки: PyMilvus, Milvus, Sentence-Transformers и gdown. Вы можете получить необходимые библиотеки из PyPi, выполнив: pip3 install pymilvus==2.2.5 sentence-transformers gdown milvus.
Подготовка набора данных речей Белого дома
Как и почти в любом проекте AI/ML, основанном на реальных наборах данных, сначала нам нужно подготовить данные. Мы используем gdown, чтобы скачать набор данных, и zipfile, чтобы извлечь его в локальную папку. После выполнения приведенного ниже кода мы ожидаем увидеть файл с названием “The white house speeches.csv” в папке с названием “white_house_2021_2022”.
import gdown
url = 'https://drive.google.com/uc?id=10_sVL0UmEog7mczLedK5s1pnlDOz3Ukf'
output = './white_house_2021_2022.zip'
gdown.download(url, output)
import zipfile
with zipfile.ZipFile("./white_house_2021_2022.zip","r") as zip_ref:
zip_ref.extractall("./white_house_2021_2022")
Мы используем pandas, чтобы загрузить и изучить данные CSV.
import pandas as pd
df = pd.read_csv("./white_house_2021_2022/The white house speeches.csv")
df.head()
Когда мы смотрим на head данных, что вы замечаете? Первое, что я замечаю, — данные имеют четыре столбца: заголовок, дату и время, место и столбец с речью. Второе — в них есть нулевые значения. Нулевые значения не всегда являются проблемой, но для наших данных они являются проблемой.
Очистка набора данных
Речи без какого-либо содержания (нулевые значения в столбце “Speech”) для нас совершенно бесполезны. Давайте удалим наши нулевые значения и снова изучим данные.
df = df.dropna()
df
Теперь мы видим, что на самом деле есть вторая проблема, которая не была сразу очевидна при просмотре только head данных. Если вы посмотрите на последнюю запись, то увидите, что эта запись — всего лишь время. “12:18 P.M. EST” едва ли можно назвать речью. Нет смысла сохранять эту запись. Мы не можем извлечь никакой ценности из сохранения векторного вложения.
Давайте избавимся от всех речей, которые короче определенной длины. Для этого примера я выбрал 50, но вы можете выбрать любое значение, которое вам подходит. Я выбрал 50, исследовав множество разных чисел. Если вы посмотрите стенограммы речей длиной от 20 до 50 символов, то увидите, что многие из них — это места или время, к которым добавлено несколько случайных предложений.
cleaned_df = df.loc[(df["Speech"].str.len() > 50)]cleaned_df
Разобравшись с короткими, бессодержательными речами, мы снова смотрим на наши данные и видим еще одну проблему. Многие речи содержат значения \r\n — переводы строк и возвраты каретки. Эти символы используются для форматирования, но не несут никакой семантической ценности. Следующий шаг в процессе очистки данных — избавиться от них.
cleaned_df["Speech"] = cleaned_df["Speech"].str.replace("\r\n", "")
cleaned_df
Так выглядит намного лучше. Последний шаг — преобразовать столбец “Date_time” в более подходящий формат для хранения в нашей векторной базе данных и сравнения с другими датами и временем. Мы используем библиотеку datetime, чтобы просто преобразовать этот формат даты и времени в универсальный формат YYYY-MM-DD.
import datetime
# Convert the 'date' column to datetime objects
cleaned_df["Date_time"] = pd.to_datetime(cleaned_df["Date_time"], format="%B %d, %Y")
cleaned_df
Настройка векторной базы данных для семантического поиска
Теперь наши данные очищены и готовы к работе. Следующий шаг — запустить векторную базу данных, чтобы фактически искать речи по их содержанию. В этом примере мы используем Milvus Lite, легкую версию Milvus, которую можно запустить без Docker, Kubernetes или необходимости иметь дело с какими-либо YAML-файлами.
Первое, что мы делаем, — определяем некоторые константы. Нам нужно имя коллекции (для векторной базы данных), количество измерений в нашем встроенном векторе, размер пакета и число, определяющее, сколько результатов мы хотим получить при поиске. В этом примере используется sentence transformer MiniLM L6 v2, который создает векторы эмбеддингов размерности 384.
COLLECTION_NAME = "white_house_2021_2022"
DIMENSION = 384
BATCH_SIZE = 128
TOPK = 3
Мы используем default_server из Milvus. Затем мы используем PyMilvus SDK для подключения к нашему локальному серверу Milvus. Если в нашей векторной базе данных есть коллекция с тем же именем, что и имя коллекции, которое мы определили ранее, мы удаляем эту коллекцию, чтобы гарантировать, что начинаем с чистого листа.
from milvus import default_server
from pymilvus import connections, utility
default_server.start()
connections.connect(host="127.0.0.1", port=default_server.listen_port)
if utility.has_collection(COLLECTION_NAME):
utility.drop_collection(COLLECTION_NAME)
Как и в большинстве других баз данных, нам нужна схема для загрузки данных в векторную базу данных Milvus. Сначала мы определяем поля данных, которые должен иметь каждый объект. Хорошо, что мы ранее посмотрели на данные. Мы используем пять полей данных: четыре столбца, которые у нас были ранее, и столбец ID. Только на этот раз вместо фактического текста мы используем векторный эмбеддинг речи.
from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
# object should be inserted in the format of (title, date, location, speech embedding)
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=500),
FieldSchema(name="date", dtype=DataType.VARCHAR, max_length=100),
FieldSchema(name="location", dtype=DataType.VARCHAR, max_length=200),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=DIMENSION)
]
schema = CollectionSchema(fields=fields)
collection = Collection(name=COLLECTION_NAME, schema=schema)
Последнее, что нам нужно определить перед тем, как быть готовыми загрузить данные в векторную базу данных, — это индекс. Существует множество векторных индексов и паттернов, но в этом примере мы используем индекс IVF_FLAT со 128 кластерами. Более крупные приложения обычно используют более 128 кластеров, но у нас в любом случае лишь немного больше 600 записей. Для нашего расстояния мы используем норму L2. После того как мы определим параметры индекса, мы создаём индекс в нашей коллекции и загружаем её для использования.
index_params = {
"index_type": "IVF_FLAT",
"metric_type": "L2",
"params": {"nlist": 128},
}
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()
Получение векторных эмбеддингов из речей
Большая часть того, что мы уже рассмотрели, применима при работе практически с любой базой данных. Мы очистили некоторые данные, запустили экземпляр базы данных и определили схему для нашей базы данных. Помимо определения индекса, ещё одна вещь, которую нам нужно сделать именно для векторных баз данных, — это получить эмбеддинги.
Сначала мы получаем модель преобразователя предложений MiniLM L6 v2, как упоминалось выше. Затем мы создаём функцию, которая выполняет преобразование данных и вставляет их в коллекцию. Эта функция принимает пакет данных, получает эмбеддинги для расшифровок речей, создаёт объект для вставки и вставляет его в коллекцию.
Для контекста: эта функция выполняет пакетное обновление. В этом примере мы пакетно вставляем 128 записей за раз. Единственное преобразование данных, которое мы выполняем при вставке, — это превращение текста речи в эмбеддинг.
from sentence_transformers import SentenceTransformer
transformer = SentenceTransformer('all-MiniLM-L6-v2')
# expects a list of (title, date, location, speech)
def embed_insert(data: list):
embeddings = transformer.encode(data[3])
ins = [
data[0],
data[1],
data[2],
[x for x in embeddings]
]
collection.insert(ins)
Заполнение вашей векторной базы данных
Когда функция, которая создаёт пакетные эмбеддинги и вставки, готова, мы готовы заполнить базу данных. В этом примере мы проходим циклом по каждой строке в нашем dataframe и добавляем данные в список списков, который используем для пакетной обработки данных. Как только мы достигаем нашего размера пакета, мы вызываем функцию embed_insert и сбрасываем наш пакет.
Если после завершения цикла в пакете данных остаются какие-либо данные, мы создаём эмбеддинги и вставляем оставшиеся данные. Наконец, чтобы завершить заполнение нашей векторной базы данных, мы вызываем flush, чтобы гарантировать, что база данных обновлена и проиндексирована.
data_batch = [[], [], [], []]
for index, row in cleaned_df.iterrows():
data_batch[0].append(row["Title"])
data_batch[1].append(str(row["Date_time"]))
data_batch[2].append(row["Location"])
data_batch[3].append(row["Speech"])
if len(data_batch[0]) % BATCH_SIZE == 0:
embed_insert(data_batch)
data_batch = [[], [], [], []]
# Embed and insert the remainder
if len(data_batch[0]) != 0:
embed_insert(data_batch)
# Call a flush to index any unsealed segments.
collection.flush()
Семантический поиск речей Белого дома на основе описаний
Допустим, я хочу найти речь, в которой президент говорил о влиянии возобновляемой энергии в Национальной лаборатории возобновляемой энергии (NREL), и речь, в которой выступают вице-президент и премьер-министр Канады. Я могу найти названия наиболее похожих речей, произнесённых членами Белого дома в 2021–2022 годах, используя векторную базу данных, которую мы только что создали.
Мы можем искать в нашей векторной базе данных речи, наиболее похожие на наши описания. Затем всё, что нам нужно сделать, — это преобразовать описание в векторный эмбеддинг с помощью той же модели, которую мы использовали для получения эмбеддингов речей, а затем выполнить поиск в векторной базе данных.
После того как мы преобразуем описания в векторное вложение, мы используем функцию search в нашей коллекции. Мы передаем вложения в качестве данных для поиска, указываем поле, в котором выполняем поиск, добавляем некоторые параметры для того, как выполнять поиск, лимит на количество результатов и поле, которое хотим вернуть. В этом примере параметры поиска, которые нам нужно передать, — это тип метрики, который должен быть того же типа, что мы использовали при создании индекса (норма L2), и количество кластеров, по которым мы хотим выполнить поиск (устанавливая nprobe равным 10).
import time
search_terms = ["The President speaks about the impact of renewable energy at the National Renewable Energy Lab.", "The Vice President and the Prime Minister of Canada both speak."]
# Search the database based on input text
def embed_search(data):
embeds = transformer.encode(data)
return [x for x in embeds]
search_data = embed_search(search_terms)
start = time.time()
res = collection.search(
data=search_data, # Embeded search value
anns_field="embedding", # Search across embeddings
param={"metric_type": "L2",
"params": {"nprobe": 10}},
limit = TOPK, # Limit to top_k results per search
output_fields=["title"] # Include title field in result
)
end = time.time()
for hits_i, hits in enumerate(res):
print("Title:", search_terms[hits_i])
print("Search Time:", end-start)
print("Results:")
for hit in hits:
print( hit.entity.get("title"), "----", hit.distance)
print()
Когда мы выполняем поиск по предложениям в этом примере, мы ожидаем увидеть вывод, похожий на изображение ниже. Это был успешный поиск, потому что заголовки — именно те, которые мы ожидаем увидеть. Первое описание возвращает заголовок речи, произнесенной президентом Байденом в NREL, а второе описание возвращает заголовок, отражающий речь, произнесенную вице-президентом Харрис и премьер-министром Трюдо.
Резюме
В этом руководстве мы узнали, как использовать векторную базу данных для семантического поиска по речам, произнесенным администрацией Байдена перед промежуточными выборами 2022 года. Семантический поиск позволяет нам взять краткое описание и искать семантически похожий текст, а не только синтаксически похожий текст. Это позволяет нам искать речь по общему описанию, а не на основе конкретных предложений или цитат. Для большинства из нас это значительно упрощает поиск речей, которые могли бы нас заинтересовать.
Читать далее
Zilliz Cloud BYOC Now Available Across AWS, GCP, and Azure
Zilliz Cloud BYOC is now generally available on all three major clouds. Deploy fully managed vector search in your own AWS, GCP, or Azure account — your data never leaves your VPC.
Why AI Databases Don't Need SQL
Whether you like it or not, here's the truth: SQL is destined for decline in the era of AI.
OpenAI o1: What Developers Need to Know
In this article, we will talk about the o1 series from a developer's perspective, exploring how these models can be implemented for sophisticated use cases.