client.py 3.21 KB
Newer Older
1
import argparse
2
import config as cfg
3
4
5
from src.data import parse_and_clean_data, make_screenings
from src.data import make_clusters, make_dataset_full
from src.data import make_dataset_count, make_dataset_emb
6
from src.model import make_xgb_models
7

8
9
10
11
12
13
14
def str2bool(v):
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
15
    else:
16
        raise argparse.ArgumentTypeError('Boolean value expected.')
17
18
19

def parse_arguments():
    parser = argparse.ArgumentParser(description='A client for AIR')
20
21
22
23
24
25
26
27
28
29
    parser.add_argument('--dataset-year', type=str, default="2020",
                        choices=['2019', '2020'], help='string indicating dataset year')
    parser.add_argument('--dataset-version', type=str, default="emb",
                        choices=['emb', 'ohe'], help='string indicating dataset version')
    parser.add_argument("--enable-visualization", type=str2bool, nargs='?',
                        const=True, default=False,
                        help="bolean indicating if visualization should be enabled.")
    parser.add_argument("--use-real-ats-names", type=str2bool, nargs='?',
                        const=True, default=False,
                        help="bolean indicating if we should use real ats names.")
30
31
32
33
    return parser.parse_args()

def main():
    parsed_args = parse_arguments()
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
    dataset_year = parsed_args.dataset_year
    dataset_version = parsed_args.dataset_version
    enable_visualization = parsed_args.enable_visualization
    use_real_ats_names = parsed_args.use_real_ats_names
    print(f"Client started. Using this configuration:")
    print(f"Raw data dictionary: {cfg.RAW_DATA_DIR_2020}")
    print(f"Dataset year: {dataset_year}")
    print(f"Dataset version: {dataset_version}")
    print(f"Visualization enabled: {enable_visualization}")
    print(f"Use real ATS names: {use_real_ats_names}\n")
    print("Now parsing and cleaning data ...")
    if dataset_year == '2019':
        parse_and_clean_data.main(year=dataset_year)
    else:
        parse_and_clean_data.main()
        
50
51
52
53
54
    print("Extracting screenings ...")
    make_screenings.main()
    
    print("Making clusters ...")
    make_clusters.main()
55
    print(f"Completed making cluster model. It can be found at: {cfg.CLUSTERS_DIR}\n")
56
57
    
    print("Making full dataset ...")
58
59
60
61
62
63
64
65
66
67
68
69
    make_dataset_full.main(use_real_ats_names)
    
    if dataset_version == "emb":
        print("Making dataset with embedded ats ...")
        make_dataset_emb.main(enable_visualization)
    else:
        print("Making dataset with one-hot encoded ats ...")
        make_dataset_count.main()
    
    print("\nCompleted generating datasets at:")
    print(f"Interim data dictionary: {cfg.INTERIM_DATA_DIR}")
    print(f"Processed data dictionary: {cfg.PROCESSED_DATA_DIR}\n")
70
    
71
    print(f"Now making 4 XGBoost models based on version: {dataset_version} ...")
72
    make_xgb_models.main()
73
    
74
75
76
    print(f"Completed making models. Models and SHAP plots can be found at:\n" +
          f"{cfg.COMPLETE_XGB_DIR}\n" + f"{cfg.COMPLIANCE_XGB_DIR}\n" +
          f"{cfg.FALL_XGB_DIR}\n" + f"{cfg.FALL_TEST_XGB_DIR}" + "\n")
77
78
79
    
if __name__ == "__main__":
    main()