Update train.py

train.py

@@ -7,10 +7,18 @@ from nemo.utils import exp_manager, logging
 from omegaconf import OmegaConf, open_dict
 
 
-def train_model(train_manifest: str, val_manifest: str, accelerator: str, batch_size: int, num_epochs: int, model_save_path: str = None,) -> None:
+def train_model(
+        train_manifest: str = None, 
+        val_manifest: str = None, 
+        accelerator: str = "cpu", 
+        batch_size: int = 1, 
+        num_epochs: int = 1, 
+        model_save_path: str = None,
+) -> None:
 
     # Loading a STT Quartznet 15x5 model
     model = nemo_asr.models.ASRModel.from_pretrained("stt_en_quartznet15x5")
+    
     # New vocabulary for a model
     new_vocabulary = [
         " ",
@@ -56,7 +64,8 @@ def train_model(train_manifest: str, val_manifest: str, accelerator: str, batch_
         "ә",
         "ө",
     ]
-
+    
+    # Configurations
     with open_dict(model.cfg):
         # Setting up the labels and sample rate
         model.cfg.labels = new_vocabulary
@@ -104,7 +113,7 @@ def train_model(train_manifest: str, val_manifest: str, accelerator: str, batch_
 
     # Trainer
     trainer = ptl.Trainer(
-        accelerator="gpu",
+        accelerator=accelerator,
         max_epochs=num_epochs,
         accumulate_grad_batches=1,
         enable_checkpointing=False,
@@ -133,10 +142,10 @@ def train_model(train_manifest: str, val_manifest: str, accelerator: str, batch_
     print(OmegaConf.to_yaml(model.cfg))
     print("-----------------------------------------------------------")
 
-    #     # Fitting the model
+    # Fitting the model
     trainer.fit(model)
 
-    #     # Saving the model
+    # Saving the model
     if model_save_path:
         model.save_to(f"{model_save_path}")
         print(f"Model saved at path : {os.getcwd() + os.path.sep + model_save_path}")
@@ -145,9 +154,9 @@ def train_model(train_manifest: str, val_manifest: str, accelerator: str, batch_
 if __name__ == "__main__":
     # Parse command line arguments
     parser = argparse.ArgumentParser()
-    parser.add_argument("--train_manifest", help="Path for train manifest JSON file.")
-    parser.add_argument("--val_manifest", help="Path for validation manifest JSON file.")
-    parser.add_argument("--accelerator", help="What accelerator type to use (cpu, gpu, tpu, etc.).")
+    parser.add_argument("--train_manifest", default = None, help="Path for train manifest JSON file.")
+    parser.add_argument("--val_manifest", default = None, help="Path for validation manifest JSON file.")
+    parser.add_argument("--accelerator", default="cpu", help="What accelerator type to use (cpu, gpu, tpu, etc.).")
     parser.add_argument("--batch_size", type=int, default=1, help="Batch size of the dataset to train.")
     parser.add_argument("--num_epochs", type=int, default=1, help="Number of epochs to train for.")
     parser.add_argument("--model_save_path", default=None, help="Path for saving a trained model.")
@@ -160,4 +169,4 @@ if __name__ == "__main__":
         batch_size=args.batch_size,
         num_epochs=args.num_epochs,
         model_save_path=args.model_save_path,
-    )
+    )