conjunto de datos de dbpedia

Conjunto de datos que contiene 1 millón de artículos de Wikipedia y sus embeddings vectoriales

El conjunto de datos de dbpedia contiene 1 millón de artículos de Wikipedia y sus embeddings vectoriales, generados con el modelo text-embedding-3-large de OpenAI.

Este conjunto de datos es un excelente punto de partida para entender los embeddings vectoriales, la búsqueda de similitud vectorial y la IA generativa. Usamos este conjunto de datos para demostrar la búsqueda aproximada de vecinos más cercanos en ClickHouse y una aplicación de preguntas y respuestas simple pero potente.

​

Detalles del conjunto de datos

El conjunto de datos contiene 26 archivos Parquet alojados en huggingface.co. Los archivos se llaman 0.parquet, 1.parquet, …, 25.parquet. Para ver algunas filas de ejemplo del conjunto de datos, visita esta página de Hugging Face.

​

Crear tabla

Cree la tabla dbpedia para almacenar el id, el título, el texto y el vector de embeddings del artículo:

CREATE TABLE dbpedia
(
  id      String,
  title   String,
  text    String,
  vector  Array(Float32) CODEC(NONE)
) ENGINE = MergeTree ORDER BY (id);

​

Cargar tabla

Para cargar el conjunto de datos a partir de todos los archivos Parquet, ejecuta el siguiente comando de shell:

for i in $(seq 0 25); do
  echo "Procesando archivo ${i}..."
  clickhouse client -q "INSERT INTO dbpedia SELECT _id, title, text, \"text-embedding-3-large-1536-embedding\" FROM url('https://huggingface.co/api/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-large-1536-1M/parquet/default/train/${i}.parquet') SETTINGS max_http_get_redirects=5,enable_url_encoding=0;"
  echo "Archivo ${i} completado."
done

Como alternativa, se pueden ejecutar Sentencias SQL individuales, como se muestra a continuación, para cargar cada uno de los 25 archivos Parquet:

INSERT INTO dbpedia SELECT _id, title, text, "text-embedding-3-large-1536-embedding" FROM url('https://huggingface.co/api/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-large-1536-1M/parquet/default/train/0.parquet') SETTINGS max_http_get_redirects=5,enable_url_encoding=0;
INSERT INTO dbpedia SELECT _id, title, text, "text-embedding-3-large-1536-embedding" FROM url('https://huggingface.co/api/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-large-1536-1M/parquet/default/train/1.parquet') SETTINGS max_http_get_redirects=5,enable_url_encoding=0;
...
INSERT INTO dbpedia SELECT _id, title, text, "text-embedding-3-large-1536-embedding" FROM url('https://huggingface.co/api/datasets/Qdrant/dbpedia-entities-openai3-text-embedding-3-large-1536-1M/parquet/default/train/25.parquet') SETTINGS max_http_get_redirects=5,enable_url_encoding=0;

Verifique que se muestren 1 millón de filas en la tabla dbpedia:

SELECT count(*)
FROM dbpedia

   ┌─count()─┐
1. │ 1000000 │
   └─────────┘

​

Búsqueda semántica

Lectura recomendada: “Embeddings vectoriales ” en la guía de OpenAPI

La búsqueda semántica (también denominada búsqueda por similitud) mediante embeddings vectoriales implica los siguientes pasos:

• Recibir una consulta de búsqueda de un usuario en lenguaje natural; por ejemplo, “Háblame de algunos viajes panorámicos en tren”, “Novelas de suspense ambientadas en Europa”, etc.
• Generar un vector de embedding para la consulta de búsqueda mediante un modelo LLM
• Encontrar los vecinos más cercanos al vector de embedding de la consulta en el conjunto de datos

Los vecinos más cercanos son documentos, imágenes o contenidos que constituyen resultados relevantes para la consulta del usuario. Los resultados recuperados son la entrada clave para la generación aumentada mediante recuperación (RAG) en aplicaciones de IA generativa.

​

Ejecutar una búsqueda de similitud vectorial por fuerza bruta

La búsqueda KNN (k - vecinos más cercanos), o búsqueda por fuerza bruta, consiste en calcular la distancia de cada vector del conjunto de datos al vector de embedding de búsqueda y, a continuación, ordenar esas distancias para obtener los vecinos más cercanos. Con el conjunto de datos dbpedia, una forma rápida de observar visualmente la búsqueda semántica es usar vectores de embedding del propio conjunto de datos como vectores de búsqueda. Por ejemplo:

