Jupyter batch processing embeddings for chunks

jupyter_embeddings_batch_chunks.py

# Input: folder chunks/chunk_*
# Output: folder chunks/chunk_*_embeddings
# Deletes trash/cache for each iteration to free disk space

import subprocess

def run_shell_command(command):
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    output, error = process.communicate()
    
    if process.returncode != 0:
        print(f"Error executing command: {command}")
        print(f"Error message: {error.decode('utf-8')}")
    else:
        print(f"Command executed successfully: {command}")
        print(f"Output: {output.decode('utf-8')}")

def delete_trash_and_cache():
    # Check disk usage before deletion
    run_shell_command("du -sh ~/.local/share/Trash")

    # Run the deletion command
    run_shell_command("rm -r ~/.local/share/Trash/*")

    # Check disk usage after deletion
    run_shell_command("du -sh ~/.local/share/Trash")

    # Check disk usage before deletion
    run_shell_command("du -sh ~/.cache")

    # Run the deletion command
    run_shell_command("rm -r ~/.cache/*")

    # Check disk usage after deletion
    run_shell_command("du -sh ~/.cache")

# Run the function
delete_trash_and_cache()

from tqdm import tqdm 
import pandas as pd 
import numpy as np
import os
tqdm.pandas()
from FlagEmbedding import FlagModel

model = FlagModel('BAAI/bge-base-en-v1.5')
chunks = [8,9,10,11,12,13,14]

for chunk in chunks:
    print(f"chunk: {chunk}")

    df = pd.read_parquet(f"chunks/chunk_{chunk}.parquet") # luckily, parsing worked just fine!
    #df = df.iloc[:10]

    def split_into_chunks(text, chunk_size=300):
        # Check for NaN values
        if text is None or not isinstance(text, str) or text == "":
            return []

        words = text.split()

        # Handle empty text
        if not words:
            return []

        chunks = [words[i:i + chunk_size] for i in range(0, len(words), chunk_size)]
        return [' '.join(chunk) for chunk in chunks]


    df["content_chunk"] = df["content"].progress_apply(lambda x: split_into_chunks(x, 300)) # 35 sec

    # Explode the list of chunks into separate rows
    df_expanded = df.explode('content_chunk')

    # Convert 'content_chunk' column to string
    df_expanded['content_chunk'] = df_expanded['content_chunk'].astype(str)

    # Drop the original 'content' column if needed
    df_expanded = df_expanded.drop('content', axis=1)

    # Display the resulting DataFrame
    #df_expanded
    print("embedding generation")
    embeddings = model.encode(df_expanded["content_chunk"].to_list())#.to_list())


    df_expanded["embeddings"] = list(embeddings)

    # Assuming 'embeddings' is the name of the column with arrays
    df_expanded['embeddings'] = df_expanded['embeddings'].apply(lambda x: x.astype(np.float32) if x is not None else None)

    print("saving")
    df_expanded.to_parquet( f"chunks/chunk_{chunk}_embeddings.parquet")

    file_path = f"chunks/chunk_{chunk}.parquet"
    if os.path.exists(file_path):
        os.remove(file_path)
        print(f"Deleted: {file_path}")
    else:
        print(f"File not found: {file_path}")
    delete_trash_and_cache()


    #pd.read_parquet( f"chunks/chunk_{chunk}_embeddings.parquet")