00-Document-Ingestion-and-Index.py

# Databricks notebook source
# MAGIC %md
# MAGIC # Document Ingestion and Index

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

pip install langchain transformers mlflow[genai]>=2.9.0 databricks-vectorsearch

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

dbutils.library.restartPython()

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

import os
import time
import requests
import pandas as pd

from pyspark.sql.functions import pandas_udf, explode
from pyspark.sql.types import StringType

from langchain.text_splitter import RecursiveCharacterTextSplitter
from transformers import AutoTokenizer

import mlflow
import mlflow.deployments
from databricks.vector_search.client import VectorSearchClient
from databricks.sdk import WorkspaceClient
import databricks.sdk.service.catalog as c

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

# MAGIC %sql
# MAGIC CREATE CATALOG IF NOT EXISTS leahey_sandbox;
# MAGIC USE CATALOG leahey_sandbox;
# MAGIC CREATE SCHEMA IF NOT EXISTS tgn_llm_qa;
# MAGIC USE SCHEMA tgn_llm_qa;

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

# MAGIC %md
# MAGIC ## Download episode transcripts
# MAGIC
# MAGIC Download transcripts from https://www.phfactor.net/tgn and save each in DBFS

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

urls = [f"https://www.phfactor.net/tgn/{i}.0/episode.txt" for i in range(1,265)] 
df_urls = spark.createDataFrame(urls, StringType()).toDF("url")

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

def download_transcript(url: str):
    transcript = requests.get(url, verify=False).text
    transcript.replace("\'", "'").replace('"',"").replace("\n","")
    return transcript

@pandas_udf("string")
def download_transcripts_udf(urls: pd.Series) -> pd.Series:
    return urls.apply(download_transcript)

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

# MAGIC %md
# MAGIC ## Define UDFs:
# MAGIC - Split transcripts into chunks using Langchain RecursiveCharacterTextSplitter
# MAGIC - Create embeddings from chunks

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

tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(tokenizer, chunk_size=500, chunk_overlap=50)

def split_text(text: str):
    return text_splitter.split_text(text)

@pandas_udf("array<string>")
def split_text_udf(text: pd.Series) -> pd.Series:
    return text.apply(split_text)

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

deploy_client = mlflow.deployments.get_deploy_client("databricks")

@pandas_udf("array<float>")
def get_embedding(contents: pd.Series) -> pd.Series:
    def get_embeddings(batch):
        response = deploy_client.predict(endpoint="databricks-bge-large-en", inputs={"input": batch})
        return [e["embedding"] for e in response.data]

    max_batch_size = 150
    batches = [contents.iloc[i:i + max_batch_size] for i in range(0, len(contents), max_batch_size)]

    all_embeddings = []
    for batch in batches:
        all_embeddings += get_embeddings(batch.tolist())

    return pd.Series(all_embeddings)

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

# MAGIC %md
# MAGIC ## Create Delta table containing embeddings

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

# MAGIC %sql
# MAGIC CREATE TABLE IF NOT EXISTS tgn_scripts (
# MAGIC   id BIGINT GENERATED BY DEFAULT AS IDENTITY,
# MAGIC   url STRING,
# MAGIC   content STRING,
# MAGIC   embedding ARRAY <FLOAT>
# MAGIC ) TBLPROPERTIES (delta.enableChangeDataFeed = true); 

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

(
   df_urls
    .withColumn("transcript", download_transcripts_udf("url"))
    .withColumn("content", explode(split_text_udf("transcript")))
    .withColumn("embedding", get_embedding("content"))
    .drop("transcript")
    .write.mode("overwrite").saveAsTable("tgn_scripts")
)

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

# MAGIC %md
# MAGIC ## Create Vector Search endpoint

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

vsc = VectorSearchClient()

vector_search_endpoint_name = "tgn_vs_endpoint"

if vector_search_endpoint_name not in [e['name'] for e in vsc.list_endpoints().get('endpoints', [])]:
    vsc.create_endpoint(name=vector_search_endpoint_name, endpoint_type="STANDARD")

for i in range(180):
    endpoint = vsc.get_endpoint(vector_search_endpoint_name)
    status = endpoint.get("endpoint_status", endpoint.get("status"))["state"].upper()
    if "ONLINE" in status:
        break
    elif "PROVISIONING" in status or i <6:
        if i % 20 == 0:
            print(f"Waiting for endpoint to be ready: {endpoint}")
        time.sleep(10)
    else:
        raise Exception(f"Error with the endpoint {vector_search_endpoint_name}.")

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

# MAGIC %md
# MAGIC ## Create or Sync Vector Search Index

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

table_full_name = "leahey_sandbox.tgn_llm_qa.tgn_scripts"
index_full_name = "leahey_sandbox.tgn_llm_qa.tgn_vs_index"

def index_exists():
    try:
        dict_vsindex = vsc.get_index(vector_search_endpoint_name, index_full_name).describe()
        return dict_vsindex.get('status').get('ready', False)
    except Exception as e:
        if 'RESOURCE_DOES_NOT_EXIST' not in str(e):
            print(f'Unexpected error occurred. Could be a permissions error.')
            raise e
    return False

if not index_exists():
  vsc.create_delta_sync_index(
    endpoint_name=vector_search_endpoint_name,
    index_name=index_full_name,
    source_table_name=table_full_name,
    pipeline_type="TRIGGERED",
    primary_key="id",
    embedding_dimension=1024, #Match model embedding size (bge)
    embedding_vector_column="embedding"
  )
else:
  vsc.get_index(vector_search_endpoint_name, index_full_name).sync()

def wait_for_index_to_be_ready():
  for i in range(180):
    idx = vsc.get_index(vector_search_endpoint_name, index_full_name).describe()
    index_status = idx.get('status', idx.get('index_status', {}))
    status = index_status.get('detailed_state', index_status.get('status', 'UNKNOWN')).upper()
    url = index_status.get('index_url', index_status.get('url', 'UNKNOWN'))
    if "ONLINE" in status:
      return
    if "UNKNOWN" in status:
      print(f"Can't get the status - will assume index is ready {idx} - url: {url}")
      return
    elif "PROVISIONING" in status:
      if i % 40 == 0: print(f"Waiting for index to be ready, this can take a few min... {index_status} - pipeline url:{url}")
      time.sleep(10)
    else:
        raise Exception(f"Error with the index {index_full_name}")
  raise Exception(f"Timeout, your index isn't ready yet: {vsc.get_index(index_full_name, vector_search_endpoint_name)}")

wait_for_index_to_be_ready()