@@ -194,7 +194,7 @@ def get_parameters(
 ):
     """Get the parameters from the configuration file."""
     # Check file extension
-    if params_file.endswith(".yml"):
+    if params_file.endswith(".yml") or params_file.endswith(".yaml"):
         loader = yaml.safe_load
     elif params_file.endswith(".json"):
         loader = json.load
@@ -509,7 +509,7 @@ def reorder_columns(df: pd.DataFrame, desired_order: list) -> pd.DataFrame:
     )
 
 
-def save_file(file_path, data):
+def save_file(file_path: str, data):
     """
     Save data to a file based on the file extension.
 
@@ -520,7 +520,7 @@ def save_file(file_path, data):
     Raises:
         ValueError: If the file format is not supported.
     """
-    if file_path.endswith(".yml"):
+    if file_path.endswith(".yml") or file_path.endswith(".yaml"):
         dumper = yaml.safe_dump
     elif file_path.endswith(".json"):
         dumper = json.dump
 
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langtest"
-version = "2.0.0"
+version = "2.1.0"
 description = "John Snow Labs provides a library for delivering safe & effective NLP models."
 authors = ["John Snow Labs <support@johnsnowlabs.com>"]
 readme = "README.md"