Mayank Aamseek
54b3b74675
Daip2 111 * indirect proc codes added * bug fix * bug fix * Merge branch 'main' into proc-code-indirect * bug fix * diag cd indirect added * Merge branch 'main' into daip2-111 * bug fix * deleted missed print statements * Merge branch 'main' into daip2-111 merged with main * Merge branch 'main' into daip2-111 pulled from main * sync embeddings from s3 * print statement removed * typehint added Approved-by: Alex Galarce
54 lines
2.8 KiB
Python
54 lines
2.8 KiB
Python
|
|
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")
|
|
) |