Fix CLI argument passing

marcai/cli.py

@@ -33,9 +33,7 @@ def main():
     find_matches_parser.set_defaults(func=find_matches.main)
 
     args = parser.parse_args()
-
     args.func(args)
-   
 
 
 if __name__ == "__main__":

marcai/find_matches.py

@@ -35,10 +35,7 @@ def args_parser():
     return parser
 
 
-def main():
-
-    args = args_parser().parse_args()
-
+def main(args):
     config_path = f"{args.model_dir}/config.yaml"
     model_onnx = f"{args.model_dir}/model.onnx"
 
@@ -59,9 +56,9 @@ def main():
     with open(args.pair_indices, "r") as indices_file:
         reader = csv.reader(indices_file)
         # Process records
-        for df in tqdm(multiprocess_pairs(
-            records_df, reader, args.chunksize, args.processes
-        )):
+        for df in tqdm(
+            multiprocess_pairs(records_df, reader, args.chunksize, args.processes)
+        ):
             input_df = df[config["model"]["features"]]
             prediction = predict_onnx(model_onnx, input_df)
             df.loc[:, "prediction"] = prediction.squeeze()
@@ -74,6 +71,8 @@ def main():
                     written = True
                 else:
                     df.to_csv(args.output, index=False, mode="a", header=False)
-            
+
+
 if __name__ == "__main__":
-    main()
+    args = args_parser().parse_args()
+    main(args)

marcai/predict.py

@@ -44,10 +44,7 @@ def args_parser():
     return parser
 
 
-def main():
-
-    args = args_parser().parse_args()
-
+def main(args):
     config_path = f"{args.model_dir}/config.yaml"
     model_onnx = f"{args.model_dir}/model.onnx"
 
@@ -75,4 +72,5 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    args = args_parser().parse_args()
+    main(args)

marcai/process.py

@@ -47,7 +47,7 @@ def multiprocess_pairs(
 
                 for future in done:
                     # Get job's output
-                    df  = future.result()
+                    df = future.result()
 
                     # Yield output
                     yield df
@@ -58,7 +58,7 @@ def multiprocess_pairs(
 
                 if pairs_chunk is None:
                     break
-            
+
                 indices = np.array(pairs_chunk).astype(int)
 
                 left_indices = indices[:, 0]
@@ -127,11 +127,7 @@ def process(df0, df1):
     result_df["author"] = comps.maximum(authors, null_value=0.5)
 
     # Weighted title comparison
-    weights = {
-        "title_a": 1,
-        "raw": 0,
-        "title_p": 1
-    }
+    weights = {"title_a": 1, "raw": 0, "title_p": 1}
 
     result_df["title_agg"] = comps.column_aggregate_similarity(
         df0[weights.keys()], df1[weights.keys()], weights.values(), null_value=0
@@ -142,8 +138,6 @@ def process(df0, df1):
         df0["title"], df1["title"], null_value=0.5
     )
 
-
-
     # Token set similarity
     result_df["title_tokenset"] = comps.token_set_similarity(
         df0["title"], df1["title"], null_value=0
@@ -220,10 +214,8 @@ def args_parser():
     return parser
 
 
-def main():
-    
+def main(args):
     start = time.time()
-    args = args_parser().parse_args()
 
     # Load records
     print("Loading records...")
@@ -258,4 +250,5 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    args = args_parser().parse_args()
+    main(args)

marcai/train.py

@@ -93,12 +93,11 @@ def args_parser():
     parser.add_argument("-n", "--run-name", help="Name for training run", required=True)
     return parser
 
-def main():
-    
-    args = args_parser().parse_args()
 
+def main(args):
     train(args.run_name)
 
 
 if __name__ == "__main__":
-    main()
+    args = args_parser().parse_args()
+    main(args)