Como Implementar Busca Semântica com PyTorch e FAISS

No mundo da programação, a eficiência e a inovação andam de mãos dadas. Hoje, vou levá-lo através de uma jornada de como você pode transformar palavras em vetores semânticos e realizar buscas de conteúdo significativas utilizando PyTorch e FAISS. Este é um guia passo a passo do código que escrevi, então prepare-se para mergulhar em um exemplo prático que você pode executar em sua própria máquina.

Configuração Inicial

Antes de tudo, você precisa fazer o download do arquivo .csv que foi usado para esta busca. Segue o link do site kaggle

Agora, precisamos preparar nosso ambiente de trabalho. Isso significa verificar a disponibilidade de uma GPU para acelerar o processamento e definir um diretório para armazenar dados em cache:

import torch
import pandas as pd
import time

cache_dir = 'cache_dir'
start_time = time.time()

device = torch.device('cuda') if torch.cuda.is_available() else 'cpu'
num_gpus = torch.cuda.device_count()

Com o PyTorch, é fácil alternar entre CPUs e GPUs com base na disponibilidade. Em seguida, identificamos as GPUs disponíveis e imprimimos suas informações:

if num_gpus > 0:
    print(f"GPUs available: {num_gpus}")
    for i in range(num_gpus):
        gpu_name = torch.cuda.get_device_name(i)
        print(f"GPU {i}: {gpu_name}")
else:
    print("There is no GPU available. Using CPU")

Carregamento e Preparação dos Dados

A seguir, utilizamos o pandas para ler os dados de um arquivo CSV e adicionamos uma coluna única para identificação:

pdf = pd.read_csv(f"data/facebook_reviews.csv")
pdf['id'] = pdf.index

É importante notar que o caminho para o arquivo CSV deve ser atualizado para refletir a localização do arquivo em seu sistema.

Processamento de Linguagem Natural (NLP)

Com nossos dados prontos, é hora de transformar o texto em vetores semânticos. Aqui, carregamos um modelo pré-treinado do sentence_transformers e codificamos nossos textos:

from sentence_transformers import InputExample, SentenceTransformer

model = SentenceTransformer('all-MiniLM-L6-v2', cache_folder=cache_dir)
model.to(device)

pdf_subset = pdf.head(1000)
faiss_review_text_embedding = model.encode(pdf_subset.review_text.values.tolist())

O que é FAISS?

FAISS (Facebook AI Similarity Search) é uma biblioteca eficiente de busca por similaridade desenvolvida pelo Facebook AI Research (FAIR). Ela é projetada para ajudar na busca rápida de vetores de alta dimensão e suporta indexação de grandes conjuntos de dados. FAISS é especialmente útil em contextos onde se deseja comparar itens ou documentos com base em suas características codificadas em vetores, como em sistemas de recomendação, agrupamento de dados e busca semântica.

O poder do FAISS reside em sua capacidade de realizar buscas por vizinhança aproximada em larga escala. Ele faz isso de forma mais eficiente do que as buscas exatas tradicionais, que podem ser muito lentas quando lidamos com grandes volumes de dados. FAISS utiliza técnicas de quantização para reduzir o tamanho da memória necessária para armazenar os vetores e acelerar as operações de busca.

A biblioteca é altamente otimizada para arquiteturas de GPU, o que permite que operações de busca e agrupamento sejam realizadas com incrível velocidade. No entanto, ela também pode ser executada em CPUs para flexibilidade em ambientes onde GPUs não estão disponíveis ou são limitadas. Com uma API simples e uma série de índices pré-construídos, FAISS torna a implementação de sistemas de busca por similaridade uma tarefa acessível e menos onerosa.

Criação do Índice com FAISS

Com os vetores de texto em mãos, o próximo passo é criar um índice para permitir buscas eficientes. Utilizamos o FAISS para isso, normalizando os vetores e adicionando identificadores únicos:

import numpy as np
import faiss

pdf_to_index = pdf_subset.set_index(['id'], drop=False)
id_index = np.array(pdf_to_index.id.values).flatten().astype('int')

content_encoded_normalized = faiss_review_text_embedding.copy()
faiss.normalize_L2(content_encoded_normalized)

index_content = faiss.IndexIDMap(faiss.IndexFlatIP(len(faiss_review_text_embedding[0])))
index_content.add_with_ids(content_encoded_normalized, id_index)

Realizando a Busca

Definimos uma função de busca que leva em conta a similaridade dos vetores semânticos:

def search_content(query, pdf_to_index, k=5):
    query_vector = model.encode([query])
    faiss.normalize_L2(query_vector)

    top_k = index_content.search(query_vector, k)
    ids = top_k[1][0].tolist()
    similarities = top_k[0][0].tolist()

    results = pdf_to_index.loc[ids]
    results['similarities'] = similarities

    return results

print(search_content('annoying ads', pdf_to_index))

Conclusão

Por fim, avaliamos a performance do nosso script medindo o tempo de execução total. Isso nos dá uma noção da eficiência do nosso código:

end_time = time.time()
total_time = end_time - start_time
total_time = time.strftime("%H:%M:%S", time.gmtime(total_time))
print("Script execution time:", total_time)

Este post é uma introdução à busca semântica com ferramentas de ponta, gerado pelo ChatGPT (GPT-4) com báse no código que eu passei pra ele e pedi para transformar neste post.

Categorias: Inteligência Artificial
Tags: embedding, ia, vectors