How to Split Text For Vector Embeddings in Snowflake

				
					
create or replace function fn_token_based_text_chunker(textdata string, max_token NUMBER , overlap_size NUMBER)
returns ARRAY
language python
runtime_version = '3.9'
handler = 'bert_text_chunker'
IMPORTS=('@TOKENIZER/google-bert.zip')
packages = ('transformers')
as
$$
from transformers import AutoTokenizer

import fcntl
import os
import sys
import threading
import zipfile

# File lock class for synchronizing write access to /tmp
class FileLock:
   def __enter__(self):
      self._lock = threading.Lock()
      self._lock.acquire()
      self._fd = open('https://i0.wp.com/www.phdata.io/tmp/lockfile.LOCK', 'w+')
      fcntl.lockf(self._fd, fcntl.LOCK_EX)

   def __exit__(self, type, value, traceback):
      self._fd.close()
      self._lock.release()
      
# Get the location of the import directory. Snowflake sets the import
# directory location so code can retrieve the location via sys._xoptions.
IMPORT_DIRECTORY_NAME = "snowflake_import_directory"
import_dir = sys._xoptions[IMPORT_DIRECTORY_NAME]

zip_file_path = import_dir + "google-bert.zip"
extracted = '/tmp/google_bert_dir'

# Extract the contents of the ZIP. This is done under the file lock
# to ensure that only one worker process unzips the contents.
with FileLock():
   if not os.path.isdir(extracted + '/google_bert_dir'):
      with zipfile.ZipFile(zip_file_path, 'r') as myzip:
         myzip.extractall(extracted)

        
def bert_text_chunker(textdata, max_tokens=512, chunk_overlap=50):
	# Initialize the BERT tokenizer
	model_name= 'google-bert/bert-base-uncased'
	tokenizer = AutoTokenizer.from_pretrained(extracted +'/'+model_name)
	#tokenizer=
	# Tokenize the text
	tokens = tokenizer.tokenize(textdata)
	
	# Initialize list to hold chunks
	chunks = []
	
	# Loop through tokens and create chunks with overlap
	for i in range(0, len(tokens), max_tokens - chunk_overlap):
		# Get the current chunk with overlap
		chunk = tokens[i:i + max_tokens]
		chunks.append(chunk)
		
		# Stop if we have reached the end of the tokens
		if i + max_tokens >= len(tokens):
			break
	
	# Detokenize each chunk back to text
	text_chunks = [tokenizer.convert_tokens_to_string(chunk) for chunk in chunks]
	
	return text_chunks

$$;