Usando um Banco de Dados Vetorial para Pesquisar Discursos da Casa Branca
Este artigo foi publicado originalmente no The New Stack.
A temporada de campanha está chegando para a campanha presidencial dos EUA. É um bom momento para olhar para trás e analisar alguns dos discursos feitos pelo governo Biden durante seus dois primeiros anos no cargo. Não seria ótimo pesquisar algumas transcrições de discursos para aprender mais sobre as mensagens da Casa Branca sobre determinados temas até agora?
Digamos que queiramos pesquisar o conteúdo de um discurso. Como faríamos isso? Poderíamos usar busca semântica. A busca semântica é um dos tópicos mais quentes em inteligência artificial (IA) no momento. Ela se tornou mais importante à medida que vimos o aumento da popularidade de aplicações de processamento de linguagem natural (NLP), como o ChatGPT. Em vez de consultar repetidamente o GPT, o que é caro tanto econômica quanto ecologicamente, podemos usar um banco de dados vetorial para armazenar os resultados em cache (como com GPTCache).
Neste tutorial, vamos iniciar um banco de dados vetorial localmente para podermos pesquisar os discursos de Biden de 2021 a 2022 por conteúdo. O conjunto de dados que usamos é o conjunto de dados “The White House (Speeches and Remarks) 12/10/2022”, que encontramos no Kaggle e disponibilizamos para download via Google Drive para este exemplo. Um notebook com o passo a passo deste tutorial está disponível no GitHub.
Antes de mergulharmos no código, certifique-se de baixar os pré-requisitos. Precisamos de quatro bibliotecas: PyMilvus, Milvus, Sentence-Transformers e gdown. Você pode obter as bibliotecas necessárias no PyPi executando: pip3 install pymilvus==2.2.5 sentence-transformers gdown milvus.
Preparando o conjunto de dados de discursos da Casa Branca
Como em quase qualquer projeto de IA/ML baseado em conjuntos de dados do mundo real, primeiro precisamos preparar os dados. Usamos gdown para baixar o conjunto de dados e zipfile para extraí-lo para uma pasta local. Depois de executar o código abaixo, esperamos ver um arquivo intitulado “The white house speeches.csv” em uma pasta intitulada “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")
Usamos pandas para carregar e inspecionar os dados CSV.
import pandas as pd
df = pd.read_csv("./white_house_2021_2022/The white house speeches.csv")
df.head()
Quando damos uma olhada no head dos dados, o que você nota? A primeira coisa que noto é que os dados têm quatro colunas: uma coluna de título, data e hora, localização e discurso. A segunda coisa é que há valores nulos. Valores nulos nem sempre são um problema, mas são para os nossos dados.
Limpando o conjunto de dados
Discursos sem qualquer substância (valores nulos na coluna “Speech”) são totalmente inúteis para nós. Vamos remover nossos valores nulos e reexaminar os dados.
df = df.dropna()
df
Agora vemos que, na verdade, há um segundo problema que não era imediatamente óbvio ao olhar apenas para o head dos dados. Se você olhar para a última entrada, verá que essa entrada é apenas um horário. “12:18 P.M. EST” dificilmente é um discurso. Não faz sentido salvar esta entrada. Não conseguimos derivar nenhum valor ao salvar uma incorporação vetorial.
Vamos eliminar todos os discursos que têm menos do que um determinado comprimento. Para este exemplo, escolhi 50, mas você pode escolher o que fizer sentido para você. Escolhi 50 explorando muitos números diferentes. Se você procurar transcrições de discursos entre 20 e 50 caracteres, verá que muitas são locais ou horários com algumas frases aleatórias no meio.
cleaned_df = df.loc[(df["Speech"].str.len() > 50)]cleaned_df
Com os discursos curtos e sem substância resolvidos, mais uma vez olhamos para nossos dados para ver mais um problema. Muitos dos discursos contêm valores \r\n - novas linhas e retornos. Esses caracteres são usados para formatação, mas não contêm nenhum valor semântico. O próximo passo em nosso processo de limpeza de dados é nos livrar deles.
cleaned_df["Speech"] = cleaned_df["Speech"].str.replace("\r\n", "")
cleaned_df
Isso está ficando muito melhor. O passo final é converter a coluna “Date_time” para um formato melhor para armazenar em nosso banco de dados vetorial e ser comparado a outros datetimes. Usamos a biblioteca datetime para simplesmente converter esse formato de datetime em um formato universal 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
Configurando um Banco de Dados Vetorial para Busca Semântica
Nossos dados agora estão limpos e prontos para trabalhar. O próximo passo é iniciar um banco de dados vetorial para realmente pesquisar os discursos pelo seu conteúdo. Para este exemplo, usamos o Milvus Lite, uma versão lite do Milvus que você pode colocar em funcionamento sem Docker, Kubernetes ou lidar com qualquer tipo de arquivo YAML.
A primeira coisa que fazemos é definir algumas de nossas constantes. Precisamos de um nome de coleção (para o banco de dados vetorial), o número de dimensões em nosso vetor incorporado, um tamanho de lote e um número que define quantos resultados queremos de volta quando pesquisamos. Este exemplo usa o sentence transformer MiniLM L6 v2, que produz vetores de embedding de 384 dimensões.
COLLECTION_NAME = "white_house_2021_2022"
DIMENSION = 384
BATCH_SIZE = 128
TOPK = 3
Usamos o default_server do Milvus. Em seguida, usamos o SDK PyMilvus para nos conectar ao nosso servidor Milvus local. Se houver uma coleção em nosso banco de dados vetorial que tenha o mesmo nome que o nome da coleção que definimos anteriormente, descartamos essa coleção para garantir que comecemos com uma página em branco.
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)
Como a maioria dos outros bancos de dados, precisamos de um schema para carregar dados no banco de dados vetorial Milvus. Primeiro, definimos os campos de dados que queremos que cada objeto tenha. Ainda bem que olhamos os dados antes. Usamos cinco campos de dados, as quatro colunas que tínhamos anteriormente e uma coluna de ID. Exceto que desta vez usamos o embedding vetorial do discurso em vez do texto real.
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)
A última coisa que precisamos definir antes de estarmos prontos para carregar dados no banco de dados vetorial é o índice. Existem muitos índices e padrões vetoriais, mas, para este exemplo, usamos o índice IVF_FLAT com 128 clusters. Aplicações maiores geralmente usam mais de 128 clusters, mas, de qualquer forma, temos apenas um pouco mais de 600 entradas. Para nossa distância, medimos usando a norma L2. Depois de definirmos nossos parâmetros de índice, criamos o índice em nossa coleção e o carregamos para uso.
index_params = {
"index_type": "IVF_FLAT",
"metric_type": "L2",
"params": {"nlist": 128},
}
collection.create_index(field_name="embedding", index_params=index_params)
collection.load()
Obtendo Embeddings Vetoriais de Discursos
Muito do que abordamos até agora se aplica ao trabalhar com praticamente qualquer banco de dados. Limpamos alguns dados, inicializamos uma instância de banco de dados e definimos um esquema para nosso banco de dados. Além de definir um índice, outra coisa que precisamos fazer para bancos de dados vetoriais em particular é obter os embeddings.
Primeiro, obtemos o modelo de transformador de sentenças MiniLM L6 v2, conforme mencionado acima. Em seguida, criamos uma função que realiza uma transformação nos dados e os insere na coleção. Esta função recebe um lote de dados, obtém os embeddings das transcrições dos discursos, cria um objeto para inserir e o insere na coleção.
Para contextualizar, esta função realiza uma atualização em lote. Neste exemplo, estamos inserindo em lote 128 entradas de uma só vez. A única transformação de dados que fazemos em nossa inserção é transformar o texto do discurso em um embedding.
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)
Populando Seu Banco de Dados Vetorial
Com uma função que cria embeddings em lote e realiza inserções concluída, estamos prontos para popular o banco de dados. Para este exemplo, percorremos cada linha em nosso dataframe e adicionamos a uma lista de listas que usamos para agrupar nossos dados em lotes. Assim que atingimos o tamanho do nosso lote, chamamos a função embed_insert e redefinimos nosso lote.
Se houver dados restantes no lote de dados depois de terminarmos o loop, geramos embeddings e inserimos os dados restantes. Por fim, para concluir o preenchimento do nosso banco de dados vetorial, chamamos flush para garantir que o banco de dados seja atualizado e indexado.
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()
Busca Semântica em Discursos da Casa Branca com Base em Descrições
Digamos que eu esteja interessado em encontrar um discurso em que o presidente falou sobre o impacto da energia renovável no National Renewable Energy Lab (NREL) e um discurso em que a vice-presidente e o primeiro-ministro do Canadá falam. Posso encontrar os títulos dos discursos mais semelhantes feitos pelos membros da Casa Branca em 2021-2022 usando o banco de dados vetorial que acabamos de criar.
Podemos pesquisar em nosso banco de dados vetorial os discursos mais semelhantes às nossas descrições. Então, tudo o que precisamos fazer é converter a descrição em um embedding vetorial usando o mesmo modelo que usamos para obter os embeddings dos discursos e, em seguida, pesquisar no banco de dados vetorial.
Depois de convertermos as descrições em uma incorporação vetorial, usamos a função search em nossa coleção. Passamos as incorporações como os dados de busca, passamos o campo que estamos procurando, adicionamos alguns parâmetros sobre como buscar, um limite para o número de resultados e o campo que queremos retornar. Neste exemplo, os parâmetros de busca que precisamos passar são o tipo de métrica, que deve ser o mesmo tipo que usamos ao criar o índice (norma L2), e o número de clusters que queremos pesquisar (definindo nprobe como 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()
Quando buscamos pelas frases neste exemplo, esperamos ver uma saída como a imagem abaixo. Esta foi uma busca bem-sucedida porque os títulos são o que esperamos ver. A primeira descrição retorna o título de um discurso feito pelo Presidente Biden no NREL, e a segunda descrição retorna um título que reflete um discurso feito pela Vice-Presidente Harris e pelo Primeiro-Ministro Trudeau.
Resumo
Neste tutorial, aprendemos como usar um banco de dados vetorial para buscar semanticamente em discursos feitos pela administração Biden antes das eleições de meio de mandato de 2022. A busca semântica nos permite pegar um resumo e procurar texto semanticamente semelhante, não apenas texto sintaticamente semelhante. Isso nos permite buscar uma descrição geral de um discurso em vez de procurar um discurso com base em frases ou citações específicas. Para a maioria de nós, isso torna muito mais fácil encontrar discursos que seriam do nosso interesse.
Continue lendo
Announcing VDBBench 1.0: Open-Source VectorDB Benchmarking with Your Real-World Production Workloads
Discover VDBBench 1.0, an open-source tool for benchmarking vector databases with real-world production data, streaming ingestion, and concurrent workloads.
Legal Document Analysis: Harnessing Zilliz Cloud's Semantic Search and RAG for Legal Insights
Enhance legal document analysis with Zilliz Cloud’s Semantic Search and RAG. Improve accuracy, efficiency, and scalability for contracts, case law, and compliance.
Introducing DeepSearcher: A Local Open Source Deep Research
In contrast to OpenAI’s Deep Research, this example ran locally, using only open-source models and tools like Milvus and LangChain.