Vector Indexing for Retrieval-Augmented Generation with Databricks, SAP HANA, and Generative AI Hub

[yxlee12345]

Introduction

In the rapidly evolving landscape of artificial intelligence, the combination of vector indexing and retrieval-augmented generation (RAG) has emerged as a powerful approach for enhancing the accuracy and relevance of generated content. By leveraging the capabilities of Databricks, SAP HANA Database, and tools for Generative AI like the Generative AI Hub SDK and LangChain, we can create a robust framework for efficient and scalable vector indexing.

In my previous blog post, we saw how LangChain could be used with SAP HANA for vector indexing and RAG. In this blog post, we will explore the step-by-step process of setting up a vector indexing pipeline using Databricks for data processing, SAP HANA as the vector database, and the Generative AI Hub SDK alongside LangChain for seamless integration. At each step, code examples for a Databricks notebook will be provided. Whether you're a data scientist, AI researcher, or a developer looking to enhance your RAG workflows, this guide will provide you with the insights and practical knowledge needed to get started. Let's dive into the intricacies of this powerful combination and unlock new potentials for your AI-driven projects.

Setting Up Our Environment

Databricks clusters come with a good number of pre-installed packages for data science, data engineering and AI purposes, so tools like Spark and LangChain are already available by default. However, there are some packages that we will need for the example pipeline that we will be building:

• html2text: For converting text from HTML to Markdown format
• generative-ai-hub-sdk: For working with Generative AI models in the Generative AI Hub
• hdbcli: For connection to a HANA Database

The following code snippet could be executed in Databricks to install these packages:

%pip install html2text "generative-ai-hub-sdk[all]" hdbcli

dbutils.library.restartPython()

Note: While generative-ai-hub-sdk and hdbcli are required for the vector indexing process, html2text is only required for this example. Your use cases might not necessarily require html2text, but might require other packages instead, so you should adjust the pip install command accordingly to fit your use cases.

Building Our Vector Indexing Pipeline

Building a vector indexing pipeline is essentially a data engineering task, and as such, we could split the pipeline processes into the Extract, Transform, and Load (ETL) steps.

Extract

The extract step involves ingesting raw data from one or more data sources. For our example, our raw data are HTML files exported from an Atlassian Confluence wiki space, where each HTML file contains the contents of a single wiki page. These files are stored in a container within an Azure Data Lake Service (ADLS) account, and can be read using PySpark as shown:

from pyspark.sql.functions import expr

HTML_FILEPATHS_GLOB_PATTERN = "abfss://[email protected]/sandbox/example_wikis/*"

html_binary_spark_df = (
    spark
    .read
    .format("binaryFile")
    .options(pathGlobFilter="*.html")
    .load(HTML_FILEPATHS_GLOB_PATTERN)
    .withColumn("content", expr("CAST(content AS STRING)"))
)

If the extract step goes well, we should get a PySpark DataFrame containing these columns:

• path: the full paths to the HTML files, which look something like abfss://[email protected]/sandbox/example_wikis/GAI01.03.10-RAGe_4337208805.html
• modificationTime: timestamp values indicating when the files were last modified
• length: length of the file contents
• content: decoded HTML contents of the files

Note:

• The snippet above shows how we can ingest data that exists as raw binary files in ADLS, but your use cases might involve data in different formats (e.g. CSV, Parquet) or from different sources (e.g. database, API), so you will need to adjust your ingestion logic accordingly. The documentation for working with inputs and outputs using PySpark is a good place to figure out how you could use Spark for ingesting your data.
• In this example, HTML_FILEPATHS_GLOB_PATTERN points to a container in ADLS where the raw data is stored, with aitm being the container name, coredatalaketestint.dfs.core.windows.net being the address of the ADLS instance, and sandbox/example_wikis/ being the folder within the aitm container where the data is stored. For different cloud platforms (GCP, AWS etc), the URI would be different, so it is advised to check the relevant documentations from Spark or the respective cloud platforms.

Transform

After we have ingested our data, the next step is to clean and preprocess our data into a suitable format for vector indexing. For this example, we will perform the following to preprocess our data:

• Converting the contents from HTML to Markdown: Markdown is syntactically simpler compared to HTML, which makes it more understandable to Large Language Models (LLMs), thus we do this conversion
• Removing unnecessary contents: We will remove headers and footers in page contents to reduce noise in the data, which could help improve the retrieval performance of a RAG application. We also remove authors from the contents for privacy reasons.
• Getting the base file names: This is an optional step to get the base file names from the path values in the path column, for recording as metadata during the indexing step later.

These smaller steps can be done as shown:

import re
import html2text
from pyspark.sql.functions import udf, col
from pyspark.sql.types import StringType

AUTHORS_REGEX_PATTERN = re.compile(r"Created\sby.*last\smodified.*[A-Z][a-z]{2}\s\d{2}\,\s\d{4}")

convert_html_to_markdown = udf(html2text.html2text, StringType())

@udf(StringType())
def strip_header_and_footer(content: str) -> str:
    return content[53:-100]

@udf(StringType())
def remove_authors(content: str) -> str:
    return AUTHORS_REGEX_PATTERN.sub("", content)

@udf(StringType())
def get_filename(path: str) -> str:
    return path.split("/")[-1]

wiki_texts_spark_df = (
    html_binary_spark_df
    .withColumn("content", convert_html_to_markdown(col("content")))
    .withColumn("content", strip_header_and_footer(col("content")))
    .withColumn("content", remove_authors(col("content")))
    .withColumn("filename", get_filename(col("path")))
    .selectExpr("filename", "content")
)

