Tutorial de la base de datos vectorial Chroma

En este tutorial vas a aprender qué es, cómo montar y cómo usar Chroma, una de las bases de datos vectoriales más conocidas e utilizadas hoy en día. Todo ello lo haremos en Python y con un enfoque práctico. ¿Te suena bien? ¡Vamos con ello!

Qué es y cómo funciona Chroma

Chroma es una base de datos vectorial. Como tal, su objetivo es que seas capaz de guardar vectores (generalmente embeddings) para después dotar de dicha información a otros modelos (como LLMs) o, simplemente, como herramienta de búsqueda.

Así pues, a nivel global, la forma de usar Chroma es la siguiente:

1. Crear nuestra colección, lo que es el equivalente a una tabla en una base de datos relacional. En este proceso, deberemos indicar qué modelo debe usar Chroma para convertir los textos en embeddings. Como veremos más adelante, hay varios disponibles.
2. Enviar a Chroma un texto que queramos que guarde, junto con los metadatos que queramos para el filtrado del texto. Cuando Chroma reciba el texto se encargará de convertirlo a embedding.
3. Consultar Chroma enviando un texto o un embedding, recibiremos los n documentos más parecidos, sienod n un parámetro de la consulta. Además, podremos filtrar la consulta en base a metadatos para que únicamente se ejecute sobre los documentos que cumplan una serie de criterios.

Ahora que sabes cómo funciona Chroma a nivel general, vamos a empezar a trabajar con Chroma.

Chroma requiere una versión de Python superior a 3.7. Aunque seguramente la gran mayoría de lectores cumpláis con esta característica, para asegurar el funcionamiento y facilitar el despliegue de la base de datos voy a desplegar Chroma en un contenedor Docker. Puedes aprender más sobre Docker aquí.

Primeros pasos con Chroma: levantando la base de datos

Como con toda base de datos, lo primero de todo es levantar la base de datos para que luego podamos añadir colecciones, insertar nuevos textos, etc.

En este sentido, actualemente (Junio de 2023) Chroma ofrece dos formas diferentes de almacenar la información:

1. DuckDB: persiste la información guardando los datos en ficheros parquet.
2. Clickhouse:guarda la información en otra base de datos Clickhouse.

La principal diferencia entre ambos es que, DuckDB permite crear un servicio de Chroma standalone, pero será mucho menos escalable. Por el contrario, si guardamos los datos en Clickhouse, podremos escalar mejor, aunque no será un servicio standalone.

Personalmente, si vas a hacer pruebas o requieres de una base de datos vectorial que no sea muy grande, te recomiendo que utilices DuckDB, mientras que si quieres escalar más, te recomiendo implementar Chroma con Clickhouse.

Ahora bien, ¿cómo podemos montar la base de datos en cada caso? Veamos.

Cómo crear una base de datos Chroma con DuckDB como backend

Para crear una base de datos de Chroma con DuckDB como backend, tendrás que hacer dos pasos:

1. Crear la base de datos Chroma y hacerla accesible mediante una API como FastAPI.
2. Crear la imagen Docker y desplegarla. Adicionalmente, si quieres persitencia de los datos, siempre puedes crear un Docker Compose con un volumen.

Para crear la base de datos Chroma y hacerla accesible mediante una API, Chroma ofrece la clase Settings, la cual nos permite definir cómo queremos que se implemente. En nuestro caso, hay dos parámetros que deberemos definir:

• chroma_db_impl: indica cuál serál el backend que utilice Chroma. En nuestro caso, debemos indicar duckdb+parquet.
• persist_directory nos permite indicar en qué carpeta se guardarán los ficheros parquet para conseguir el almacenamiento persistente.

Si queremos que la carpeta persist_directory persista dentro del contenedor, recuerda crear un volúmen para dicha carpeta.

Por último, pasaremos los resultados de ese settings a FastAPI para crear la API con la que interactuar con Chroma.