Query

SELECT id, title
FROM dbpedia
ORDER BY cosineDistance(vector, ( SELECT vector FROM dbpedia WHERE id = '<dbpedia:The_Remains_of_the_Day>') ) ASC
LIMIT 20

Response

    ┌─id────────────────────────────────────────┬─title───────────────────────────┐
 1. │ <dbpedia:The_Remains_of_the_Day>          │ The Remains of the Day          │
 2. │ <dbpedia:The_Remains_of_the_Day_(film)>   │ The Remains of the Day (film)   │
 3. │ <dbpedia:Never_Let_Me_Go_(novel)>         │ Never Let Me Go (novel)         │
 4. │ <dbpedia:Last_Orders>                     │ Last Orders                     │
 5. │ <dbpedia:The_Unconsoled>                  │ The Unconsoled                  │
 6. │ <dbpedia:The_Hours_(novel)>               │ The Hours (novel)               │
 7. │ <dbpedia:An_Artist_of_the_Floating_World> │ An Artist of the Floating World │
 8. │ <dbpedia:Heat_and_Dust>                   │ Heat and Dust                   │
 9. │ <dbpedia:A_Pale_View_of_Hills>            │ A Pale View of Hills            │
10. │ <dbpedia:Howards_End_(film)>              │ Howards End (film)              │
11. │ <dbpedia:When_We_Were_Orphans>            │ When We Were Orphans            │
12. │ <dbpedia:A_Passage_to_India_(film)>       │ A Passage to India (film)       │
13. │ <dbpedia:Memoirs_of_a_Survivor>           │ Memoirs of a Survivor           │
14. │ <dbpedia:The_Child_in_Time>               │ The Child in Time               │
15. │ <dbpedia:The_Sea,_the_Sea>                │ The Sea, the Sea                │
16. │ <dbpedia:The_Master_(novel)>              │ The Master (novel)              │
17. │ <dbpedia:The_Memorial>                    │ The Memorial                    │
18. │ <dbpedia:The_Hours_(film)>                │ The Hours (film)                │
19. │ <dbpedia:Human_Remains_(film)>            │ Human Remains (film)            │
20. │ <dbpedia:Kazuo_Ishiguro>                  │ Kazuo Ishiguro                  │
    └───────────────────────────────────────────┴─────────────────────────────────┘
20 rows in set. Elapsed: 0.261 sec. Processed 1.00 million rows, 6.22 GB (3.84 million rows/s., 23.81 GB/s.)

Anota la latencia de la consulta para poder compararla con la latencia de la consulta de ANN (usando un índice vectorial). Registra también la latencia de la consulta con la caché del sistema de archivos del SO en frío y con max_threads=1 para determinar el uso real de cómputo y del ancho de banda de almacenamiento (¡extrapólalo a un conjunto de datos de producción con millones de vectores!)

​

Crear un índice de similitud vectorial

Ejecute la siguiente instrucción SQL para definir y crear un índice de similitud vectorial en la columna vector:

ALTER TABLE dbpedia ADD INDEX vector_index vector TYPE vector_similarity('hnsw', 'cosineDistance', 1536, 'bf16', 64, 512);

ALTER TABLE dbpedia MATERIALIZE INDEX vector_index SETTINGS mutations_sync = 2;

Los parámetros y las consideraciones de rendimiento para la creación de índices y la búsqueda se describen en la documentación.

La creación y el almacenamiento del índice pueden tardar unos minutos, según la cantidad de núcleos de CPU disponibles y el ancho de banda del almacenamiento.

​

Realizar una búsqueda ANN

Approximate Nearest Neighbours o ANN engloba un conjunto de técnicas (p. ej., estructuras de datos especiales, como grafos y bosques aleatorios) que calculan resultados mucho más rápido que la búsqueda vectorial exacta. La precisión de los resultados suele ser “suficientemente buena” para un uso práctico. Muchas técnicas aproximadas ofrecen parámetros para ajustar el compromiso entre la precisión de los resultados y el tiempo de búsqueda.

Una vez que se ha creado el índice de similitud vectorial, las consultas de búsqueda vectorial usarán automáticamente el índice:

Query

SELECT
    id,
    title
