Update train.py

train.py

@@ -993,3 +993,21 @@ def train_and_evaluate(
                 pid_data.pop("process_pids", None)
                 json.dump(pid_data, pid_file, indent=4)
             os._exit(0)
+
+            
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument("--model_type", type=str, default="rvc", choices=["rvc", "djcm"])
+parser.add_argument("--train_data", type=str, required=True)
+parser.add_argument("--output_path", type=str, required=True)
+args = parser.parse_args()
+
+if args.model_type == "djcm":
+    from models.djcm.djcm_model import DJCMModel
+    model = DJCMModel()
+else:
+    from models.rvc.rvc_model import RVCModel
+    model = RVCModel()
+
+print(f"Training {args.model_type} model with data at {args.train_data}")
+