client.py 3.72 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
    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,
26
                        help="boolean indicating if visualization should be enabled.")
27
28
    parser.add_argument("--use-real-ats-names", type=str2bool, nargs='?',
                        const=True, default=False,
29
30
31
32
33
                        help="boolean indicating if we should use real ats names.")
    parser.add_argument("--run-full-pipeline", type=str2bool, nargs='?',
                        const=True, default=True,
                        help="boolean indicating if we should run full pipeline. " +
                        "set to false to only make models")
34
35
36
37
    return parser.parse_args()

def main():
    parsed_args = parse_arguments()
38
39
40
41
    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
42
    run_full_pipeline = parsed_args.run_full_pipeline
43
44
45
46
47
    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}")
48
49
50
51
52
53
54
55
56
    print(f"Use real ATS names: {use_real_ats_names}")
    print(f"Run full pipeline: {run_full_pipeline}\n")
    
    if run_full_pipeline:
        print("Now parsing and cleaning data ...")        
        if dataset_year == '2019':
            parse_and_clean_data.main(year=dataset_year)
        else:
            parse_and_clean_data.main()
57
        
58
59
        print("Extracting screenings ...")
        make_screenings.main()
60
    
61
62
63
        print("Making clusters ...")
        make_clusters.main()
        print(f"Completed making cluster model. It can be found at: {cfg.CLUSTERS_DIR}\n")
64
    
65
66
        print("Making full dataset ...")
        make_dataset_full.main(use_real_ats_names)
67
    
68
69
70
71
72
73
        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()
74
            
75
76
77
        print("\nCompleted generating datasets at:")
        print(f"Interim data dictionary: {cfg.INTERIM_DATA_DIR}")
        print(f"Processed data dictionary: {cfg.PROCESSED_DATA_DIR}\n")
78
    
79
    print(f"Making 4 XGBoost models based on version: {dataset_version} ...\n")
80
81
    make_xgb_models.main(dataset_version)
    
82
83
84
    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")
85
86
87
    
if __name__ == "__main__":
    main()