FROM dbpedia
ORDER BY cosineDistance(vector, (
        SELECT vector
        FROM dbpedia
        WHERE id = '<dbpedia:Glacier_Express>'
    )) ASC
LIMIT 20

Response

    ┌─id──────────────────────────────────────────────┬─title─────────────────────────────────┐
 1. │ <dbpedia:Glacier_Express>                       │ Glacier Express                       │
 2. │ <dbpedia:BVZ_Zermatt-Bahn>                      │ BVZ Zermatt-Bahn                      │
 3. │ <dbpedia:Gornergrat_railway>                    │ Gornergrat railway                    │
 4. │ <dbpedia:RegioExpress>                          │ RegioExpress                          │
 5. │ <dbpedia:Matterhorn_Gotthard_Bahn>              │ Matterhorn Gotthard Bahn              │
 6. │ <dbpedia:Rhaetian_Railway>                      │ Rhaetian Railway                      │
 7. │ <dbpedia:Gotthard_railway>                      │ Gotthard railway                      │
 8. │ <dbpedia:Furka–Oberalp_railway>                 │ Furka–Oberalp railway                 │
 9. │ <dbpedia:Jungfrau_railway>                      │ Jungfrau railway                      │
10. │ <dbpedia:Monte_Generoso_railway>                │ Monte Generoso railway                │
11. │ <dbpedia:Montreux–Oberland_Bernois_railway>     │ Montreux–Oberland Bernois railway     │
12. │ <dbpedia:Brienz–Rothorn_railway>                │ Brienz–Rothorn railway                │
13. │ <dbpedia:Lauterbrunnen–Mürren_mountain_railway> │ Lauterbrunnen–Mürren mountain railway │
14. │ <dbpedia:Luzern–Stans–Engelberg_railway_line>   │ Luzern–Stans–Engelberg railway line   │
15. │ <dbpedia:Rigi_Railways>                         │ Rigi Railways                         │
16. │ <dbpedia:Saint-Gervais–Vallorcine_railway>      │ Saint-Gervais–Vallorcine railway      │
17. │ <dbpedia:Gatwick_Express>                       │ Gatwick Express                       │
18. │ <dbpedia:Brünig_railway_line>                   │ Brünig railway line                   │
19. │ <dbpedia:Regional-Express>                      │ Regional-Express                      │
20. │ <dbpedia:Schynige_Platte_railway>               │ Schynige Platte railway               │
    └─────────────────────────────────────────────────┴───────────────────────────────────────┘
20 rows in set. Elapsed: 0.025 sec. Processed 32.03 thousand rows, 2.10 MB (1.29 million rows/s., 84.80 MB/s.)

​

Generación de embeddings para la consulta de búsqueda

Las consultas de búsqueda por similitud vistas hasta ahora usan uno de los vectores existentes en la tabla dbpedia como vector de búsqueda. En aplicaciones reales, el vector de búsqueda debe generarse a partir de una consulta introducida por el usuario, que podría estar en lenguaje natural. El vector de búsqueda debe generarse con el mismo modelo LLM utilizado para generar los vectores de embedding del conjunto de datos.

A continuación se muestra un script de Python de ejemplo para demostrar cómo llamar mediante programación a las API de OpenAI para generar vectores de embedding con el modelo text-embedding-3-large. El vector de embedding de búsqueda se pasa después como argumento a la función cosineDistance() en la consulta SELECT.

Para ejecutar el script, es necesario configurar una API key de OpenAI en la variable de entorno OPENAI_API_KEY. La API key de OpenAI puede obtenerse después de registrarse en https://platform.openai.com.

import sys
from openai import OpenAI
import clickhouse_connect

ch_client = clickhouse_connect.get_client(compress=False) # Pasar las credenciales de ClickHouse
openai_client = OpenAI() # Establecer la variable de entorno OPENAI_API_KEY

def get_embedding(text, model):
  text = text.replace("\n", " ")
  return openai_client.embeddings.create(input = [text], model=model, dimensions=1536).data[0].embedding

while True:
    # Aceptar la consulta de búsqueda del usuario
    print("Enter a search query :")
    input_query = sys.stdin.readline();

    # Llamar al endpoint de la API de OpenAI para obtener el embedding
    print("Generating the embedding for ", input_query);
    embedding = get_embedding(input_query,
                              model='text-embedding-3-large')

    # Ejecutar la consulta de búsqueda vectorial en ClickHouse
    print("Querying clickhouse...")
    params = {'v1':embedding, 'v2':10}
    result = ch_client.query("SELECT id,title,text FROM dbpedia ORDER BY cosineDistance(vector, %(v1)s) LIMIT %(v2)s", parameters=params)

    for row in result.result_rows:
        print(row[0], row[1], row[2])
        print("---------------")