Así pues, el código será el siguiente:

# server.py
import chromadb
import chromadb.config
from chromadb.server.fastapi import FastAPI

settings = chromadb.config.Settings(
    chroma_db_impl="duckdb+parquet", 
    persist_directory='chroma_data'
)
server = FastAPI(settings)
app = server.app

Por úlitmo, tenemos que crear un Dockerfile que instale las librerías necesarias y ejecute un la API en un webserver.

Personalmente he tenido problemas con la última versión de FastAPI, por lo que recomiendo utilizar la versión 0.85.1 de FastAPI.

Ese proceso se realiza en el siguiente Dockerfile:

FROM python:3.8
RUN pip install uvicorn
RUN pip install chromadb 
RUN pip install --force-reinstall fastapi==0.85.1
COPY server.py .
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]

Con esto, ya tenemos nuestro Chroma listo para ser desplegado en un contenedor usando DuckDB como backend. Ahora, veamos cómo usar Chroma con Clickhouse como backend.

Cómo crear una base de datos Chroma con Clickhouse como backend

Para utilizar Chroma con Clickhouse sí o sí vamos a necesitar utilizar Docker Compose, ya que necesitaremos ejecutar Chroma en un microservico y Clickhouse en otro.

Si no sabes qué es o cómo funciona Docker Compose, puedes aprender más sobre ello aquí.

Hecho esto, vamos a modificar el script que crea el servidor. Concretamente, en la configuración deberemos indicar:

• chroma_db_impl: cómo se implementa Chroma. En este caso recibirá el valor clickhouse.
• clickhouse_host: el host donde se está ejecutando Clickhouse. Como nosotros lo ejecutaremos dentro de un Compose, indicaremos el nombre del servicio.
• clickhouse_port: puerto HTTP en el que está escuchando Clickhouse, el cual podemos configurar en el compose, aunque por defecto es 8123.

Nota: tanto el host como el puerto pueden ser definidos como una variable de entorno, lo cual es más recomendable, ya que no expones datos. El tipo de implementación de Chroma se puede definir con la variable CHROMA_DB_IMPL, el host de clickhouse se puede definir en la variable CLICKHOUSE_HOST y el puerto en la variable CLICKHOUSE_PORT.

import chromadb
import chromadb.config
from chromadb.server.fastapi import FastAPI

settings = chromadb.config.Settings(
    chroma_db_impl = 'clickhouse',
    clickhouse_host ='clickhouse',
    clickhouse_port = 8123 
)

server = FastAPI(settings)
app = server.app()

Hecho esto, vamos a usar el mismo Dockerfile que para la implementación de Chroma con DuckDB, la cual es la siguiente:

FROM python:3.8
RUN pip install uvicorn
RUN pip install chromadb 
RUN pip install --force-reinstall fastapi==0.85.1
COPY server.py .
CMD ["uvicorn", "server:app", "--host", "0.0.0.0", "--port", "8000"]

Por último, vamos a definir un Compose que levante tanto Chroma como Clickhouse. Además, tenemos que definir el network para que ambos contenedores estén conectados.

Asimismo, es recomendable también definir volúmenes tanto para Chroma como para Clickhouse. A continuación os pongo un ejemplo inspirado en el test que utiliza el propio Chroma:

services:
  chroma:
    image: chroma
    build:
      context: .
      dockerfile: Dockerfile
    ports:
      - 8000:8000
    depends_on:
      - clickhouse
    networks:
      - net

  clickhouse:
    image: clickhouse/clickhouse-server:22.9-alpine
    environment:
      - ALLOW_EMPTY_PASSWORD=yes
      - CLICKHOUSE_TCP_PORT=9000
      - CLICKHOUSE_HTTP_PORT=8123
    ports:
      - '8123:8123'
      - '9000:9000'
    volumes:
      - clickhouse_data:/var/lib/clickhouse
      - clickhouse_logs:/var/log/clickhouse-server
    networks:
      - net

volumes:
  clickhouse_data:
    driver: local
  clickhouse_logs:
    driver: local

networks:
  net:
    driver: bridge

Ejecutando el compose ya tendríamos nuestra base de datos Chroma levantada ejecutándose en el puerto 8000.

Perfecto, ahora que sabemos cómo montar Chroma, una base de datos vectorial, veamos cómo podemos trabajar con ella desde Python.

Cómo trabajar con Chroma en Python

Cuando trabajamos con Chroma tenemos que saber varias cuestiones diferentes:

1. Cómo conectarnos el cliente a nuestra base de datos Chroma.
2. Cómo crear y gestionar colecciones.
3. Cómo insertar nuevos documentos y/o vectores en las colecciones.
4. Cómo hacer consultas a nuestra base de datos para obtener los vectores más parecidos.

Así pues, vamos a ir paso a paso viendo cada uno de estos puntos. ¡Vamos con ello!

1. Cómo conectarnos el cliente a nuestra base de datos Chroma

Para conectarnos e interactuar con una base de datos Chroma lo que necesitamos es un cliente. Esto en Python lo podemos conseguir instalando la siguiente librería:

pip install chromadb

Perfecto, ahora que tenemos Chroma instalado, vamos a conectarnos a nuestra base de datos Chroma. Para ello debemos indicar:

1. Forma en la que trabajaremos con Chroma, que en este caso es 'rest', es decir, mediante API.
2. El host en el que está corriendo Chroma. Si estás ejecutando Chroma en Docker en local, deberás indicar 'localhost' y, sino, el host o IP en la que se esté ejecutando.
3. El puerto en el que se está ejecutando Chroma, que deberá coincidir con el puerto que hemos indicado en el Dockerfile y en el Compose (si procede). En nuestro caso, el puerto 8000.

Hecho esto, vamos a ejecutar el método get_version() para comprobar que se puede conectar correctamente:

import chromadb
from chromadb.config import Settings

setting = Settings(
    chroma_api_impl="rest",
    chroma_server_host= 'localhost',
    chroma_server_http_port = 8000
)

client = chromadb.Client(setting)

client.get_version()

'0.3.25'

Perfecto, ya tenemos nuestro cliente conectado a nuestra base de datos vectorial. Ahora, veamos cómo podemos crear y gestionar colecciones.

2. Cómo crear y gestionar colecciones

Aunque no lo parezca, la creación y gestión de colecciones de una base de datos vectorial no es una cosa menor. A la hora de crear la colección es cuando definiremos el modelo que la base de datos utilizará para crear los embeddings en caso de que insertemos un documento y no un vector.

Existen diferentes formas de crear una colección, aunque la más sencilla es mediante el método create_collection:

from chromadb.utils import embedding_functions

collection = client.create_collection(
    name="test",
    embedding_function= embedding_functions.DefaultEmbeddingFunction()
)

Como puedes ver, cuando creamos una colección he definido una función de embedding que debe aplicar. Veamos qué opciones nos ofrece Chroma en este sentido.

Funciones de Embedding en Chroma

Tal como se ve en la función anterior, Chroma ofrece diferentes funciones para obtener los embeddings de los documentos. Más concretamente, a Junio de 2023 ofrece las siguientes funciones:

• Cualquier modelo de la librería SentenceTransformers enlace. De hecho, por defecto utiliza el modelo all-MiniLM-L6-v2.

Estas funciones las podemos cargar de la siguiente manera:

# Change mode_name by desired model
embeder = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="model-name")

• Modelos de la librería Instructor: se trata de modelos que puedes ejecutar en local sobre GPU. Para ello, primero deberás instalar la librería mediante el siguiente comando pip install InstructorEmbedding. Después, podrás definir el embedder tal como muestro a continuación:

