1 - Preparing Data for RAG.py

# Databricks notebook source
# MAGIC %run ./Includes/Setup

# COMMAND ----------

# MAGIC %md
# MAGIC # Data Exploration

# COMMAND ----------

source_path = "s3://dalhussein-odsc/papers/"
articles_path = 'dbfs:/mnt/odsc/papers'

download_dataset(source_path, articles_path)

# COMMAND ----------

files = dbutils.fs.ls(articles_path)
display(files)

# COMMAND ----------

# MAGIC %md
# MAGIC # Bronze

# COMMAND ----------

df_raw = (spark.read.format("binaryfile")
                .load(articles_path)
                .withColumnRenamed("path", "doc_uri")
      )
display(df_raw)

# COMMAND ----------

bronze_table_name = f"{catalog_name}.{schema_name}.bronze_articles_raw"
df_raw.write.mode("overwrite").saveAsTable(bronze_table_name)

# COMMAND ----------

# MAGIC %md
# MAGIC # Silver

# COMMAND ----------

import fitz  # PyMuPDF

def pdf_to_text(pdf_content):
    """
        Convert PDF to text using PyMuPDF
    """
    doc = fitz.open(stream=pdf_content, filetype="pdf")
    text = ""
    for page in doc:
        text += page.get_text()
    return text

# COMMAND ----------


from pyspark.sql.types import StringType

pdf_to_text_udf = udf(pdf_to_text, StringType())

df_parsed = (df_raw.withColumn("text", pdf_to_text_udf("content"))
                   .select("doc_uri", "text")
            )
display(df_parsed)

# COMMAND ----------

silver_table_name = f"{catalog_name}.{schema_name}.silver_articles_parsed"
df_parsed.write.mode("overwrite").saveAsTable(silver_table_name)

# COMMAND ----------

# MAGIC %md
# MAGIC # Gold

# COMMAND ----------

import pandas as pd 
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import Document
from llama_index.core.utils import set_global_tokenizer
from transformers import AutoTokenizer
from typing import Iterator
from pyspark.sql.functions import pandas_udf, explode

@pandas_udf("array<string>")
def read_as_chunk(batch_iter: Iterator[pd.Series]) -> Iterator[pd.Series]:
    # set llama as tokenizer
    set_global_tokenizer(
      AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
    )
    # sentence splitter from llama_index to split on sentences
    splitter = SentenceSplitter(chunk_size=500, chunk_overlap=50)
    def extract_and_split(col):
      nodes = splitter.get_nodes_from_documents([Document(text=col)])
      return [n.text for n in nodes]

    for x in batch_iter:
        yield x.apply(extract_and_split)

# COMMAND ----------

gold_table_name = f"{catalog_name}.{schema_name}.gold_articles_chunks"

spark.sql(f"""
            CREATE TABLE IF NOT EXISTS {gold_table_name} (
            chunk_id BIGINT GENERATED BY DEFAULT AS IDENTITY,
            chunked_text STRING,
            doc_uri STRING
            -- NOTE: the table has to be CDC because VectorSearch is using DLT that is requiring CDC state
            ) TBLPROPERTIES (delta.enableChangeDataFeed = true);
          """)

# COMMAND ----------

df_chunks = (df_parsed.withColumn("chunked_text", explode(read_as_chunk("text")))
                .select("chunked_text", "doc_uri")
            )
display(df_chunks)

# COMMAND ----------

df_chunks.write.mode("overwrite").saveAsTable(gold_table_name)

# COMMAND ----------

# MAGIC %md
# MAGIC