Búsqueda semántica con incrustaciones de PostgreSQL y OpenAI |  de Dima Timofeev |  noviembre de 2023
Imagen por Ígor Omilayev en desempaquetar

Implementar la búsqueda semántica en bases de datos corporativas puede ser un desafío y requiere un esfuerzo significativo. Sin embargo, ¿tiene que ser así? En este artículo, demuestro cómo puede utilizar PostgreSQL junto con OpenAI Embeddings para implementar la búsqueda semántica en sus datos. Si prefiere no utilizar la API de OpenAI Embeddings, le proporcionaré enlaces a modelos de incrustación gratuitos.

A un nivel muy alto, las bases de datos vectoriales con LLM permiten realizar búsquedas semánticas de los datos disponibles (almacenados en bases de datos, documentos, etc.) gracias al “Estimación eficiente de representaciones de palabras en el espacio vectorial”artículo (también conocido como “Papel Word2Vec”) escrito en coautoría por el legendario Jeff Dean, sabemos cómo representar palabras como vectores de valor real. Las incrustaciones de palabras son representaciones vectoriales densas de palabras en un espacio vectorial donde las palabras con significados similares están más cercanas entre sí. Las incrustaciones de palabras capturan relaciones semánticas entre palabras y existe más de una técnica para crearlas.

Imagen del autor

Practiquemos y usemos OpenAI incrustación-de-texto-ada ¡modelo! La elección de la función de distancia normalmente no importa mucho. OpenAI recomienda la similitud del coseno. Si no desea utilizar incrustaciones de OpenAI y prefiere ejecutar un modelo diferente localmente en lugar de realizar llamadas API, le sugiero que considere una de las siguientes Modelos preentrenados de SentenceTransformers. Elige sabiamente tu modelo.

import os

import openai
from openai.embeddings_utils import cosine_similarity

openai.api_key = os.getenv("OPENAI_API_KEY")

def get_embedding(text: str) -> list:
response = openai.Embedding.create(
input=text,
model="text-embedding-ada-002"
)
return response['data'][0]['embedding']

good_ride = "good ride"
good_ride_embedding = get_embedding(good_ride)
print(good_ride_embedding)
# [0.0010935445316135883, -0.01159335020929575, 0.014949149452149868, -0.029251709580421448, -0.022591838613152504, 0.006514389533549547, -0.014793967828154564, -0.048364896327257156, -0.006336577236652374, -0.027027441188693047, ...]
len(good_ride_embedding)
# 1536

Ahora que hemos comprendido qué es una incrustación, utilicémosla para ordenar algunas reseñas.

good_ride_review_1 = "I really enjoyed the trip! The ride was incredibly smooth, the pick-up location was convenient, and the drop-off point was right in front of the coffee shop."
good_ride_review_1_embedding = get_embedding(good_ride_review_1)
cosine_similarity(good_ride_review_1_embedding, good_ride_embedding)
# 0.8300454513797334

good_ride_review_2 = "The drive was exceptionally comfortable. I felt secure throughout the journey and greatly appreciated the on-board entertainment, which allowed me to have some fun while the car was in motion."
good_ride_review_2_embedding = get_embedding(good_ride_review_2)
cosine_similarity(good_ride_review_2_embedding, good_ride_embedding)
# 0.821774476808789

bad_ride_review = "A sudden hard brake at the intersection really caught me off guard and stressed me out. I wasn't prepared for it. Additionally, I noticed some trash left in the cabin from a previous rider."
bad_ride_review_embedding = get_embedding(bad_ride_review)
cosine_similarity(bad_ride_review_embedding, good_ride_embedding)
# 0.7950041130579355

Si bien la diferencia absoluta puede parecer pequeña, considere una función de clasificación con miles y miles de reseñas. En tales casos, podemos priorizar resaltar solo los positivos en la parte superior.

Una vez que una palabra o un documento se ha transformado en una incrustación, se puede almacenar en una base de datos. Sin embargo, esta acción no clasifica automáticamente la base de datos como una base de datos vectorial. Sólo cuando la base de datos comience a admitir operaciones rápidas en el vector podremos etiquetarla legítimamente como una base de datos vectorial.

Existen numerosas bases de datos vectoriales comerciales y de código abierto, lo que lo convierte en un tema muy debatido. Demostraré el funcionamiento de bases de datos vectoriales utilizando un pgvectoruna extensión PostgreSQL de código abierto que permite funcionalidades de búsqueda de similitudes vectoriales para posiblemente la base de datos más popular.

Ejecutemos el contenedor PostgreSQL con pgvector:

docker pull ankane/pgvector

docker run --env "POSTGRES_PASSWORD=postgres" --name "postgres-with-pgvector" --publish 5432:5432 --detach ankane/pgvector

Empecemos pgcli para conectarse a la base de datos (pgcli postgres://postgres:postgres@localhost:5432) y cree una tabla, inserte las incrustaciones que calculamos anteriormente y luego seleccione elementos similares:

-- Enable pgvector extension.
CREATE EXTENSION vector;

-- Create a vector column with 1536 dimensions.
-- The `text-embedding-ada-002` model has 1536 dimensions.
CREATE TABLE reviews (text TEXT, embedding vector(1536));

-- Insert three reviews from the above. I omitted the input for your convinience.
INSERT INTO reviews (text, embedding) VALUES ('I really enjoyed the trip! The ride was incredibly smooth, the pick-up location was convenient, and the drop-off point was right in front of the coffee shop.', '[-0.00533589581027627, -0.01026702206581831, 0.021472081542015076, -0.04132508486509323, ...');
INSERT INTO reviews (text, embedding) VALUES ('The drive was exceptionally comfortable. I felt secure throughout the journey and greatly appreciated the on-board entertainment, which allowed me to have some fun while the car was in motion.', '[0.0001858668401837349, -0.004922827705740929, 0.012813017703592777, -0.041855424642562866, ...');
INSERT INTO reviews (text, embedding) VALUES ('A sudden hard brake at the intersection really caught me off guard and stressed me out. I was not prepared for it. Additionally, I noticed some trash left in the cabin from a previous rider.', '[0.00191772251855582, -0.004589076619595289, 0.004269456025213003, -0.0225954819470644, ...');

-- sanity check
select count(1) from reviews;
-- +-------+
-- | count |
-- |-------|
-- | 3 |
-- +-------+

Estamos preparados para buscar documentos similares ahora. He acortado nuevamente la incrustación para “buen viaje” porque imprimir 1536 dimensiones es excesivo.

--- The embedding we use here is for "good ride"
SELECT substring(text, 0, 80) FROM reviews ORDER BY embedding <-> '[0.0010935445316135883, -0.01159335020929575, 0.014949149452149868, -0.029251709580421448, ...';

-- +--------------------------------------------------------------------------+
-- | substring |
-- |--------------------------------------------------------------------------|
-- | I really enjoyed the trip! The ride was incredibly smooth, the pick-u... |
-- | The drive was exceptionally comfortable. I felt secure throughout the... |
-- | A sudden hard brake at the intersection really caught me off guard an... |
-- +--------------------------------------------------------------------------+
SELECT 3
Time: 0.024s

¡Terminado! Como puede observar, hemos calculado incrustaciones para múltiples documentos, los almacenamos en la base de datos y realizamos búsquedas de similitud de vectores. Las aplicaciones potenciales son enormes y van desde búsquedas corporativas hasta funciones en sistemas de registros médicos para identificar pacientes con síntomas similares. Además, este método no se limita a los textos; La similitud también se puede calcular para otros tipos de datos, como sonido, vídeo e imágenes.

¡Disfrutar!