# Change mode_name by desired model
embedder = embedding_functions.InstructorEmbeddingFunction(
    model_name="model_name",
    device="cuda" # Accepts cuda or CPU
)

• La API de OpenAI. Esto requiere que cuentes con una API Key de OpenAI, lo cual es de pago. Puedes aprender cómo conseguirlo aquí. En caso de ya disponer de ella, la forma de utilizar la API es la siguiente:

pip isntall  openai

Definición del embedder mediante la clase OpenAIEmbeddingFunction:

# Change mode_name by desired model
embeder = embedding_functions.OpenAIEmbeddingFunction(
    api_key="API_KEY",
    model_name="model-name
)

• La API de Cohere: el procedimiento es el mismo que en el caso de la API de OpenAI, primero debes registrarte, lo cual puedes aprender aquí. Desppués, debes instalar su librería ejecutando pip install cohere y por último debes definir el embedder de la siguiente forma:

# Change mode_name by desired model
embeddger  = embedding_functions.CohereEmbeddingFunction(
    api_key="YOUR_API_KEY",
    model_name="model-name"
)

• La API de Google Palm: el procedimiento es el mismo que en con el resto de APIs, primero debes registrarte y después, debes instalar su librería ejecutando pip install google-generativeai y por último debes definir el embedder de la siguiente forma:

# Change mode_name by desired model
embedder = embedding_functions.GooglePalmEmbeddingFunction(
    api_key=api_key,
    model=model_name
)

• Función de embedding personalizada: puedes crear tu propia clase de embeddings heredando la clase EmbeddingFunciton, tal como se muestra a continuación:

from chromadb.api.types import Documents, EmbeddingFunction, Embeddings

class MyEmbeddingFunction(EmbeddingFunction):
    def __call__(self, texts: Documents) -> Embeddings:
        # embed the documents somehow
        return embeddings

Como pudes ver existen muchos modelos de embeddings diferentes que podemos utilizar. Así que, visto esto, veamos cómo podemos eliminar y gestionar las colecciones.

Cómo gestionar las colecciones

Alguna cuestión interesante es que, si intentas crear una colección que ya existe, tendrás un error, tal como muestro a continuación:

client.create_collection(
    name="test",
    embedding_function= embedding_functions.DefaultEmbeddingFunction()
)

---------------------------------------------------------------------------
ValueError: Collection with name test already exists

Una forma de evitar simple, sin tener que comprobar si el nombre de la colección existe o no, es pasando el párametro get_or_create como True. De esta forma, si la colección ya existe, actualizará sus metadatos y no devolverá error. En caso de no existir, la creará. Veamos un ejemplo creando la misma colección, pero con otro modelo:

embedder = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name= 'all-mpnet-base-v2'
)

client.create_collection(
    name="test",
    embedding_function= embedder,
    get_or_create = True
)

Como ves, la función se ha ejecutado correctamente y, de hecho, ha descargado el modelo puesto que no lo tenía descargado previamente. Por último, puedes eliminar una colección con la función delete_collection, indicando el nombre de la colección:

client.delete_collection(name = 'test')

Ahora que sabemos cómo crear y gestionar colecciones, veamos cómo podemos insertar documentos y/o vectores en dichas colecciones:

Cómo insertar nuevos documentos y/o vectores en las colecciones

Para trabajar con una colección, lo primero que debemos hacer es obtener dicha colección como un objeto en Python. Esto lo podemos hacer mediante la función get_collection del cliente. En este caso también deberemos indicar la embedding function que se debería aplicar.

collection = client.get_collection(name = 'test', embedding_function= embedder)

Ahora que tenemos la colección, ya podemos insertar datos en ella. Para ello, se utiliza el método add del objeto que acabamos de crear. En cualquier caso, existen dos formas diferentes de insertar contenido a Chroma:

