Семантический поиск с использованием Airflow и Astronomer#

В данном руководстве Platform V Vector DB (далее - Vector DB) представлен в качестве провайдера в Apache Airflow будет использован открытый инструмент, который позволяет настроить рабочие процессы обработки данных.

Будет создан конвейер в виде DAG (направленный ациклический граф) на языке Python. С помощью этого инструмента можно воспользоваться мощным набором возможностей и библиотек языка Python для выполнения практически любых задач конвейера данных.

Astronomer - управляемая платформа, которая упрощает процесс разработки и развертывания проектов Airflow благодаря удобному интерфейсу командной строки и обширным возможностям автоматизации.

Airflow полезен при выполнении операций в Vector DB на основе событий данных или создании параллельных задач для генерации векторных вложений. Используя Airflow, можно настроить мониторинг и оповещения о конвейерах для полной наблюдаемости.

Предварительные условия#

Пожалуйста, убедитесь, что готовы следующие компоненты:

Запущенный экземпляр Vector DB. Будет использоваться бесплатный экземпляр из https://cloud.qdrant.io
CLI от Astronomer. Найдите инструкции по установке здесь.
Токен HuggingFace для создания вложений.

Реализация#

Необходимо построить DAG, который параллельно генерирует вложения для корпуса данных и выполняет семантическую выборку на основе пользовательского ввода.

Настройка проекта#

Интерфейс командной строки Astronomer делает настройку проекта Airflow очень простой:

mkdir qdrant-airflow-tutorial && cd qdrant-airflow-tutorial
astro dev init

Эта команда создает все файлы проекта, необходимые для локального запуска Airflow. Найдете каталог под названием dags, где можно разместить файлы Python DAG.

Чтобы использовать Vector DB внутри Airflow, установите провайдер Vector DB Airflow, добавив следующее в файл requirements.txt

apache-airflow-providers-qdrant

Настройте учетные данные#

Можно настроить подключения к провайдерам с помощью интерфейса пользователя Airflow, переменных окружения или файла airflow_settings.yml.

Добавьте следующее содержимое в файл .env в проекте. Замените значения согласно учетным данным.

HUGGINGFACE_TOKEN="<YOUR_HUGGINGFACE_ACCESS_TOKEN>"
AIRFLOW_CONN_QDRANT_DEFAULT='{
    "conn_type": "qdrant",
    "host": "xyz-example.eu-central.aws.cloud.qdrant.io:6333",
    "password": "<YOUR_QDRANT_API_KEY>"
}'

Добавление корпуса данных#

Добавьте некоторые тестовые данные для работы. Вставьте следующий контент в файл под названием books.txt в каталоге include.

| To Kill a Mockingbird (1960) | fiction | Harper Lee Pulitzer Prize-winning novel explores racial injustice and moral growth through the eyes of young Scout Finch in the Deep South.
| Harry Potter and the Sorcerer Stone (1997) | fantasy | J.K. Rowling magical tale follows Harry Potter as he discovers his wizarding heritage and attends Hogwarts School of Witchcraft and Wizardry.
| The Great Gatsby (1925) | fiction | F. Scott Fitzgerald classic novel delves into the glitz, glamour, and moral decay of the Jazz Age through the eyes of narrator Nick Carraway and his enigmatic neighbour, Jay Gatsby.
| 1984 (1949) | dystopian | George Orwell dystopian masterpiece paints a chilling picture of a totalitarian society where individuality is suppressed and the truth is manipulated by a powerful regime.
| The Catcher in the Rye (1951) | fiction | J.D. Salinger iconic novel follows disillusioned teenager Holden Caulfield as he navigates the complexities of adulthood and society expectations in post-World War II America.
| Pride and Prejudice (1813) | romance | Jane Austen beloved novel revolves around the lively and independent Elizabeth Bennet as she navigates love, class, and societal expectations in Regency-era England.
| The Hobbit (1937) | fantasy | J.R.R. Tolkien adventure follows Bilbo Baggins, a hobbit who embarks on a quest with a group of dwarves to reclaim their homeland from the dragon Smaug.
| The Lord of the Rings (1954-1955) | fantasy | J.R.R. Tolkien epic fantasy trilogy follows the journey of Frodo Baggins to destroy the One Ring and defeat the Dark Lord Sauron in the land of Middle-earth.
| The Alchemist (1988) | fiction | Paulo Coelho philosophical novel follows Santiago, an Andalusian shepherd boy, on a journey of self-discovery and spiritual awakening as he searches for a hidden treasure.
| The Da Vinci Code (2003) | mystery/thriller | Dan Brown gripping thriller follows symbologist Robert Langdon as he unravels clues hidden in art and history while trying to solve a murder mystery with far-reaching implications.

Написание DAG#

Добавьте следующее содержание в файл books_recommend.py в директории dags:

import os
import requests