If the cleaning and preprocessing steps go well, we should get a different PySpark DataFrame containing these columns:

• filename: the base name of the HTML files, which look something like GAI01.03.10-RAGe_4337208805.html
• content: the cleaned contents of the wiki pages in Markdown format

Note: Spark user-defined functions (UDFs) can be used to apply preprocessing steps to the columns of a PySpark DataFrame in a simple and clean fashion. More information about PySpark UDFs can be found in the PySpark documentation.

After cleaning our data, we will still need to split our wiki contents into chunks. This is especially important if we are working with LLMs that have lower token limits. We also typically want to have smaller chunks of text to reduce the risk of encountering the lost-in-the-middle problem during the RAG process.

For our example, we use LangChain to load the PySpark DataFrame and perform the chunking of wiki contents:

from langchain_community.document_loaders import PySparkDataFrameLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

full_contents_loader = PySparkDataFrameLoader(spark_session=spark, df=wiki_texts_spark_df, page_content_column="content")
documents_to_split = full_contents_loader.load()

document_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    model_name="gpt-4",
    chunk_size=200,
    chunk_overlap=20
)

document_chunks = document_splitter.split_documents(documents_to_split)

Note: While we use the RecursiveCharacterTextSplitter class from LangChain with chunk_size of 200 and chunk_overlap of 20 as an example, other chunking strategies might work better for your own use cases. Therefore, some exploration should be done to determine chunking strategies that work better for your use cases. For a starting place on chunking strategy, you can refer to this Medium post.

Load

Now that we have cleaned and preprocessed our data, we are ready for the indexing step itself.
As we are using an embedding model from the Generative AI Hub as well as the HANA Database for the vector indexing step, we need to load the secrets for these services. For this example, these secrets have been stored as Databricks secrets which can be loaded as shown:

import json

SECRET_SCOPE = "TEST_AITM_SCOPE"
gen_ai_hub_service_key = json.loads(dbutils.secrets.get(scope=SECRET_SCOPE, key="EXAMPLE_GENAI_HUB_SERVICE_KEY"))
hana_secrets = json.loads(dbutils.secrets.get(scope=SECRET_SCOPE, key="EXAMPLE_HANA_VECTOR_SECRETS"))

After loading the secrets, we want to initialize the client objects for Generative AI Hub and HANA Database respectively as shown:

# HANA DB
from hdbcli import dbapi

hana_conn = dbapi.connect(
    address=hana_secrets["host"],
    port=hana_secrets["port"],
    user=hana_secrets["user"],
    password=hana_secrets["password"],
    autocommit=True,
    sslTrustStore=hana_secrets["certificate"],
)


# GenAI Hub
import os
from gen_ai_hub.proxy.core.proxy_clients import get_proxy_client

os.environ["AICORE_AUTH_URL"] = gen_ai_hub_service_key["url"]
os.environ["AICORE_CLIENT_ID"] = gen_ai_hub_service_key["clientid"]
os.environ["AICORE_CLIENT_SECRET"] = gen_ai_hub_service_key["clientsecret"]
os.environ["AICORE_RESOURCE_GROUP"] = gen_ai_hub_service_key["appname"].split("!")[0]
os.environ["AICORE_BASE_URL"] = f"{gen_ai_hub_service_key['serviceurls']['AI_API_URL']}/v2"

proxy_client = get_proxy_client("gen-ai-hub")

We then want to initialize the LangChain objects for working with our embedding model and HANA DB:

from gen_ai_hub.proxy.langchain.init_models import init_embedding_model
from langchain_community.vectorstores import HanaDB

embeddings = init_embedding_model("text-embedding-ada-002", proxy_client=proxy_client)
hana_vectordb = HanaDB(embedding=embeddings, connection=hana_conn, table_name="DATABRICKS_HANA_EXAMPLE_VECTORSTORE")

Note: My previous blog post gives more explanation of this code snippet, so you may want to check it out.

The last bit of the indexing process is a simple method call to write the chunks to HANA DB. LangChain will take care of the embedding of chunks and writing to HANA under the hood.

hana_vectordb.add_documents(document_chunks)

Optionally, we can run the following code snippet to test the retrieval of document chunks, as a sanity check:

hana_vector_retriever = hana_vectordb.as_retriever()

hana_vector_retriever.get_relevant_documents("What is RAG?")

What Next?

Now that we have our pipeline notebook, the typical next step would be to make our pipeline scalable by creating a Databricks job that runs that notebook. Configurations such as trigger schedules or notifications depend on the use case and can be set in Databricks accordingly. More information can be found in the Databricks documentation for jobs.
The vectors that are indexed can be consumed by a Retrieval-Augmented Generation (RAG) application to enhance the relevance and accuracy of generated content. For an example of how to consume these vectors using LangChain, refer to my previous blog post where I demonstrate the integration and utilization of these vectors in a RAG workflow.

Conclusion

In this blog post, we've delved into the intricate process of setting up a vector indexing pipeline using Databricks, SAP HANA, and Generative AI tools like the Generative AI Hub SDK and LangChain. By carefully walking through each step—from data extraction and transformation to loading and indexing—we've demonstrated how to build a scalable and efficient pipeline that significantly enhances the accuracy and relevance of AI-generated content. This powerful combination of technologies not only streamlines your RAG workflows but also opens new avenues for innovation in AI-driven projects.  Feel free to leave your comments and share your thoughts or questions about this blog post.
Happy coding and may your AI projects reach new heights of excellence! 🥂