• Insertar un documento o una lista de documentos, de tal forma que sea el propio Chroma la que transforme esos documentos en embeddings. Para ello, deberemos pasar una lista de textos al parámetro documents.
• Insertar un vector manualmente. Ya sea bien porque hemos calculado nosotros los embeddings previamente o porque vamos a almacenar vectores de otro tipo. Al fin y al cabo, una base de datos vectorial está pensada para alojar vectores, no necesariamente tienen por qué ser emebeddings. Para ello, debemos pasar una lista de embeddings (listas) al parámetro embeddings.

Asimismo, a la hora de insertar la información podemos indicar las siguientes cuestiones:

• id: se trata del identificador del vector, que es obligatorio y debe ser único.
• metadata: se trata de información adicional que podemos llegar a usar más adelante de cara a filtrar la información. Por ejemplo, supongamos que contamos con el año del documento, se podría llegar a filtrar para obtener vectores similares para un año o periodo concreto. En este caso, es opcional.

Visto esto, vamos a extraer una serie de documentos (textos) para poder guardarlos en nuestra base de datos vectorial. El enfoque es el siguiente:

1. Descargar la lista de las 500 empresas que forman el S&P500, lo cual haremos de esta página (enlace).
2. Para cada empresa, extraer la descripcción principal de la empresa de Wikipedia. Esto lo haremos con la función summary de la librería wikipedia.

Nota: el objetivo de esta parte no es conseguir una informacuión limpia y perfecta el 100% de los casos, sino hacer una extracción rápida y simple que sirva como ejemplo. Por eso, los casos en los que la extracción no funcione los descartaré.

Así pues primero instalo las librerías:

pip install lxml wikipedia

Y ahora descargo los datos que necesito (tardará unos mínutos).

import pandas as pd
import wikipedia


def get_wikipedia_summary(name):
    try: 
        summary = wikipedia.summary(name + ' company')
    except:
        summary = None
    finally:
        return summary

url = 'https://en.wikipedia.org/wiki/List_of_S%26P_500_companies'

df = pd.read_html(url, attrs = {"id": "constituents"})[0]
df['company_summary'] = df['Security'].apply(lambda x: get_wikipedia_summary(x))

df.head()

Por último, vamos a borrar los datos con nulos y obtener:

1. Un listado con los documentos (los resúmenes de los textos).
2. Un listado con los ids que crearemos.
3. Una lista de diccionarios con los metadatos que queremos subir por documento. Estos metadatos incluirán información de filtrado que nos pueda interesar, como el Sector el sub sector, la fecha de fundación de la empresa o la fecha en la que se ha añadido al S&P500.

cleaned_df = (
    df
    .query("~company_summary.isnull()")
    .assign(
        date_founded = lambda x: x['Founded'].str.replace('\(.*\)|/.*', '', regex = True).astype(int),
        date_added = lambda x: pd.to_datetime(x['Date added'], errors='ignore') 
    )
    [['Security', 'GICS Sector', 'GICS Sub-Industry', 'date_founded', 'date_added', 'company_summary']]
    .dropna()
    .reset_index(drop=True)
)

documents = cleaned_df['company_summary'].tolist()
ids = cleaned_df.index.astype(str).tolist()
metadata = cleaned_df.drop('company_summary', axis = 1).to_dict(orient = 'records')

Ahora que tenemos esta información, vamos a añadirla a nuestra base de datos vectorial. Importante, para cada documento tiene que crear un embedding, por lo que puede tardar un poco:

collection.add(
    ids = ids, 
    documents=documents, 
    metadatas=metadata
)

Ya tenemos los documentos creados. Si en lugar de documentos quisiéramos añadir vectores el enfoque sería el mismo, pero pasando una lista de vectores al parámetro embeddings, en lugar de al parámetro documents.

Perfecto, ahora que ya sabemos cómo podemos insertar documentos en nuestra base de datos vectorial Chroma, veamos cómo podemos consultar la base de datos.

Cómo hacer consultas a nuestra base de datos Chroma para obtener los vectores más parecidos

