набор данных dbpedia

Набор данных с 1 миллионом статей из Википедии и их векторными эмбеддингами

Набор данных dbpedia содержит 1 миллион статей из Википедии и их векторные эмбеддинги, сгенерированные с помощью модели text-embedding-3-large от OpenAI.

Этот набор данных отлично подходит для знакомства с векторными эмбеддингами, поиском по сходству векторов и Generative AI. Мы используем этот набор данных, чтобы продемонстрировать приближённый поиск ближайших соседей в ClickHouse и простое, но мощное Q&A-приложение.

​

Сведения о наборе данных

Набор данных содержит 26 файлов Parquet, размещённых на huggingface.co. Файлы называются 0.parquet, 1.parquet, …, 25.parquet. Чтобы посмотреть несколько строк из набора данных в качестве примера, перейдите на эту страницу Hugging Face.

​

Создайте таблицу

Создайте таблицу dbpedia для хранения идентификатора статьи, заголовка, текста и эмбеддинг-вектора:

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

​

Загрузите таблицу

Чтобы загрузить данные из всех файлов Parquet, выполните следующую команду оболочки:

for i in $(seq 0 25); do
  echo "Обработка файла ${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 "Файл ${i} обработан."
done

Либо можно выполнить отдельные SQL-команды, как показано ниже, чтобы загрузить каждый из 25 файлов 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;

Убедитесь, что в таблице dbpedia отображается 1 миллион строк:

SELECT count(*)
FROM dbpedia

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

​

Семантический поиск

Рекомендуем к прочтению: “Векторные эмбеддинги ” — руководство OpenAPI

Семантический поиск (также называемый поиском по сходству) с использованием векторных эмбеддингов включает следующие шаги:

• Принять от пользователя поисковый запрос на естественном языке, например “Расскажи мне о живописных железнодорожных путешествиях”, “Детективные романы, действие которых происходит в Европе” и т. д.
• Сгенерировать эмбеддинг-вектор для поискового запроса с помощью LLM-модели
• Найти в наборе данных ближайших соседей для эмбеддинг-вектора поискового запроса

Ближайшие соседи — это документы, изображения или другой контент, наиболее релевантный запросу пользователя. Полученные результаты служат ключевым входом для Retrieval Augmented Generation (RAG) в приложениях Generative AI.

​

Выполните поиск по векторному сходству полным перебором

Поиск KNN (k - ближайших соседей) или поиск полным перебором заключается в вычислении расстояния от каждого вектора в наборе данных до искомого эмбеддинг-вектора с последующей сортировкой по расстоянию, чтобы получить ближайших соседей. Для набора данных dbpedia быстрый способ наглядно продемонстрировать семантический поиск — использовать эмбеддинг-векторы из самого набора данных в качестве поисковых векторов. Например:

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.)

Запишите задержку выполнения запроса, чтобы мы могли сравнить её с задержкой выполнения ANN-запроса (с использованием векторного индекса). Также запишите задержку выполнения запроса при холодном файловом кэше ОС и с max_threads=1, чтобы оценить реальное использование вычислительных ресурсов и пропускной способности хранилища (экстраполируйте это на продакшн-набор данных с миллионами векторов!)

​

Создайте индекс векторного сходства

Выполните следующий SQL-запрос, чтобы задать и создать индекс векторного сходства для столбца 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;

Параметры и особенности производительности при создании индекса и выполнении поиска описаны в документации.

Создание и сохранение индекса может занять несколько минут в зависимости от количества доступных ядер процессора и пропускной способности хранилища.

​

Выполните ANN-поиск

Approximate Nearest Neighbours (ANN) — это группа методов (например, специальные структуры данных, такие как графы и случайные леса), которые находят результаты значительно быстрее, чем точный векторный поиск. Точность результатов обычно “достаточно высока” для практического применения. Многие приближённые методы предоставляют параметры, позволяющие настраивать компромисс между точностью результатов и временем поиска.

После построения индекса векторного сходства запросы векторного поиска будут автоматически использовать этот индекс:

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 строк в наборе. Elapsed: 0.025 сек. Обработано 32.03 тысяч строк, 2.10 МБ (1.29 миллион строк/с., 84.80 МБ/с.)

​

Генерация эмбеддингов для поискового запроса

В запросах поиска по сходству, рассмотренных выше, в качестве вектора поиска используется один из существующих векторов в таблице dbpedia. В реальных приложениях вектор поиска должен генерироваться для пользовательского запроса, который может быть сформулирован на естественном языке. Вектор поиска следует генерировать с помощью той же LLM-модели, которая использовалась для создания эмбеддинг-векторов для набора данных.

Ниже приведён пример Python-скрипта, демонстрирующий, как программно вызывать API OpenAI для генерации эмбеддинг-векторов с использованием модели text-embedding-3-large. Затем поисковый эмбеддинг-вектор передаётся в качестве аргумента в функцию cosineDistance() в запросе SELECT.

Для запуска скрипта необходимо, чтобы API-ключ OpenAI был задан в переменной окружения OPENAI_API_KEY. API-ключ OpenAI можно получить после регистрации на https://platform.openai.com.

import sys
from openai import OpenAI
import clickhouse_connect

ch_client = clickhouse_connect.get_client(compress=False) # Передайте учётные данные ClickHouse
openai_client = OpenAI() # Задайте переменную окружения 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:
    # Получить поисковый запрос от пользователя
    print("Enter a search query :")
    input_query = sys.stdin.readline();

    # Обратиться к конечной точке OpenAI API для получения эмбеддинга
    print("Generating the embedding for ", input_query);
    embedding = get_embedding(input_query,
                              model='text-embedding-3-large')

    # Выполнить запрос векторного поиска в 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("---------------")

​

Демо-приложение Q&A

В приведённых выше примерах были показаны семантический поиск и извлечение документов с использованием ClickHouse. Далее представлен очень простой, но весьма перспективный пример приложения с генеративным ИИ.

Приложение выполняет следующие шаги:

1. Принимает от пользователя тему в качестве входных данных
2. Генерирует эмбеддинг-вектор для темы, вызывая API OpenAI с моделью text-embedding-3-large
3. Извлекает наиболее релевантные статьи/документы из Википедии, используя поиск по векторному сходству в table dbpedia
4. Принимает от пользователя вопрос в свободной форме на естественном языке, относящийся к теме
5. Использует Chat API OpenAI gpt-3.5-turbo, чтобы ответить на вопрос на основе сведений из документов, извлечённых на шаге #3. Документы, извлечённые на шаге #3, передаются в Chat API как контекст и являются ключевым связующим звеном в генеративном ИИ.

Ниже сначала приведена пара примеров диалогов при запуске Q&A-приложения, а затем — код для Q&A-приложения. Для запуска приложения необходимо, чтобы в переменной окружения OPENAI_API_KEY был задан API-ключ OpenAI. API-ключ OpenAI можно получить после регистрации на 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.

$

Код:

import sys
import time
from openai import OpenAI
import clickhouse_connect

ch_client = clickhouse_connect.get_client(compress=False) # Передайте учётные данные ClickHouse здесь
openai_client = OpenAI() # Задайте переменную окружения 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:
    # Получить тему от пользователя
    print("Enter a topic : ", end="", flush=True)
    input_query = sys.stdin.readline()
    input_query = input_query.rstrip()

    # Сгенерировать эмбеддинг-вектор для темы поиска и выполнить запрос к 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)

    # Собрать все подходящие статьи/документы
    results = ""
    for row in result.result_rows:
        results = results + row[2]

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

    # Промпт для OpenAI Chat API
    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,
    )

    # Вывести ответ на вопрос!
    print(response.choices[0].message.content)
    print("\n")