doczyai-pipelines/src/utils/embedding_utils.py

import os
import pickle

import faiss
import numpy as np
import src.config as config
from src.utils.string_utils import datetime_str


def create_faiss_index(
    choices,
    model,
    save_path="faiss_index.bin",
    embedding_path="embeddings.npy",
    choices_path="choices.pkl",
):
    embeddings = model.encode(choices, normalize_embeddings=True).astype("float32")

    # Create FAISS index
    index = faiss.IndexFlatIP(embeddings.shape[1])
    index.add(embeddings)

    # Save index
    faiss.write_index(index, save_path)

    # Save embeddings and choices for future use
    np.save(embedding_path, embeddings)
    with open(choices_path, "wb") as f:
        pickle.dump(choices, f)

    return index


def load_faiss_index(
    index_path="faiss_index.bin",
    embedding_path="embeddings.npy",
    choices_path="choices.pkl",
):
    index = faiss.read_index(index_path)
    embeddings = np.load(embedding_path)
    with open(choices_path, "rb") as f:
        choices = pickle.load(f)
    return index, embeddings, choices


def load_embeddings():
    """loads all embedding files required for processing indirect codes (e.g. proc codes)"""

    # list all embeddings present in s3
    s3_client = config.S3_CLIENT
    response = s3_client.list_objects_v2(Bucket="doczy-investment", Prefix="embeddings")
    file_list = [
        obj["Key"]
        for obj in response.get("Contents", [])
        if not obj["Key"].endswith("/")
    ]

    # list embeddings already present locally
    if not os.path.exists("embeddings"):
        os.makedirs("embeddings")
    local_files = []
    for root, dirs, files in os.walk("embeddings"):
        for name in files:
            local_files.append(os.path.join(root, name))

    # download embeddings not present locally
    for file in file_list:
        local_path = file

        # create directory structure for this file
        directory = os.path.dirname(local_path)
        if not os.path.exists(directory):
            os.makedirs(directory)

        # check if file already exists
        if not os.path.exists(local_path):
            try:
                print(f"{datetime_str()} Downloading {file} to {local_path}...")
                s3_client.download_file("doczy-investment", file, local_path)
            except Exception as e:
                print(f"{datetime_str()} Error downloading {file}: {e}")