Implementando Busca Semântica no ChromaDB
Crie seu embedding, armazene no banco e faça buscas por similaridade
Me siga no LinkedIn | Apoie a Newsletter | Curso: Converse com seus documentos
As aplicações de IA Generativa dependem fundamentalmente de embeddings vetoriais. Embeddings transformam conteúdo digital (que pode ser textual, de imagens, em voz) em um vetor numérico.
Mais importante, embeddings tentam encapsular o significado do texto, possibilitando assim que modelos de IA compreendam quais textos são semanticamente similares. Dessa forma, uma das funções do embeddings, é ajudar na busca de embeddings similares.
O processo de extrair os textos mais semelhantes em uma base de dados é conhecido como busca por vetores ou busca vetorial. Nesse texto vamos implementar um mecanismo de busca usando o ChromaDB.
Vetores e suas dimensões
Em termos simples, um vetor é uma lista de números que pode representar um ponto no espaço.
Por exemplo, em um espaço 2D, um vetor [2,3]
representa um ponto que está 2 unidades ao longo do eixo x e 3 unidades ao longo do eixo y. Em um espaço 3D, um vetor [2,3,4] adiciona uma dimensão adicional, o eixo z.
A beleza dos vetores é que eles podem ter muitas mais dimensões do que conseguimos visualizar. É comum ter vetores de embeddings com centenas, ou até milhares de dimensões.
Vetores de embeddings
A palavra “embedding” é apenas uma maneira sofisticada de dizer “representação” ou “tradução”.
Quando falamos em “embeddings vetoriais”, estamos nos referindo a representações de dados complexos, como textos ou imagens, como vetores.
Por exemplo, pense na palavra gato. Podemos representá-la como um ponto em um espaço multidimensional usando um vetor.
[0.8108, 0.6671, 0.5565, 0.5449, 0.4466]
Essa representação pode capturar a essência ou significado da palavra em relação a outras palavras.
Palavras com significados semelhantes estariam mais próximas neste espaço, enquanto as diferentes estariam mais distantes.
Por que usar um vetor de embeddings?
Há algumas vantanges em se utilizar vetores de embeddings.
Primeiro, os vetores de embeddings nos permitem converter diversas formas de dados em um formato comum (vetores) que os algoritmos de aprendizado de máquina podem entender e processar. Essa representação compacta captura características essenciais dos dados, facilitando a manipulação e compreensão por modelos de aprendizado de máquina.
Ao fazer essa tradução em embedding, os embeddings tendem a capturar relações semânticas entre dados. Por exemplo, palavras com significados semelhantes têm representações vetoriais mais próximas entre si. Isso facilita a realização de operações como identificação de sinônimos, análise de similaridade entre documentos ou até mesmo agrupamento de informações relacionadas.
Mas essas operações de similaridade só podem ser implementadas uma vez que a representação vetorial permite a realização de operações matemáticas nos dados. Por exemplo, pode-se calcular a distância entre dois vetores, indicando assim a sua similaridade.
Como criar vetores de embeddings?
A verdade é que hoje em dia existem diversas algoritmos e modelos, incorporados nas mais diversas bibliotecas, que são capazes de gerar embeddings. Existe, inclusive, um benchmark que tenta avaliar a qualidade dos embeddings criados por esses algoritmos/modelos, chamado MTEB. Esse benchmark avaliar os embeddings em critérios como por exemplo: classificação, ranqueamento, ou a similaridade de sentenças.
O modelo SetenceBert, implementado na biblioteca open source sbert, é uma boa opção, uma vez que tem fácil utilização, e seu desempenho ainda faz frente a diversos modelos proprietários. Para instalar a biblioteca, usamos:
pip install -U sentence-transformers
Usamos o seguinte comando para criar um embedding a partir de uma string, usando sbert:
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("all-MiniLM-L6-v2")
embedding = model.encode("gato")
E pronto, nossa variável embedding
agora armazena a representação vetorial da palavra “gato”.
Algumas características de um embedding:
Os modelos de embedding operam com um tamanho máximo de tokens. Ou seja, só podem processar uma determinada quantidade de conteúdo. De acordo com o MTEB, essa quantidade pode variar de 512 tokens até 8k tokens.
Além disso, os novos modelos de embedding fornecem um sofisticado mecanismo de compressão que garante que o tamanho final do vetor de embedding será o mesmo, independente da quantidade de tokens de entrada. Ou seja, o vetor da palavra “gato” terá o mesma quantidade do elementos do vetor da palavra “não atire o pau no gato”, que terá a mesma quantidade de elementos de um livro sobre gatos.
embedding = model.encode("gato")
print(len(embedding)) # 384
print(embedding.shape) # (384,)
embedding = model.encode("não atire o pau no gato")
print(len(embedding)) # 384
print(embedding.shape) # (384,)
Busca no ChromaDB
Uma vez que nossos embeddings foram criados, é uma prática comum inseri-los em um banco de dados com suporte a vetores multidimensionais. Há diversas opções no mercado, como o ChromaDB, Pinecole e até o Postgres, através da extensão PGVector.
Nesse texto, vamos fazer um simples cadastro usando o ChromaDB. Para instalar, basta usar o comando
!pip install chromadb
Para instanciar um cliente e criar uma base, podemos usar o comando:
import chromadb
chroma_client = chromadb.Client()
Felizmente, o chromadb já possui uma integração com diversos modelos de embedding. Como usuário, basta informar qual modelo deseja-se instanciar no chromadb.
No caso, caso queira-se utilizar o mesmo sbert, que usamos anteriormente, podemos utilizar a classe SentenceTransformerEmbeddingFunction
, que é um wrapper do modelo.
from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
model = SentenceTransformerEmbeddingFunction()
print(model("gato"))
Em seguida, precisamos criar a base de dados (chamada gentilmente de ml4se
) e inserir os embeddings, junto do conteúdo original, através do método add
.
chroma_collection = chroma_client.create_collection("ml4se", embedding_function=model)
docs = [
"o rato roeu a roupa do rei de roma",
"o sapo não lava o pé, não lava porque não quer",
"Gosto muito de te ver, leãozinho",
"machine learning é maneiro",
"adoro ler os textos da newsletter ml4se",
"o cachorro late quando faz au au",
"Foi na loja do Mestre André que eu comprei um pianinho"
]
ids = [str(i) for i in range(len(docs))]
chroma_collection.add(ids=ids, documents=docs)
Há dois pontos que precisam de atenção no código acima.
O primeiro é que a lista de docs, de forma proposital, conta com vários itens relacionados a animais.
Segundo, perceba, no entanto, que não foi necessário explicitamente fazer uma transformação de texto para embedding. Ao criar a base de dados, passamos a função de embedding que desejamos utilizar (embedding_function=model
), e todo o resto é abstraído pelo banco de dados.
Podemos ainda verificar quantos registros foram inseridos na base, usando o método count
.
chroma_collection.count() # retorna 3 registros
Por fim, podemos finalmente fazer uma busca de itens semelhantes no banco de dados, usando o método query
.
query = "cante uma música de animal"
results = chroma_collection.query(query_texts=[query], n_results=3)
retrieved_documents = results['documents'][0]
Essa consulta nos retorna três registros: 'o rato roeu a roupa do rei de roma', 'o cachorro late quando faz au au', 'Gosto muito de te ver, leãozinho'.
Perceba que a busca foi capaz de selecionar somente os registros que tinham alguma relação com algum animal. Curioso, não é? Isso acontece pois o banco de dados realizou um calculo para buscar registros do banco que são mais similares a query do usuário.
Se imprimirmos o objeto results
teremos outras informações, como os ids do usuário, e as distâncias de cada resultado, quando comparado a query do usuário.
Seria agora interessante experimentar com outras queries e dados!
O código deste texto está disponível neste link, licenciado como MIT.
Minicurso: converse com seus documentos
Se você é dev e quer aprender um pouco mais desse conjunto de técnicas para criar aplicações baseada em LLMs, considere comprar o minicurso “Converse com seus Documentos Usando LLM”, disponível por apenas R$ 99.