Busca de Similaridade Vetorial de Artigos Científicos do ArXiv com Milvus 2.1
Introdução
Uma das melhores maneiras de aprender qualquer tópico mais recente de Ciência de Dados é lendo artigos de pesquisa open-source no arxiv.org. No entanto, o vasto número de artigos de pesquisa pode ser esmagador até mesmo para o pesquisador mais experiente classificar. Ferramentas como connected papers podem ajudar, mas elas medem a similaridade com base nas citações e bibliografia compartilhadas entre os artigos, não no significado semântico do texto nesses documentos.
Neste post, eu me propus a criar um mecanismo de busca por similaridade semântica, que recebe um único artigo de “consulta” como entrada e usa NLP de última geração para encontrar os top-K artigos mais similares do corpus do arxiv de aproximadamente 640K artigos de ciência da computação! A busca roda com latência <50ms em um único laptop! Especificamente, neste post, vou abordar
- Configurar o ambiente e baixar os dados do arXiv do Kaggle
- Carregar os dados no Python usando Dask
- Implementar uma aplicação de busca por similaridade semântica de artigos científicos usando o banco de dados vetorial Milvus
As técnicas usadas neste post podem ser usadas como modelo para criar qualquer mecanismo de busca por similaridade semântica em NLP, não apenas para artigos científicos. A única diferença seria o modelo pré-treinado usado.
Para este post, usaremos o arXiv Dataset from Kaggle, que os autores lançaram sob a licença CC0: Public Domain.
Eu descrevi as considerações de busca por similaridade vetorial em escala de produção no meu post anterior. Todas essas considerações também se aplicam a este projeto. O banco de dados vetorial Milvus é tão bem projetado que muitas das etapas são exatamente as mesmas e são replicadas aqui apenas por completude.
Configure o ambiente e baixe os dados do arxiv do Kaggle.
A Cornel University carregou todo o corpus do arXiv em um Kaggle dataset e o licenciou sob a licença CC0: Public Domain. Podemos baixar diretamente o conjunto de dados usando a API do Kaggle. Se você ainda não fez isso, configure a API do Kaggle no seu sistema seguindo estas instruções.
Usaremos um ambiente conda para este post chamado semantic_similarity. Se você não instalou o conda no seu sistema, pode fazer isso instalando o mini forge open-source a partir de seu repositório do GitHub. As etapas abaixo criam os diretórios necessários e o ambiente conda, instalam as bibliotecas Python necessárias e baixam o arxiv dataset from Kaggle.
# Create the necessary directories
mkdir -p semantic_similarity/notebooks semantic_similarity/data semantic_similarity/milvus
# CD into the data directory
cd semantic_similarity/data
# Create and activate a conda environment
conda create -n semantic_similarity python=3.9
conda activate semantic_similarity
## Create Virtual Environment using venv if not using conda
# python -m venv semantic_similarity
# source semantic_similarity/bin/activate
# Pip install the necessary libraries
pip install jupyterlab kaggle matplotlib scikit-learn tqdm ipywidgets
pip install "dask[complete]" sentence-transformers
pip install pandas pyarrow pymilvus protobuf==3.20.0
# Download data using the kaggle API
kaggle datasets download -d Cornell-University/arxiv
# Unzip the data into the local directory
unzip arxiv.zip
# Delete the Zip file
rm arxiv.zip
Carregue os dados no Python usando Dask
Os dados que baixamos do Kaggle são um arquivo JSON de 3,3 GB contendo cerca de 2 milhões de artigos! Para processar eficientemente um conjunto de dados tão grande, não é uma boa ideia carregar o conjunto de dados inteiro na memória usando pandas. Em vez disso, podemos usar o Dask para dividir os dados em várias partições e carregar apenas algumas partições na memória em um determinado momento.
Dask
Dask é uma biblioteca de código aberto que nos permite aplicar computação paralela com uma API semelhante à do pandas facilmente. É simples de configurar na sua máquina local executando,pip install dask[complete] como mostrado na seção de configuração. Vamos começar importando primeiro as bibliotecas necessárias.
import dask.bag as db
import json
from datetime import datetime
import time
data_path = '../data/arxiv-metadata-oai-snapshot.json'
Usaremos dois componentes do Dask para processar eficientemente o grande arquivo JSON do arxiv.
- Dask Bag: Ele nos permite carregar o arquivo JSON em blocos de tamanho fixo e executar algumas funções de pré-processamento em cada linha de dados.
- Dask DataFrame: Podemos converter um dask bag em um dask dataframe para obter acesso a APIs semelhantes às do pandas
Passo 1: Carregar o arquivo JSON em um Dask bag
Vamos carregar o arquivo JSON em um dask bag onde cada bloco tem 10 MB de tamanho. Você pode ajustar o argumento blocksize para controlar o tamanho que deseja que cada bloco tenha. Em seguida, aplicamos a função json.loads a cada linha do dask bag usando a função .map() para analisar a string JSON em um dicionário Python.
# Read the file in blocks of 10MB and parse the JSON.
papers_db = db.read_text(data_path, blocksize="10MB").map(json.loads)
# Print the first row
papers_db.take(1)
Imagem do Autor
Passo 2: Escrever funções auxiliares de pré-processamento
A partir da saída impressa, vemos que cada linha contém vários metadados relacionados a um artigo. Vamos escrever três funções auxiliares para nos ajudar a pré-processar o conjunto de dados.
v1_date(): Esta função serve para extrair a data em que os autores enviaram a primeira versão do artigo para o arXiv. Converteremos a data para tempo UNIX e a armazenaremos como um novo campo nessa linha.text_col(): Esta função serve para combinar os campos “title” e “abstract” usando um token “[SEP]” para que possamos alimentar esses textos no modelo de embeddings SPECTRE. Falaremos mais sobre o SPECTRE na próxima seção.filters(): Esta função mantém apenas as linhas que atendem a alguns critérios, como comprimento máximo de texto em várias colunas e artigos na categoria Ciência da Computação.
def v1_date(row):
"""
For each row in the dask bag,
find the date of the first version of the paper
and add it to the row as a new column
Args:
row: a row of the dask bag
Returns:
A row of the dask bag with added "unix_time" column
"""
versions = row["versions"]
date = None
for version in versions:
if version["version"] == "v1":
date = datetime.strptime(version["created"], "%a, %d %b %Y %H:%M:%S %Z")
date = int(time.mktime(date.timetuple()))
row["unix_time"] = date
return row
def text_col(row):
"""
It takes a row of a dataframe, adds a new column called 'text'
that is the concatenation of the 'title' and 'abstract' columns
Args:
row: the row of the dataframe
Returns:
A row with the text column added.
"""
row["text"] = row["title"] + "[SEP]" + row["abstract"]
return row
def filters(row):
"""
Para cada linha no dask bag, mantenha a linha somente se ela atender aos critérios de filtro
Args:
row: a linha do dataframe
Returns:
Máscara booleana
"""
return ((len(row["id"])<16) and
(len(row["categories"])<200) and
(len(row["title"])<4096) and
(len(row["abstract"])<65535) and
("cs." in row["categories"]) # Manter apenas artigos de CS
)
Etapa 3: Execute as funções auxiliares de pré-processamento no Dask bag
Podemos usar facilmente as funções .map() e .filter() para executar as funções auxiliares em cada linha do Dask bag, como mostrado abaixo. Como o Dask suporta encadeamento de métodos, aproveitamos esta oportunidade para manter apenas algumas colunas essenciais em nosso Dask bag e descartar o restante.
# Especifique as colunas a manter na tabela final
cols_to_keep = ["id", "categories", "title", "abstract", "unix_time", "text"]
# Aplique o pré-processamento
papers_db = (
papers_db.map(lambda row: v1_date(row))
.map(lambda row: text_col(row))
.map(
lambda row: {
key: value
for key, value in row.items()
if key in cols_to_keep
}
)
.filter(filters)
)
# Imprima a primeira linha
papers_db.take(1)
Imagem do Autor
Etapa 4: Converta o Dask Bag em um Dask DataFrame
A etapa final do carregamento de dados é converter o Dask Bag em um Dask Dataframe para usar APIs semelhantes às do pandas em cada bloco ou partição dos dados.
# Converta o Dask Bag em um Dask Dataframe
schema = {
"id": str,
"title": str,
"categories": str,
"abstract": str,
"unix_time": int,
"text": str,
}
papers_df = papers_db.to_dataframe(meta=schema)
# Exiba as primeiras 5 linhas
papers_df.head()
Imagem do Autor
Implementando uma aplicação de busca por similaridade semântica de artigos científicos usando o banco de dados vetorial Milvus
Milvus é um dos bancos de dados vetoriais de código aberto mais populares, criado para buscas por similaridade altamente escaláveis e extremamente rápidas. Usaremos o Milvus Standalone para esta publicação, pois estamos executando o Milvus apenas em nossa máquina local.
Etapa 1: Instale o banco de dados vetorial Milvus localmente
Instalar o banco de dados vetorial Milvus é muito fácil usando Docker, então primeiro precisamos instalar o Docker e o Docker Compose. Depois, tudo o que precisamos fazer é baixar um docker-compose.yml e iniciar os contêineres docker, como mostrado no trecho de código abaixo! O site milvus.io fornece muitas outras opções para instalar tanto o Milvus standalone quanto o Milvus Cluster; confira se você precisar instalá-lo em um cluster Kubernetes ou instalá-lo offline.
# Entre no diretório milvus
cd semantic_similarity/milvus
# Baixe a versão Standalone do docker compose do Milvus
wget https://github.com/milvus-io/milvus/releases/download/v2.1.0/milvus-standalone-docker-compose.yml -O ./docker-compose.yml
# Execute o contêiner docker do servidor Milvus localmente
sudo docker-compose up -d
Etapa 2: Crie uma coleção Milvus
Agora que temos o servidor do banco de dados vetorial Milvus em execução em nossa máquina local, podemos interagir com ele usando a biblioteca pymilvus. Primeiro, vamos importar os módulos necessários e conectar ao servidor Milvus em execução em localhost. Fique à vontade para alterar os parâmetros alias e collection_name. O modelo que usamos para converter nosso texto em embeddings determina o valor do parâmetro emb_dim. No caso do SPECTRE, os embeddings têm 768 dimensões.
# Certifique-se de que um servidor Milvus já esteja em execução
from pymilvus import connections, utility
from pymilvus import Collection, CollectionSchema, FieldSchema, DataType
# Conectar ao servidor Milvus
connections.connect(alias="default", host="localhost", port="19530")
# Nome da coleção
collection_name = "arxiv"
# Tamanho do embedding
emb_dim = 768
# # Verificar se há uma coleção existente e descartá-la se existir
# if utility.has_collection(collection_name):
# print(utility.list_collections())
# utility.drop_collection(collection_name)
Opcionalmente, você pode verificar se a coleção especificada por collection_name já está presente no seu servidor Milvus. Para este exemplo, se a coleção já estiver disponível, eu a excluo. Mas em um servidor de produção, você não faria isso e, em vez disso, pularia o código de criação da coleção abaixo.
Uma coleção Milvus é análoga a uma tabela em um banco de dados tradicional. Para criar uma coleção para armazenar dados, primeiro precisamos especificar o schema da coleção. Neste exemplo, estamos aproveitando a capacidade do Milvus 2.1 de armazenar índice e campos de string para armazenar todos os metadados necessários relacionados a cada artigo. A chave primária idx e outros campos categories, title, abstract têm datatype VARCHAR com comprimentos máximos razoáveis, enquanto o embedding é um campo FLOAT_VECTORcontendo os embeddings de dimensão emb_dim. O Milvus oferece suporte a uma ampla variedade de tipos de dados, conforme mostrado em nossa página de documentação.
# Criar um schema para a coleção
idx = FieldSchema(name="id", dtype=DataType.VARCHAR, is_primary=True, max_length=16)
categories = FieldSchema(name="categories", dtype=DataType.VARCHAR, max_length=200)
title = FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=4096)
abstract = FieldSchema(name="abstract", dtype=DataType.VARCHAR, max_length=65535)
unix_time = FieldSchema(name="unix_time", dtype=DataType.INT64)
embedding = FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=emb_dim)
# Campos na coleção
fields = [idx, categories, title, abstract, unix_time, embedding]
schema = CollectionSchema(
fields=fields, description="Similaridade Semântica de Artigos Científicos"
)
# Criar uma coleção com o schema
collection = Collection(
name=collection_name, schema=schema, using="default", shards_num=10
)
Depois que uma coleção tiver sido criada, agora estamos prontos para carregar nossos textos e vetores nela.
Etapa 3: Iterar pelas partições do nosso dataframe Dask, incorporar os textos usando SPECTER e carregá-los no banco de dados vetorial Milvus
Primeiro, precisamos converter os textos no dataframe Dask em um vetor de embedding para executar uma busca por similaridade semântica. Minha publicação abaixo compartilha como podemos converter textos em embeddings. Em particular, usaremos um modelo SBERT Bi-Encoder chamado SPECTRE para converter artigos científicos em embeddings.
SPECTER [artigo] [Github]: Scientific Paper Embeddings using Citation-informed TransformERs é um modelo para converter artigos científicos em embeddings.
- Os textos de Título e Resumo de cada artigo são concatenados com o token [SEP] e convertidos em embeddings usando o token [CLS] de um modelo Transformer pré-treinado (SciBERT).
- Use citações como um sinal proxy para a relação entre documentos. Se um artigo cita outro, podemos inferir que ambos estão relacionados.
- Objetivo de treinamento com perda tripla: Treinamos o modelo Transformer, de modo que artigos com citações compartilhadas fiquem mais próximos no espaço de embeddings.
- Em outras palavras, um artigo Positivo é um artigo citado no artigo Consulta, enquanto um artigo Negativo é um artigo não citado pelo artigo Consulta. Negativos amostrados aleatoriamente são negativos “fáceis”.
- Para melhorar o desempenho, criamos negativos “difíceis” usando artigos que NÃO são citados pelo artigo Consulta, mas SÃO citados pelo artigo Positivo.
- Precisamos apenas do Título e do Resumo durante a inferência. Nenhuma citação é necessária, então o SPECTER pode produzir embeddings até mesmo para novos artigos que ainda não têm nenhuma citação!
- O SPECTER oferece excelente desempenho (melhor que o SciBERT) em classificação de tópicos, previsão de citações e recomendação de Artigos Científicos.
Imagem do Autor usando capturas de tela do artigo SPECTER de código aberto
Usar o modelo SPECTRE pré-treinado é simples com a biblioteca Sentence Transformer. Podemos baixar o modelo pré-treinado com apenas uma linha de código, como mostrado abaixo. Também escrevemos uma função auxiliar simples para converter uma coluna inteira de textos da partição do dataframe Dask em embeddings.
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
# Scientific Papers SBERT Model
model = SentenceTransformer('allenai-specter')
def emb_gen(partition):
return model.encode(partition['text']).tolist()
Precisamos iterar sobre as partições do dataframe Dask para carregar os dados em nossa coleção Milvus. Durante cada iteração, carregamos apenas as linhas dessa partição na memória e adicionamos os dados das colunas de metadados a uma variável data. Podemos usar a API dask .map_partitions() para aplicar a geração de embeddings a cada linha na partição e anexar os resultados de volta à mesma variável data. Por fim, podemos carregar os dados no Milvus com collection.insert.
# Initialize
collection = Collection(collection_name)
for partition in tqdm(range(papers_df.npartitions)):
# Get the dask dataframe for the partition
subset_df = papers_df.get_partition(partition)
# Check if dataframe is empty
if len(subset_df.index) != 0:
# Metadata
data = [
subset_df[col].values.compute().tolist()
for col in ["id", "categories", "title", "abstract", "unix_time"]
]
# Embeddings
data += [
subset_df
.map_partitions(emb_gen)
.compute()[0]
]
# Insert data
collection.insert(data)
Observe que a ordem das colunas adicionadas à variável data deve seguir a mesma ordem da variável fields que definimos durante a criação do esquema!
Etapa 4: Crie um índice de Vizinhos Mais Próximos Aproximados (ANN) nos dados carregados
Depois de inserirmos todos os embeddings no banco de dados vetorial Milvus, precisamos criar um índice ANN para acelerar a busca. Neste exemplo, estou usando o tipo de índice HNSW, um dos índices ANN mais rápidos e precisos. Consulte a documentação do Milvus para obter mais informações sobre o índice HNSW e seus parâmetros.
# Add an ANN index to the collection
index_params = {
"metric_type": "L2",
"index_type": "HNSW",
"params": {"efConstruction": 128, "M": 8},
}
collection.create_index(field_name="embedding", index_params=index_params)
Etapa 5: Execute suas consultas de Busca por Similaridade Vetorial!
Finalmente, os dados em nossa coleção do Milvus estão prontos para serem consultados. Primeiro, precisamos carregar a coleção na memória para executar consultas nela.
# Load the collection into memory
collection = Collection(collection_name)
collection.load()
Em seguida, criei uma função auxiliar simples que recebe um query_text, converte-o para o embedding SPECTRE, executa uma busca ANN em toda a coleção do Milvus e imprime os resultados. Podemos controlar a qualidade e a velocidade da busca usando os search_params descritos na página de documentação do HNSW.
def query_and_display(query_text, collection, num_results=10):
# Embed the Query Text
query_emb = [model.encode(query_text)]
# Search Params
search_params = {"metric_type": "L2", "params": {"ef": 128}}
# Search
query_start = datetime.now()
results = collection.search(
data=query_emb,
anns_field="embedding",
param=search_params,
limit=num_results,
expr=None,
output_fields=["title", "abstract"],
)
query_end = datetime.now()
# Print Results
print(f"Query Speed: {(query_end - query_start).total_seconds():.2f} s")
print("Results:")
for res in results[0]:
title = res.entity.get("title").replace("\n ", "")
print(f"➡️ ID: {res.id}. L2 Distance: {res.distance:.2f}")
print(f"Title: {title}")
print(f"Abstract: {res.entity.get('abstract')}")
Agora podemos usar a função auxiliar com apenas uma linha de código para executar uma busca semântica de artigos do arXiv em todos os ~640 mil artigos de Ciência da Computação armazenados em nossa coleção do Milvus. Por exemplo, estou procurando alguns artigos semelhantes ao artigo SimCSE que discuti em detalhes em minha postagem anterior. Os 10 principais resultados são bastante relevantes para minha consulta de busca, pois estão, em sua maioria, relacionados ao aprendizado contrastivo de embeddings de sentenças! É ainda mais impressionante que toda a busca tenha levado apenas 30 ms rodando no meu laptop, o que está bem dentro dos requisitos típicos de uso para a maioria das aplicações!
# Query for papers that are similar to the SimCSE paper
title = "SimCSE: Simple Contrastive Learning of Sentence Embeddings"
abstract = """This paper presents SimCSE, a simple contrastive learning framework that greatly advances state-of-the-art sentence embeddings. We first describe an unsupervised approach, which takes an input sentence and predicts itself in a contrastive objective, with only standard dropout used as noise. This simple method works surprisingly well, performing on par with previous supervised counterparts. We find that dropout acts as minimal data augmentation, and removing it leads to a representation collapse. Then, we propose a supervised approach, which incorporates annotated pairs from natural language inference datasets into our contrastive learning framework by using "entailment" pairs as positives and "contradiction" pairs as hard negatives. We evaluate SimCSE on standard semantic textual similarity (STS) tasks, and our unsupervised and supervised models using BERT base achieve an average of 76.3% and 81.6% Spearman's correlation respectively, a 4.2% and 2.2% improvement compared to the previous best results. We also show -- both theoretically and empirically -- that the contrastive learning objective regularizes pre-trained embeddings' anisotropic space to be more uniform, and it better aligns positive pairs when supervised signals are available."""
query_text = f"{title}[SEP]{abstract}"
query_and_display(query_text, collection, num_results=10)
Imagem do Autor
Se não precisarmos executar mais consultas, podemos liberar a coleção para liberar a memória da nossa máquina. Remover uma coleção da memória não causa perda de dados, pois ela ainda está armazenada em nosso disco e pode ser carregada novamente quando necessário.
# Release the collection from memory when it's not needed anymore
collection.release()
Se você quiser parar o servidor Milvus e excluir todos os dados do disco, pode seguir as instruções para parar o Milvus. Cuidado! Esta operação é irreversível e excluirá todos os dados no seu cluster Milvus.
Conclusão
Neste post, implementamos um serviço de Busca Semântica ultrarrápida e escalável de Artigos Científicos usando embeddings SPECTRE e o banco de dados vetorial Milvus em algumas etapas simples. Essa abordagem é escalável em produção para centenas de milhões ou até bilhões de vetores. Testamos a busca usando uma consulta de artigo de exemplo que retornou os 10 principais resultados em apenas 30 ms! A reputação do Milvus como um banco de dados de busca por similaridade vetorial altamente escalável e extremamente rápido é bem merecida!
Para mais inspiração sobre as aplicações do Milvus, acesse as Milvus demonstrações de banco de dados vetorial e demos e Bootcamp.
Continue lendo
What Is a Vector Lakebase?
A Vector Lakebase is a unified, lake-native data architecture for AI that combines vector-database-grade serving with open lake storage, reusable lake-level indexes, and a shared semantic layer.
Build for the Boom: Why AI Agent Startups Should Build Scalable Infrastructure Early
Explore strategies for developing AI agents that can handle rapid growth. Don't let inadequate systems undermine your success during critical breakthrough moments.
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.