​

Aplicación de demostración de preguntas y respuestas

Los ejemplos anteriores demostraron la búsqueda semántica y la recuperación de documentos con ClickHouse. A continuación se presenta una aplicación de ejemplo de IA generativa muy sencilla, pero con gran potencial.

La aplicación realiza los siguientes pasos:

1. Acepta un tema como entrada del usuario
2. Genera un vector de embedding para el tema invocando la API de OpenAI con el modelo text-embedding-3-large
3. Recupera artículos/documentos de Wikipedia muy relevantes mediante búsqueda de similitud vectorial en la tabla dbpedia
4. Acepta del usuario una pregunta abierta en lenguaje natural relacionada con el tema
5. Usa la API de Chat gpt-3.5-turbo de OpenAI para responder a la pregunta basándose en el conocimiento contenido en los documentos recuperados en el paso n.º 3. Los documentos recuperados en el paso n.º 3 se pasan como contexto a la API de Chat y constituyen el vínculo clave en la IA generativa.

A continuación se muestran primero un par de ejemplos de conversación al ejecutar la aplicación de preguntas y respuestas, seguidos por el código de la aplicación de preguntas y respuestas. Para ejecutar la aplicación, es necesario configurar una API key de OpenAI en la variable de entorno OPENAI_API_KEY. La API key de OpenAI se puede obtener tras registrarse en https://platform.openai.com.

$ python3 QandA.py

Enter a topic : FIFA world cup 1990
Generating the embedding for 'FIFA world cup 1990' and collecting 100 articles related to it from ClickHouse...

Enter your question : Who won the golden boot
Salvatore Schillaci of Italy won the Golden Boot at the 1990 FIFA World Cup.

Enter a topic : Cricket world cup
Generating the embedding for 'Cricket world cup' and collecting 100 articles related to it from ClickHouse...

Enter your question : Which country has hosted the world cup most times
England and Wales have hosted the Cricket World Cup the most times, with the tournament being held in these countries five times - in 1975, 1979, 1983, 1999, and 2019.

$

Código:

import sys
import time
from openai import OpenAI
import clickhouse_connect

ch_client = clickhouse_connect.get_client(compress=False) # Pase las credenciales de ClickHouse aquí
openai_client = OpenAI() # Establezca la variable de entorno OPENAI_API_KEY

def get_embedding(text, model):
  text = text.replace("\n", " ")
  return openai_client.embeddings.create(input = [text], model=model, dimensions=1536).data[0].embedding

while True:
    # Obtener el tema de interés del usuario
    print("Enter a topic : ", end="", flush=True)
    input_query = sys.stdin.readline()
    input_query = input_query.rstrip()

    # Generar un vector de embedding para el tema de búsqueda y consultar ClickHouse
    print("Generating the embedding for '" + input_query + "' and collecting 100 articles related to it from ClickHouse...");
    embedding = get_embedding(input_query,
                              model='text-embedding-3-large')

    params = {'v1':embedding, 'v2':100}
    result = ch_client.query("SELECT id,title,text FROM dbpedia ORDER BY cosineDistance(vector, %(v1)s) LIMIT %(v2)s", parameters=params)

    # Recopilar todos los artículos/documentos coincidentes
    results = ""
    for row in result.result_rows:
        results = results + row[2]

    print("\nEnter your question : ", end="", flush=True)
    question = sys.stdin.readline();

    # Prompt para la API de Chat de OpenAI
    query = f"""Use the below content to answer the subsequent question. If the answer cannot be found, write "I don't know."

Content:
\"\"\"
{results}
\"\"\"

Question: {question}"""

    GPT_MODEL = "gpt-3.5-turbo"
    response = openai_client.chat.completions.create(
        messages=[
        {'role': 'system', 'content': "You answer questions about {input_query}."},
        {'role': 'user', 'content': query},
       ],
       model=GPT_MODEL,
       temperature=0,
    )

    # Imprimir la respuesta a la pregunta
    print(response.choices[0].message.content)
    print("\n")