Ahora que tenemos nuestros datos ya subidos, podemos buscar los n documentos que más se parezcan a un texto o un embedding. Para ello, usaremos el método query de nuestra colección. Este método permite recibir los siguientes parámetros:

• query_texts: input en formato texto sobre el que queremos encontrar vectores parecidos.
• query_embeddings: input en formato vector sobre el que queremos encontrar vectores parecidos.
• n_results: número de resultados que debe devolver la búsqueda.
• where: filtrado de los vectores en base a metadatos.
• where_document: filtrado de los vectores en base a que los documentos contengan un contenido específico.
• include: qué debereía devolver la búsqueda. Por defecto devuelve toda la información disponible: («metadatas», «documents», «distances»).

Así pues, vamos a extraer los documentos que más se parezcan al prompt phone manufacturer:

results = collection.query(
    query_texts="phone manufacturers", 
    n_results=5
)

results

{'ids': [['375', '62', '124', '48', '314']],
 'distances': [[1.03172767162323,
   1.1286187171936035,
   1.164528250694275,
   1.1696319580078125,
   1.1702216863632202]],
 'embeddings': None,
 'metadatas': [[{'Security': 'Qualcomm',
    'GICS Sector': 'Information Technology',
    'GICS Sub-Industry': 'Semiconductors',
    'date_founded': 1985,
    'date_added': '1999-07-22'},
   {'Security': 'Best Buy',
    'GICS Sector': 'Consumer Discretionary',
    'GICS Sub-Industry': 'Computer & Electronics Retail',
    'date_founded': 1966,
    'date_added': '1999-06-29'},
   {'Security': 'Corning Inc.',
    'GICS Sector': 'Information Technology',
    'GICS Sub-Industry': 'Electronic Components',
    'date_founded': 1851,
    'date_added': '1995-02-27'},
   {'Security': 'AT&T',
    'GICS Sector': 'Communication Services',
    'GICS Sub-Industry': 'Integrated Telecommunication Services',
    'date_founded': 1983,
    'date_added': '1983-11-30'},
   {'Security': 'Motorola Solutions',
    'GICS Sector': 'Information Technology',
    'GICS Sub-Industry': 'Communications Equipment',
    'date_founded': 1928,
    'date_added': '1957-03-04'}]],
 }

Como podemos ver, ante esta búsquda la base de datos nos ha devuelto las empresas Qualcomm, AT&T, Motorola Solutions, Best Buy y T-Mobile US. Si bien es cierto que en realidad no son empresas que fabriquen sus propios teléfonos en todos los casos, sí que son empresas que están muy relacionadas con ello.

La calidad de las respuestas que obtengamos dependerá en gran medida tanto del prompt que introduzcamos como del funcionamiento del modelo utilizado para crear los embeddings.

Así pues, podríamos filtrar los resultados para que solo nos incluyera empresas cuyo GICS Sector sea Information Technology. Veamoslo:

collection.query(
    query_texts="phone manufacturers", 
    n_results=5, 
    where = {'GICS Sector': 'Information Technology'}
)

{'ids': [['375', '124', '314', '373', '333']],
 'distances': [[1.03172767162323,
   1.164528250694275,
   1.1702216863632202,
   1.1854088306427002,
   1.2617770433425903]],
 'embeddings': None,
 'metadatas': [[{'Security': 'Qualcomm',
    'GICS Sector': 'Information Technology',
    'GICS Sub-Industry': 'Semiconductors',
    'date_founded': 1985,
    'date_added': '1999-07-22'},
   {'Security': 'Corning Inc.',
    'GICS Sector': 'Information Technology',
    'GICS Sub-Industry': 'Electronic Components',
    'date_founded': 1851,
    'date_added': '1995-02-27'},
   {'Security': 'Motorola Solutions',
    'GICS Sector': 'Information Technology',
    'GICS Sub-Industry': 'Communications Equipment',
    'date_founded': 1928,
    'date_added': '1957-03-04'},
   {'Security': 'Qorvo',
    'GICS Sector': 'Information Technology',
    'GICS Sub-Industry': 'Semiconductors',
    'date_founded': 2015,
    'date_added': '2015-06-11'},
   {'Security': 'Nvidia',
    'GICS Sector': 'Information Technology',
    'GICS Sub-Industry': 'Semiconductors',
    'date_founded': 1993,
    'date_added': '2001-11-30'}]],
 }

