Embedding Text with local (per node) NVIDIA TensorRT accelerator and GPU based Aproximate Nearest Neighbor (ANN)

The demo extends the existing Azure OpenAI based demo where encoding is processed by OpenAI requests and KNN was using GPU based brute force search. This tutorial shows how to perform fast local embeddings using multilingual E5 text embeddings and fast approximate Nearest Neighbor search using the IVFFlat algorithm. All tutorial stages are accelerated by an NVIDIA GPU using NVIDIA TensorRT and Spark Rapids ML. The tutorial folder contains two benchmark notebooks to demonstrate advantages of the presented GPU based approach compared to the previous CPU based demo

The key prerequisites for this quickstart include a working Azure OpenAI resource, and an Apache Spark cluster with SynapseML installed. We suggest creating a Synapse workspace, but the notebook was run on a Databricks GPU based cluster using Standard_NC24ads_A100_v4 with 6 workers. Databricks Runtime was 13.3 LTS ML (includes Apache Spark 3.4.1, GPU, Scala 2.12) with the related init_script to install all required packages.

Step 1: Prepare Environment

It will import the required libraries and get the initial settings.

import torch
import sys
import pyspark.sql.functions as F
from pyspark.sql.types import StructType, StructField, IntegerType, StringType
from pyspark.ml.linalg import Vectors
from pyspark.ml.linalg import VectorUDT
from spark_rapids_ml.knn import (
    ApproximateNearestNeighbors,
    ApproximateNearestNeighborsModel,
)
from synapse.ml.hf import HuggingFaceSentenceEmbedder
from synapse.ml.nn import KNN

Step 2: Load Input Data

It will load a public dataset and generate extra synthetic rows if set by the size parameter.

The loaded dataset has 1000 rows. If you specify number_of_input_rows in [1..1000] it will cut extra rows if needed

If number_of_input_rows is in [1000..1000000], it will generate extra rows using a cross join of the original data

file_path = "wasbs://publicwasb@mmlspark.blob.core.windows.net/fine_food_reviews_1k.csv"

df = spark.read.options(inferSchema="True", delimiter=",", header=True).csv(file_path)
df = df.withColumn(
    "data",
    F.format_string("Title: %s; Content: %s", F.trim(df.Summary), F.trim(df.Text)),
)

# Size of DF
number_of_input_rows = 100


# Check if the row count is less than 10
if number_of_input_rows <= 0 or number_of_input_rows >= 1000000:
    raise ValueError(f"Limit is {number_of_input_rows}, which should be less than 1M.")

if number_of_input_rows > 1000:

    # Cross-join the DataFrame with itself to create n x n pairs for string concatenation (synthetic data)
    cross_joined_df = df.crossJoin(df.withColumnRenamed("data", "data_"))

    # Create a new column 'result_vector' by concatenating the two source vectors
    tmp_df = cross_joined_df.withColumn(
        "result_vector",
        F.concat(F.col("data"), F.lit(". \n"), F.col("data_")),
    )

    # Select only the necessary columns and show the result
    tmp_df = tmp_df.select("result_vector")

    # Shuffle the DataFrame with a fixed seed to have close strings spread
    seed = 42

    df = (
        tmp_df.withColumnRenamed("result_vector", "data")
        .withColumn("id", F.monotonically_increasing_id())
        .orderBy(F.rand(seed))
    )

df = df.limit(number_of_input_rows).repartition(10).cache()

print(f"Loaded: {number_of_input_rows} rows")

Step 3: Generate Embeddings

We will first generate embeddings using NVIDIA TensorRT optimized SentenceTransformer. In the demo you can use two different HF models: intfloat/e5-large-v2 or sentence-transformers/all-MiniLM-L6-v2.

# To create embedder with different models, uncomment the following line
# embedder = HuggingFaceSentenceEmbedder(modelName="intfloat/e5-large-v2", inputCol="data", outputCol="embeddings", runtime="tensorrt")
embedder = HuggingFaceSentenceEmbedder(
    modelName="sentence-transformers/all-MiniLM-L6-v2",
    inputCol="data",
    outputCol="embeddings",
    runtime="tensorrt",
)

embeddings = embedder.transform(df).select("id", "embeddings").cache()

Step 4: Build the query against embeddings

Get query embeddings running standard SentenceTransformer just on the driver. Convert embedding results to a data frame

# Sample query
queries = ["desserts", "disgusting"]
ids = [1, 2]

# Create DataFrame directly from the data and schema
query_df = spark.createDataFrame(
    list(zip(ids, queries)),
    StructType(
        [
            StructField("id", IntegerType(), nullable=False),
            StructField("data", StringType(), nullable=False),
        ]
    ),
)

query_embeddings = embedder.transform(query_df).select("id", "embeddings").cache()

Step 5: Build a fast vector index to over review embeddings

We will use fast NVIDIA Rapids indexer. This KNN implementation will work only on GPU. If you want to use CPU then switch to synapse.ml.nn CPU based KNN implementation

RUN_ON_GPU = torch.cuda.is_available()

if RUN_ON_GPU:
    rapids_knn_model = (
        ApproximateNearestNeighbors(k=5)
        .setInputCol("embeddings")
        .setIdCol("id")
        .fit(embeddings)
    )
else:
    array_to_vector_udf = udf(lambda array: Vectors.dense(array), VectorUDT())
    df_with_vectors = embeddings.withColumn(
        "features", array_to_vector_udf(embeddings["embeddings"])
    )
    knn = (
        KNN()
        .setFeaturesCol("features")
        .setValuesCol("id")
        .setOutputCol("output")
        .setK(10)
    )
    knn_model = knn.fit(df_with_vectors)

Step 6: Find top k Nearest Neighbors ON GPU

We will use fast ANN IVFFlat algorithm from Rapids

if RUN_ON_GPU:
    (_, _, knn_df) = rapids_knn_model.kneighbors(query_embeddings)
else:
    array_to_vector_udf = udf(lambda array: Vectors.dense(array), VectorUDT())
    df_with_vectors = query_embeddings.withColumn(
        "features", array_to_vector_udf(query_embeddings["embeddings"])
    )
    knn_df = knn_model.transform(df_with_vectors)

Step 7: Collect and display results

if RUN_ON_GPU:
    result_df = (
        knn_df.withColumn(
            "zipped", F.explode(F.arrays_zip(F.col("indices"), F.col("distances")))
        )
        .select(
            F.col("query_id"),
            F.col("zipped.indices").alias("id"),
            F.col("zipped.distances").alias("distance"),
        )
        .join(df, on="id", how="inner")
        .select("query_id", "id", "data", "distance")
    )
else:
    knn_df = knn_df.withColumnRenamed("data", "original_data")
    result_df = (
        knn_df.withColumn("match", F.explode("output"))
        .join(df, df["id"] == F.col("match.value"))
        .select("original_data", F.col("data"), "match.distance")
    )

display(result_df)

Results

The goal of this demo is to showcase two acceleration techniques: local (per node) embedding generation and approximate KNN. Compared to the original method, which relies on HTTP requests to the OpenAI model and CPU-based KNN, the new approach is significantly more scalable and provides substantial acceleration, especially for large input datasets.

These are the duration comparison results on 10 T4 GPU nodes for both approaches: