Similarity Search 101: Crea una base de datos vectorial para la búsqueda de similitud de texto con Chroma y Langchain

Similarity Search 101: Crea una base de datos vectorial para la búsqueda de similitud de texto con Chroma y Langchain

Guía paso a paso para crear una base de datos vectorial para la búsqueda de similitud de texto usando Chroma y Langchain

¿Qué queremos lograr?

Crear un sistema de búsqueda de similitud de texto. El sistema va a recibir un texto y va a devolver los textos más similares a ese texto.

¿Cómo lo vamos a construir?

Así va la lógica:

  • Para almacenar el texto y poder encontrar los textos similares, usaremos una base de datos vectorial, para este proyecto se usará Chroma (Open Source y se puede instalar en local).

  • La base de datos vectorial almacena los textos como embeddings(representación de texto como vectores aka arrays de números). La transformación de texto a embeddings se hace con un modelo de lenguaje, en este caso usaremos un modelo de lenguaje de Hugging Face.

  • La clase que transforma texto a embeddings usa un tipo de dato llamado Document, que es un objeto que contiene el texto y metadata del texto.

El flujo de trabajo será Documentos -> Embeddings -> Base de datos vectorial -> Búsqueda de similitud

Costos para este proyecto: Ninguno, usaremos librerías y ejecutaremos todo en local.
Con esto en mente, a codear!

Code

Documentos

Instalar dependencias
Psdt: El proyecto usa Poetry como manejador de dependencias.

poetry add langchain

El texto para este proyecto es este blog de Hugging Face sobre embeddings: https://huggingface.co/blog/getting-started-with-embeddings.

De su repo de Github descargué el archivo markdown y lo guardé en la carpeta data de este proyecto.

Como es un archivo markdown importamos UnstructuredMarkdownLoader para crear el Document.

import os
from langchain.document_loaders import UnstructuredMarkdownLoader

# Para usar UnstructuredMarkdownLoader, se debe instalar las librerías unstructured y markdown -> poetry add unstructured markdown

proyect_root = os.path.abspath(os.path.join(os.getcwd(), os.pardir))
py
file_path = f"{proyect_root}/data/getting-started-with-embeddings.md"
loader = UnstructuredMarkdownLoader(file_path)
document = loader.load() # devuelve una lista de Documents, en este caso solo hay uno

Examinando el Document

document[0].page_content[:50] # primeros 1000 caracteres del documento

Se obtiene

"title: 'Getting Started With Embeddings'\nthumbnail"

Para obtener la cantidad de caracteres del documento

# Tamaño del documento
len(document[0].page_content)

El Documento tiene 11997 caracteres, para que la base de datos vectorial devuelva el fragmento de texto más similar a la búsqueda se trabaja con fragmentos de texto más pequeños.
Langchain tiene un separador de texto llamado RecursiveCharacterTextSplitter.

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 250, # tamaño del chunk, tomar en cuenta el número de tokens que el modelo de lengua puede manejar
    chunk_overlap  = 25, # entre 10% y 20% del tamaño del chunk
    length_function = len # función que se aplica a cada chunk para obtener su longitud
)

chunks = text_splitter.split_documents(document)

Visualizar datos de chunks (lista de Documents)

print(f"El documento se dividió en {len(chunks)} chunks")
print(f"El primer chunk tiene {len(chunks[0].page_content)} caracteres")

Así descubrimos que:

  • El documento se dividió en 70 chunks

  • El primer chunk tiene 229 caracteres

Embeddings

Para iniciar con los embeddings, instalar la librería Sentence Transformers:

poetry add sentence-transformers

Usaremos un modelo de embeddings de HuggingFace, aquí encuentras todos los disponibles: https://python.langchain.com/docs/integrations/text_embedding

from langchain.embeddings.huggingface import HuggingFaceEmbeddings

model_name = "sentence-transformers/all-MiniLM-L6-v2"
model = HuggingFaceEmbeddings(model_name=model_name)

*Más modelos en: https://huggingface.co/sentence-transformers

Obtenemos información del modelo con

model

Y se ve así

HuggingFaceEmbeddings(client=SentenceTransformer(
  (0): Transformer({'max_seq_length': 256, 'do_lower_case': False}) with Transformer model: BertModel 
  (1): Pooling({'word_embedding_dimension': 384, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})
  (2): Normalize()
), model_name='sentence-transformers/all-MiniLM-L6-v2', cache_folder=None, model_kwargs={}, encode_kwargs={}, multi_process=False)

El modelo tiene un máximo de 256 caracteres, así que la longitud de los Documentos debe ser menor a ese número.

Y en Polling vemos la dimensionalidad de los embeddings resultantes, en este caso 384.

Ejemplo de embedding

embedding_example = model.embed_documents([chunks[0].page_content])[0]

print(f"El embedding tiene {len(embedding_example)} dimensiones")

*El embedding tiene 384 dimensiones

Base de datos vectorial

Con poetry instalar chromadb -> poetry add chromadb

La siguiente función devuelve un objeto Chroma que se puede utilizar para la búsqueda de similitud.

Tiene los siguentes parámetros:

  • documentos (list): Una lista de documentos para usar en la creación del objeto Chroma.

  • modelo_embedding (function): Una función que toma un documento como entrada y devuelve su embedding.

  • ruta (str): La ruta al directorio donde se debe persistir el objeto Chroma.

  • recrear (bool): Si es True, recrea el objeto Chroma desde cero. Si es False, carga el objeto Chroma existente desde el disco.

Retorna: Un objeto Chroma que se puede utilizar para la búsqueda de similitud.

from langchain.vectorstores import Chroma

def get_chroma_db(documentos, modelo_embedding, ruta, recrear):
  if recrear:
    # Si recrear es True, crea un nuevo objeto Chroma desde cero
    chroma = Chroma.from_documents(
      documents=documentos,
      embedding=modelo_embedding,
      persist_directory=ruta,
    )

    chroma.persist()
    return chroma
  else:
    # Si recrear es False, carga el objeto Chroma existente desde el disco
    chroma = Chroma(
      persist_directory=ruta,
      embedding_function=modelo_embedding,
    )
    return chroma

Usamos la función para crear la bd vectorial

CHROMA_INDEX_NAME = "similarity-search-101-chroma" #nombre de la BD

vector_store = get_chroma_db(
    documentos=chunks,
    modelo_embedding=model,
    ruta=f"{proyect_root}/db/{CHROMA_INDEX_NAME}",
    recrear=True 
)

Ahora podemos usar la BD vectorial

query = "What is a embedding model?"

docs = vector_store.similarity_search_with_score(query, k=5)

k es el número de documentos más similares que se devuelven

Para conocer los resultados ejecutamos

# número de documentos devueltos
print(f"Se devolvieron {len(docs)} documentos")
# El primer documento devuelto: score y contenido
print(f"El primer documento devuelto tiene un score de {docs[0][1]} y su contenido es:")
print(docs[0][0].page_content)

Retornando...

Se devolvieron 5 documentos
El primer documento devuelto tiene un score de 0.424575540971641 y su contenido es:
An embedding is a numerical representation of a piece of information, for example, text, documents, images, audio, etc. The representation captures the semantic meaning of what is being embedded, making it robust for many industry applications.

Y viendo el contenido de la respuesta, sí es un texto que ayudaría a un modelo de lenguaje a como gpt-3 a responder preguntas la pregunta "What is a embedding model?".

Puedes encontrar el código completo en: https://github.com/nestorxyz/ai-learning-notebooks/blob/main/notebooks/similarity-search-101.ipynb

Referencias