from airflow.decorators import dag, task
from airflow.models.baseoperator import chain
from airflow.models.param import Param
from airflow.providers.qdrant.hooks.qdrant import QdrantHook
from airflow.providers.qdrant.operators.qdrant import QdrantIngestOperator
from pendulum import datetime
from qdrant_client import models


QDRANT_CONNECTION_ID = "qdrant_default"
DATA_FILE_PATH = "include/books.txt"
COLLECTION_NAME = "airflow_tutorial_collection"

EMBEDDING_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
EMBEDDING_DIMENSION = 384
SIMILARITY_METRIC = models.Distance.COSINE


def embed(text: str) -> list:
    HUGGINFACE_URL = f"https://api-inference.huggingface.co/pipeline/feature-extraction/{EMBEDDING_MODEL_ID}"
    response = requests.post(
        HUGGINFACE_URL,
        headers={"Authorization": f"Bearer {os.getenv('HUGGINGFACE_TOKEN')}"},
        json={"inputs": [text], "options": {"wait_for_model": True}},
    )
    return response.json()[0]


@dag(
    dag_id="books_recommend",
    start_date=datetime(2023, 10, 18),
    schedule=None,
    catchup=False,
    params={"preference": Param("Something suspenseful and thrilling.", type="string")},
)
def recommend_book():
    @task
    def import_books(text_file_path: str) -> list:
        data = []
        with open(text_file_path, "r") as f:
            for line in f:
                _, title, genre, description = line.split("|")
                data.append(
                    {
                        "title": title.strip(),
                        "genre": genre.strip(),
                        "description": description.strip(),
                    }
                )

        return data

    @task
    def init_collection():
        hook = QdrantHook(conn_id=QDRANT_CONNECTION_ID)
        if not  hook.conn..collection_exists(COLLECTION_NAME):
            hook.conn.create_collection(
                COLLECTION_NAME,
                vectors_config=models.VectorParams(
                    size=EMBEDDING_DIMENSION, distance=SIMILARITY_METRIC
                ),
            )

    @task
    def embed_description(data: dict) -> list:
        return embed(data["description"])

    books = import_books(text_file_path=DATA_FILE_PATH)
    embeddings = embed_description.expand(data=books)

    qdrant_vector_ingest = QdrantIngestOperator(
        conn_id=QDRANT_CONNECTION_ID,
        task_id="qdrant_vector_ingest",
        collection_name=COLLECTION_NAME,
        payload=books,
        vectors=embeddings,
    )

    @task
    def embed_preference(**context) -> list:
        user_mood = context["params"]["preference"]
        response = embed(text=user_mood)

        return response

    @task
    def search_qdrant(
        preference_embedding: list,
    ) -> None:
        hook = QdrantHook(conn_id=QDRANT_CONNECTION_ID)

        result = hook.conn.query_points(
            collection_name=COLLECTION_NAME,
            query=preference_embedding,
            limit=1,
            with_payload=True,
        ).points

        print("Book recommendation: " + result[0].payload["title"])
        print("Description: " + result[0].payload["description"])

    chain(
        init_collection(),
        qdrant_vector_ingest,
        search_qdrant(embed_preference()),
    )


recommend_book()

Внимательно изучите, что делает каждая задача:

import_books: эта задача считывает текстовый файл, содержащий информацию о книгах (например, название, жанр и описание), и возвращает данные в виде списка словарей.
init_collection: эта задача инициализирует коллекцию в базе данных Vector DB, где будут храниться векторные представления описаний книг.
embed_description: это динамическая задача, создающая один отображаемый экземпляр задачи для каждой книги в списке. Задача использует функцию embed для генерации векторных вложений для каждого описания. Чтобы использовать другую модель встраивания, можно изменить параметры EMBEDDING_MODEL_ID, EMBEDDING_DIMENSION.
embed_user_preference: здесь берется ввод пользователя и преобразуется в вектор с помощью той же предварительно обученной модели, используемой для описаний книг.
qdrant_vector_ingest: эта задача импортирует данные о книгах в коллекцию Vector DB с помощью оператора QdrantIngestOperator, связывая каждое описание книги с соответствующими ей векторными вложениями.
search_qdrant: эта задача выполняет поиск в базе данных Vector DB, используя векторизированное предпочтение пользователя. Она находит наиболее релевантную книгу в коллекции на основе сходства векторов.

Выполнение DAG#

Откройте терминал и выполните команду astro dev start.

Должен запуститься контейнер локальной версии Airflow. Теперь можно получить доступ к интерфейсу Airflow по адресу http://localhost:8080. Перейдите к DAG, нажав на books_recommend.

DAG

Нажмите кнопку PLAY справа, чтобы запустить DAG. Необходимо ввести предпочтения, причем значение по умолчанию будет заполнено.

Предпочтения

После завершения выполнения DAG можно увидеть результат поиска в логах задачи search_qdrant.

Результат