Como puedes ver en este caso solo ha devuelto información de empresas de dicho sector. Esto sería el equivalente a una búsqueda por match exacto. Sin embargo, esta no es el único tipo de filtrado que se puede hacer. Concretamente existen los siguientes operadores:

• $eq: igual a (utilizado previamente).
• $ne: no igual a.
• $gt: mayor a.
• $gte: mayor o igual a.
• $lt: menor a.
• $lte: menor o igual a.

Es importante no olvidar el símbolo $, puesto que sino la query no funcionará.

Asimismo, si queremos que se apliquen varias condiciones (ya sea una de las dos o las dos a la vez), deberemos incluir el filtro en un diccionario con el operador $or o el operador $and.

Así pues, apliquemos la misma query para empresas del sector Information Technology pero que, además, su fecha de creación sea superior al 1990.

where_clause = {
    "$and" : [
        {"GICS Sector" : {"$eq": "Information Technology"}}, 
        {"date_founded": {"$gte": 1990}}
    ]
}

results = collection.query(
    query_texts="phone manufacturers", 
    n_results=5, 
    where = where_clause
)

[(el['Security'], el['GICS Sector'], el['date_founded']) for el in results['metadatas'][0]]

[('Qorvo', 'Information Technology', 2015),
 ('Nvidia', 'Information Technology', 1993),
 ('ON Semiconductor', 'Information Technology', 1999),
 ('Palo Alto Networks', 'Information Technology', 2005),
 ('Arista Networks', 'Information Technology', 2004)]

Asimismo, también podríamos añadir filtros, como que el documento contenga una palabra específica, como por ejemplo Apple. Probemos:

results = collection.query(
    query_texts="Phone", 
    n_results=5, 
    where = {"GICS Sector" : {"$eq": "Information Technology"}},
    where_document= {"$contains": "Apple"}
)

queried_companies = [el['Security'] for el in results['metadatas'][0]]
queried_companies

['Corning Inc.', 'Microsoft', 'Apple Inc.', 'ON Semiconductor', 'Adobe Inc.']

Por último vamos a comprobar si, efectivamente, esas empresas cuentan con la palabra Apple en su descripcción:

(
    cleaned_df
    .query('Security in @queried_companies')
    .assign(
        contains_apple = lambda x: x['company_summary'].str.contains('Apple')
    )
    [['Security', 'contains_apple']]
)

Como podéis ver, efectivamente, todas las empresas que devuelve efectivamente cuentan con la palabra «Apple» en su descripcción.

Conclusion

Espero que este post te haya servido para entender mejor qué es una base de datos vectorial, cómo puedes montarla y cómo puedes trabajar con ella. En este caso, he centrado el post en Chroma, puesto que es una base de datos open-source que considero que es bastante flexible y muy muy sencilla, aunque la idea que hay por detrás de esta es la misma para todas las bases de datos vectoriales.

Como ves, crear, alimentar y consultar una base de datos vectorial no es excesivamente complicado. En cualquier caso, es importante siempre tener un conocimiento profundo tanto del contenido que se está incluyendo a la base de datos vectorial como del funcionamiento del sistema de embeddings.

Por último, comentar que, este post pretende ser una explicación profunda pero no exhaustiva de todo lo que ofrece Chroma. Si quieres profundizar aún más en su comportamiento, te recomiendo que leas su documentación (enlace) y también el propio código en su repositorio de Github (enlace).