doczyai-pipelines/fieldExtraction/scripts/create_proc_code_embeddings.py

import src.utils.embedding_utils as embedding_utils
from src.investment.code_funcs import get_proc_crosswalk, get_clean_value
from crosswalk.crosswalk_utils import CrosswalkBuilder

import faiss
import pickle

import os
import pandas as pd
import numpy as np

from sentence_transformers import SentenceTransformer

# Load model
roberta_model = SentenceTransformer("all-roberta-large-v1")

# If running proc_code for the first time
"""
mapping_dir = "crosswalk/mapping_csvs/proc_cd"
pkl_dir = "embeddings"
for filename in os.listdir(mapping_dir):
    if filename.endswith(".csv"):
        print(f"Creating embeddings for {filename}")
        stripped_filename = filename.strip(".csv")
        mapping_df = pd.read_csv(os.path.join(mapping_dir, filename))
        proc_crosswalk = CrosswalkBuilder().from_df(mapping_df, from_col="Code", to_col="Description")
        proc_choices = [get_clean_value(x) for x in proc_crosswalk.mapping.values()]
        os.makedirs(os.path.join(pkl_dir, stripped_filename), exist_ok=True)
        proc_index = embedding_utils.create_faiss_index(choices=proc_choices, 
                                                        model=roberta_model, 
                                                        save_path=os.path.join(pkl_dir, stripped_filename, "faiss_index.bin"),
                                                        embedding_path=os.path.join(pkl_dir, stripped_filename, "embeddings.npy"),
                                                        choices_path=os.path.join(pkl_dir, stripped_filename, "choices.pkl")
                                                        )        
"""
# If running diag_code for the first time
mapping_dir = "crosswalk/mapping_csvs/diag_cd"
pkl_dir = "embeddings"
for filename in os.listdir(mapping_dir):
    if filename.endswith(".csv"):
        print(f"Creating embeddings for {filename}")
        stripped_filename = filename.strip(".csv")
        mapping_df = pd.read_csv(os.path.join(mapping_dir, filename))
        mapping_df = mapping_df.astype(str)
        diag_crosswalk = CrosswalkBuilder().from_df(mapping_df, from_col="Code", to_col="Description")
        diag_choices = [get_clean_value(x) for x in diag_crosswalk.mapping.values()]
        os.makedirs(os.path.join(pkl_dir, stripped_filename), exist_ok=True)
        proc_index = embedding_utils.create_faiss_index(choices=diag_choices,
                                                        model=roberta_model,
                                                        save_path=os.path.join(pkl_dir, stripped_filename, "faiss_index.bin"),
                                                        embedding_path=os.path.join(pkl_dir, stripped_filename, "embeddings.npy"),
                                                        choices_path=os.path.join(pkl_dir, stripped_filename, "choices.pkl")
                                                        )