Cree una búsqueda semántica con incrustaciones de LLM

En este artículo, aprenderá cómo crear un motor de búsqueda semántico simple utilizando incrustaciones de oraciones y vecinos más cercanos.

Los temas que cubriremos incluyen:

Comprender las limitaciones de la búsqueda basada en palabras clave. Generación de incrustaciones de texto con un modelo de transformador de oraciones. Implementación de un canal de búsqueda semántica del vecino más cercano en Python.

Empecemos.

Cree una búsqueda semántica con incrustaciones de LLM
Imagen del editor

Introducción

Históricamente, los motores de búsqueda tradicionales se han basado en la búsqueda por palabras clave. En otras palabras, ante una consulta como “mejores templos y santuarios para visitar en Fukuoka, Japón”, los resultados se obtienen basándose en la concordancia de palabras clave, de modo que los documentos de texto que contienen coincidencias de palabras como “templo”, “santuario” y “Fukuoka” se consideran más relevantes.

Sin embargo, este enfoque clásico es notoriamente rígido, ya que se basa en gran medida en coincidencias exactas de palabras y pasa por alto otros matices semánticos importantes, como sinónimos o frases alternativas (por ejemplo, “perro joven” en lugar de “cachorro”). Como resultado, es posible que se omitan documentos muy relevantes sin darse cuenta.

La búsqueda semántica aborda esta limitación centrándose en el significado más que en la redacción exacta. Los modelos de lenguaje grande (LLM) juegan un papel clave aquí, ya que algunos de ellos están entrenados para traducir texto en representaciones vectoriales numéricas llamadas incrustaciones, que codifican la información semántica detrás del texto. Cuando dos textos como “los perros pequeños son muy curiosos por naturaleza” y “los cachorros son curiosos por naturaleza” se convierten en vectores incrustados, esos vectores serán muy similares debido a su significado compartido. Mientras tanto, los vectores de incorporación de “los cachorros son curiosos por naturaleza” y “Dazaifu es un santuario característico en Fukuoka” serán muy diferentes, ya que representan conceptos no relacionados.

Siguiendo este principio, que puede explorar con más profundidad aquí, el resto de este artículo le guiará a través del proceso completo de creación de un motor de búsqueda semántica compacto pero eficiente. Si bien es minimalista, funciona de manera efectiva y sirve como punto de partida para comprender cómo se construyen los sistemas modernos de búsqueda y recuperación, como las arquitecturas de recuperación de generación aumentada (RAG).

El código que se explica a continuación se puede ejecutar sin problemas en una instancia de Google Colab o Jupyter Notebook.

Guía paso a paso

Primero, realizamos las importaciones necesarias para este ejemplo práctico:

Usaremos un conjunto de datos públicos de juguete llamado "ag_news", que contiene textos de artículos de noticias. El siguiente código carga el conjunto de datos y selecciona los primeros 1000 artículos.

Ahora cargamos el conjunto de datos y extraemos la columna "texto", que contiene el contenido del artículo. Luego, imprimimos una breve muestra del primer artículo para inspeccionar los datos:

El siguiente paso es obtener vectores de incrustación (representaciones numéricas) para nuestros 1000 textos. Como se mencionó anteriormente, algunos LLM están capacitados específicamente para traducir texto en vectores numéricos que capturan características semánticas. Los modelos de transformadores de frases Hugging Face, como "all-MiniLM-L6-v2", son una opción común. El siguiente código inicializa el modelo y codifica el lote de documentos de texto en incrustaciones.

A continuación, inicializamos un objeto NearestNeighbors, que implementa una estrategia de vecino más cercano para encontrar los k documentos más similares a una consulta determinada. En términos de incrustaciones, esto significa identificar los vectores más cercanos (distancia angular más pequeña). Usamos la métrica del coseno, donde los vectores más similares tienen distancias de coseno más pequeñas (y valores de similitud de coseno más altos).

La lógica central de nuestro motor de búsqueda se resume en la siguiente función. Toma una consulta de texto sin formato, especifica cuántos resultados principales recuperar mediante top_k, calcula la incrustación de la consulta y recupera los vecinos más cercanos del índice.

El bucle dentro de la función imprime los k resultados principales clasificados por similitud:

Y eso es todo. Para probar la función, podemos formular un par de consultas de búsqueda de ejemplo:

Los resultados están clasificados por similitud (aquí truncados para mayor claridad):

Resumen

Lo que hemos construido aquí puede verse como una puerta de entrada a sistemas de generación aumentada de recuperación. Si bien este ejemplo es intencionalmente simple, los motores de búsqueda semántica como este forman la capa de recuperación fundamental en las arquitecturas modernas que combinan la búsqueda semántica con grandes modelos de lenguaje.

Ahora que sabe cómo crear un motor de búsqueda semántico básico, es posible que desee explorar los sistemas de recuperación de generación aumentada con más profundidad.