from snowflake.snowpark.types import PandasSeries, PandasDataFrame

def read_file(filename):
    import joblib
    import sys
    import os

    #where all imports located at
    import_dir = sys._xoptions.get("snowflake_import_directory")

    if import_dir:
        with open(os.path.join(import_dir, filename), 'rb') as file:
            m = joblib.load(file)
            return m

#register UDF
@F.udf(name = 'predict_risk_score', is_permanent = True, replace = True, stage_location = '@ML_MODELS')
def predict_risk_score(ds: PandasSeries[dict]) -> PandasSeries[float]:

    # later we will input train data as JSON object
    # hance, we have to convert JSON object as pandas DF
    df = pd.io.json.json_normalize(ds)[feature_cols]
    pipeline = read_file('predict_risk_score.joblib')
    return pipeline.predict_proba(df)[:,1]