Repository: Galileo-Galilei/kedro-mlflow
Pull request #439 (2023-08-01 15:06:46)

Changing DataSet to Dataset#391
## Description
Changing DataSet to Dataset




Galileo-Galilei: Hi, you should likely rebase on  because you have a lot of duplicated commits.
Galileo-Galilei: Hi @ShubhamZoro, thank you very much for opening this PR, and once again sorry for not answering fast and give good support to your PR. I'll open another PR to fix the git history, unless you want to give it a try. I'll keep your commits though so you will be credited for the works. Thank you very much for contributing!

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 88f9ded..2eae232 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,11 +2,20 @@
 
 ## [Unreleased]
 
+- :sparkles: Add support for python 3.11 ([#450,rxm7706](https://github.com/Galileo-Galilei/kedro-mlflow/pull/450))
+- :boom: :recycle: Rename the following ``DataSets`` with the ``Dataset`` suffix (without capitalized ``S``) to match new kedro conventions from ``kedro>=0.19`` and onwards ([#439, ShubhamZoro](https://github.com/Galileo-Galilei/kedro-mlflow/pull/439)):
+  - ``MlflowArtifactDataSet``->``MlflowArtifactDataset``
+  - ``MlflowAbstractModelDataSet``->``MlflowAbstractModelDataset``
+  - ``MlflowModelRegistryDataSet``->``MlflowModelRegistryDataset``
+  - ``MlflowMetricDataSet``->``MlflowMetricDataset``
+  - ``MlflowMetricHistoryDataSet``->``MlflowMetricHistoryDataset``
+  - ``MlflowMetricsDataSet``->``MlflowMetricsDataset``
+
 ## [0.11.10] - 2023-10-03
 
 ### Fixed
 
--   :bug: Avoid error when using kedro==0.18.1 with `TemplatedConfigLoader` and no `mlflow.yml` configuration file ([#452](https://github.com/Galileo-Galilei/kedro-mlflow/issues/452))
+-   :bug: Avoid error when using kedro==0.18.1 with `TemplatedConfigLoader` and no `mlflow.yml` configuration file ([#452, sami-sweng](https://github.com/Galileo-Galilei/kedro-mlflow/issues/452))
 
 ## [0.11.9] - 2023-07-23
 
diff --git a/README.md b/README.md
index 6f961b2..2604251 100644
--- a/README.md
+++ b/README.md
@@ -54,8 +54,8 @@ The documentation contains:
 
 Some frequently asked questions on more advanced features:
 
-- You want to log additional metrics to the run? -> [Try ``MlflowMetricsDataSet``](https://kedro-mlflow.readthedocs.io/en/stable/source/04_experimentation_tracking/05_version_metrics.html) !
-- You want to log nice dataviz of your pipeline that you register with ``MatplotlibWriter``? -> [Try ``MlflowArtifactDataSet`` to log any local files (.png, .pkl, .csv...) *automagically*](https://kedro-mlflow.readthedocs.io/en/stable/source/04_experimentation_tracking/03_version_datasets.html)!
+- You want to log additional metrics to the run? -> [Try ``MlflowMetricsDataset``](https://kedro-mlflow.readthedocs.io/en/stable/source/04_experimentation_tracking/05_version_metrics.html) !
+- You want to log nice dataviz of your pipeline that you register with ``MatplotlibWriter``? -> [Try ``MlflowArtifactDataset`` to log any local files (.png, .pkl, .csv...) *automagically*](https://kedro-mlflow.readthedocs.io/en/stable/source/04_experimentation_tracking/03_version_datasets.html)!
 - You want to create easily an API to share your awesome model to anyone? -> [See if ``pipeline_ml_factory`` can fit your needs](https://github.com/Galileo-Galilei/kedro-mlflow/issues/16)
 - You want to do something that is not straigthforward with current implementation? Open an issue, and let's see what happens!
 
diff --git a/docs/source/01_introduction/02_motivation.md b/docs/source/01_introduction/02_motivation.md
index e1c50e1..8e03592 100644
--- a/docs/source/01_introduction/02_motivation.md
+++ b/docs/source/01_introduction/02_motivation.md
@@ -39,9 +39,9 @@ Above implementations have the advantage of being very straightforward and *mlfl
 | :------------------------ | :-------------- | :------------------------------------------------------- |
 | Set up configuration      | ``mlflow.yml``  | ``MlflowHook``                                           |
 | Logging parameters        | ``mlflow.yml``  | ``MlflowHook``                                           |
-| Logging artifacts         | ``catalog.yml`` | ``MlflowArtifactDataSet``                                |
+| Logging artifacts         | ``catalog.yml`` | ``MlflowArtifactDataset``                                |
 | Logging models            | ``catalog.yml`` | `MlflowModelLoggerDataSet` and `MlflowModelSaverDataSet` |
-| Logging metrics           | ``catalog.yml`` | ``MlflowMetricsDataSet``                                 |
+| Logging metrics           | ``catalog.yml`` | ``MlflowMetricsDataset``                                 |
 | Logging Pipeline as model | ``hooks.py``    | ``KedroPipelineModel`` and ``pipeline_ml_factory``       |
 
 `kedro-mlflow` does not currently provide interface to set tags outside a Kedro ``Pipeline``. Some of above decisions are subject to debate and design decisions (for instance, metrics are often updated in a loop during each epoch / training iteration and it does not always make sense to register the metric between computation steps, e.g. as a an I/O operation after a node run).
diff --git a/docs/source/02_installation/03_migration_guide.md b/docs/source/02_installation/03_migration_guide.md
index 00afdb8..991a23c 100644
--- a/docs/source/02_installation/03_migration_guide.md
+++ b/docs/source/02_installation/03_migration_guide.md
@@ -100,8 +100,8 @@ Replace the following entries:
 
 | old                                     | new                                               |
 | :-------------------------------------- | :------------------------------------------------ |
-| `kedro_mlflow.io.MlflowArtifactDataSet` | `kedro_mlflow.io.artifacts.MlflowArtifactDataSet` |
-| `kedro_mlflow.io.MlflowMetricsDataSet`  | `kedro_mlflow.io.metrics.MlflowMetricsDataSet`    |
+| `kedro_mlflow.io.MlflowArtifactDataset` | `kedro_mlflow.io.artifacts.MlflowArtifactDataset` |
+| `kedro_mlflow.io.MlflowMetricsDataset`  | `kedro_mlflow.io.metrics.MlflowMetricsDataset`    |
 
 ### Hooks
 
diff --git a/docs/source/03_getting_started/02_first_steps.md b/docs/source/03_getting_started/02_first_steps.md
index 5d3e686..11465a9 100644
--- a/docs/source/03_getting_started/02_first_steps.md
+++ b/docs/source/03_getting_started/02_first_steps.md
@@ -129,7 +129,7 @@ example_iris_data:
 
 ```
 
-And persist the model as a pickle with the ``MlflowArtifactDataSet`` class:
+And persist the model as a pickle with the ``MlflowArtifactDataset`` class:
 
 ```yaml
 # This is a data set used by the "Hello World" example pipeline provided with the project
@@ -140,7 +140,7 @@ example_iris_data:
   filepath: data/01_raw/iris.csv
 
 example_model:
-  type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+  type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
   data_set:
     type: pickle.PickleDataSet
     filepath: data/06_models/trained_model.pkl
diff --git a/docs/source/04_experimentation_tracking/03_version_datasets.md b/docs/source/04_experimentation_tracking/03_version_datasets.md
index 4275471..aa5fcf2 100644
--- a/docs/source/04_experimentation_tracking/03_version_datasets.md
+++ b/docs/source/04_experimentation_tracking/03_version_datasets.md
@@ -14,7 +14,7 @@ Artifacts are a very flexible and convenient way to "bind" any data type to your
 
 ## How to version data in a kedro project?
 
-``kedro-mlflow`` introduces a new ``AbstractDataSet`` called ``MlflowArtifactDataSet``. It is a wrapper for any ``AbstractDataSet`` which decorates the underlying dataset ``save`` method and logs the file automatically in mlflow as an artifact each time the ``save`` method is called.
+``kedro-mlflow`` introduces a new ``AbstractDataSet`` called ``MlflowArtifactDataset``. It is a wrapper for any ``AbstractDataSet`` which decorates the underlying dataset ``save`` method and logs the file automatically in mlflow as an artifact each time the ``save`` method is called.
 
 Since it is an ``AbstractDataSet``, it can be used with the YAML API. Assume that you have the following entry in the ``catalog.yml``:
 
@@ -28,7 +28,7 @@ You can change it to:
 
 ```yaml
 my_dataset_to_version:
-    type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
     data_set:
         type: pandas.CSVDataSet  # or any valid kedro DataSet
         filepath: /path/to/a/LOCAL/destination/file.csv # must be a local file, wherever you want to log the data in the end
@@ -38,13 +38,13 @@ and this dataset will be automatically versioned in each pipeline execution.
 
 ## Frequently asked questions
 
-### Can I pass extra parameters to the ``MlflowArtifactDataSet`` for finer control?
+### Can I pass extra parameters to the ``MlflowArtifactDataset`` for finer control?
 
-The ``MlflowArtifactDataSet`` takes a ``data_set`` argument which is a python dictionary passed to the ``__init__`` method of the dataset declared in ``type``. It means that you can pass any argument accepted by the underlying dataset in this dictionary. If you want to pass ``load_args`` and ``save_args`` in the previous example, add them in the ``data_set`` argument:
+The ``MlflowArtifactDataset`` takes a ``data_set`` argument which is a python dictionary passed to the ``__init__`` method of the dataset declared in ``type``. It means that you can pass any argument accepted by the underlying dataset in this dictionary. If you want to pass ``load_args`` and ``save_args`` in the previous example, add them in the ``data_set`` argument:
 
 ```yaml
 my_dataset_to_version:
-    type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
     data_set:
         type: pandas.CSVDataSet  # or any valid kedro DataSet
         filepath: /path/to/a/local/destination/file.csv
@@ -55,12 +55,12 @@ my_dataset_to_version:
         # ... any other valid arguments for data_set
 ```
 
-### Can I use the ``MlflowArtifactDataSet`` in interactive mode?
+### Can I use the ``MlflowArtifactDataset`` in interactive mode?
 
-Like all Kedro ``AbstractDataSet``, ``MlflowArtifactDataSet`` is callable in the python API:
+Like all Kedro ``AbstractDataSet``, ``MlflowArtifactDataset`` is callable in the python API:
 
 ```python
-from kedro_mlflow.io.artifacts import MlflowArtifactDataSet
+from kedro_mlflow.io.artifacts import MlflowArtifactDataset
 from kedro.extras.datasets.pandas import CSVDataSet
 
 csv_dataset = MlflowArtifactDataSet(
@@ -85,11 +85,11 @@ You can refer to [this issue](https://github.com/Galileo-Galilei/kedro-mlflow/is
 
 ### Can I log an artifact in a specific run?
 
-The ``MlflowArtifactDataSet`` has an extra attribute ``run_id`` which specifies the run you will log the artifact in. **Be cautious, because this argument will take precedence over the current run** when you call ``kedro run``, causing the artifact to be logged in another run that all the other data of the run.
+The ``MlflowArtifactDataset`` has an extra attribute ``run_id`` which specifies the run you will log the artifact in. **Be cautious, because this argument will take precedence over the current run** when you call ``kedro run``, causing the artifact to be logged in another run that all the other data of the run.
 
 ```yaml
 my_dataset_to_version:
-    type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
     data_set:
         type: pandas.CSVDataSet  # or any valid kedro DataSet
         filepath: /path/to/a/local/destination/file.csv  # must be a local filepath, no matter what is your actual mlflow storage (S3 or other)
@@ -98,13 +98,13 @@ my_dataset_to_version:
 
 ### Can I reload an artifact from an existing run to use it in another run ?
 
-You may want to reuse th artifact of a previous run to reuse it in another one, e.g. to continue training from a pretrained model, or to select the best model among several runs created during an hyperparamter tuning. The ``MlflowArtifactDataSet`` has an extra attribute ``run_id`` you can use to specify from which run you will load the artifact from. **Be cautious**, because:
+You may want to reuse th artifact of a previous run to reuse it in another one, e.g. to continue training from a pretrained model, or to select the best model among several runs created during an hyperparamter tuning. The ``MlflowArtifactDataset`` has an extra attribute ``run_id`` you can use to specify from which run you will load the artifact from. **Be cautious**, because:
 - this argument will take precedence over the current run** when you call ``kedro run``, causing the artifact to be loaded from another run that all the other data of the run
 - the artifact will be downloaded and erase the existing file at your local filepath
 
 ```yaml
 my_dataset_to_reload:
-    type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
     data_set:
         type: pandas.CSVDataSet  # or any valid kedro DataSet
         filepath: /path/to/a/local/destination/file.csv # must be a local filepath, no matter what is your actual mlflow storage (S3 or other)
@@ -113,13 +113,13 @@ my_dataset_to_reload:
 
 ### Can I create a remote folder/subfolders architecture to organize the artifacts?
 
-The ``MlflowArtifactDataSet`` has an extra argument ``artifact_path`` which specifies a remote subfolder where the artifact will be logged. It must be a relative path.
+The ``MlflowArtifactDataset`` has an extra argument ``artifact_path`` which specifies a remote subfolder where the artifact will be logged. It must be a relative path.
 
 With below example, the artifact will be logged in mlflow within a `reporting` folder.
 
 ```yaml
 my_dataset_to_version:
-    type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
     data_set:
         type: pandas.CSVDataSet  # or any valid kedro DataSet
         filepath: /path/to/a/local/destination/file.csv
diff --git a/docs/source/04_experimentation_tracking/04_version_models.md b/docs/source/04_experimentation_tracking/04_version_models.md
index 1c8bb6b..80c506f 100644
--- a/docs/source/04_experimentation_tracking/04_version_models.md
+++ b/docs/source/04_experimentation_tracking/04_version_models.md
@@ -11,7 +11,7 @@ MLflow allows to serialize and deserialize models to a common format, track thos
 - the ``MlflowModelLoggerDataSet`` is used to load from and save to from the mlflow artifact store. It uses optional `run_id` argument to load and save from a given `run_id` which must exists in the mlflow server you are logging to.
 - the ``MlflowModelSaverDataSet`` is used to load from and save to a given path. It uses the standard `filepath` argument in the constructor of Kedro DataSets. Note that it **does not log in mlflow**.
 
-*Note: If you use ``MlflowModelLoggerDataSet``, it will be saved during training in your current run. However, you will need to specify the run id to predict with (since it is not persisted locally, it will not pick the latest model by default). You may prefer to combine ``MlflowModelSaverDataSet`` and ``MlflowArtifactDataSet`` to make persist it both locally and remotely, see further.*
+*Note: If you use ``MlflowModelLoggerDataSet``, it will be saved during training in your current run. However, you will need to specify the run id to predict with (since it is not persisted locally, it will not pick the latest model by default). You may prefer to combine ``MlflowModelSaverDataSet`` and ``MlflowArtifactDataset`` to make persist it both locally and remotely, see further.*
 
 Suppose you would like to register a `scikit-learn` model of your `DataCatalog` in mlflow, you can use the following yaml API:
 
@@ -54,11 +54,11 @@ my_custom_model:
 
 ### How can I save model locally and log it in MLflow in one step?
 
-If you want to save your model both locally and remotely within the same run, you can leverage `MlflowArtifactDataSet`:
+If you want to save your model both locally and remotely within the same run, you can leverage `MlflowArtifactDataset`:
 
 ```yaml
 sklearn_model:
-    type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
     data_set:
         type: kedro_mlflow.io.models.MlflowModelSaverDataSet
         flavor: mlflow.sklearn
diff --git a/docs/source/04_experimentation_tracking/05_version_metrics.md b/docs/source/04_experimentation_tracking/05_version_metrics.md
index 30af424..8310448 100644
--- a/docs/source/04_experimentation_tracking/05_version_metrics.md
+++ b/docs/source/04_experimentation_tracking/05_version_metrics.md
@@ -7,20 +7,20 @@ MLflow defines a metric as "a (key, value) pair, where the value is numeric". Ea
 ## How to version metrics in a kedro project?
 
 `kedro-mlflow` introduces 3 ``AbstractDataSet`` to manage metrics:
-- ``MlflowMetricDataSet`` which can log a float as a metric
-- ``MlflowMetricHistoryDataSet`` which can log the evolution over time of a given metric, e.g. a list or a dict of float.
-- ``MlflowMetricsDataSet``. It is a wrapper around a dictionary with metrics which is returned by node and log metrics in MLflow.
+- ``MlflowMetricDataset`` which can log a float as a metric
+- ``MlflowMetricHistoryDataset`` which can log the evolution over time of a given metric, e.g. a list or a dict of float.
+- ``MlflowMetricsDataset``. It is a wrapper around a dictionary with metrics which is returned by node and log metrics in MLflow.
 
-### Saving a single float as a metric with ``MlflowMetricDataSet``
+### Saving a single float as a metric with ``MlflowMetricDataset``
 
-The ``MlflowMetricDataSet`` is an ``AbstractDataSet`` which enable to save or load a ``float`` as a mlflow metric. You must specify the ``key`` (i.e. the name to display in mlflow) when creating the dataset. Somes examples follow:
+The ``MlflowMetricDataset`` is an ``AbstractDataSet`` which enable to save or load a ``float`` as a mlflow metric. You must specify the ``key`` (i.e. the name to display in mlflow) when creating the dataset. Somes examples follow:
 
 - The most basic usage is to create the dataset and save a a value:
 
 ```python
-from kedro_mlflow.io.metrics import MlflowMetricDataSet
+from kedro_mlflow.io.metrics import MlflowMetricDataset
 
-metric_ds = MlflowMetricDataSet(key="my_metric")
+metric_ds = MlflowMetricDataset(key="my_metric")
 with mlflow.start_run():
     metric_ds.save(
         0.3
@@ -34,9 +34,9 @@ Unlike mlflow default behaviour, if there is no active run, no run is created.
 - You can also specify a ``run_id`` instead of logging in the active run:
 
 ```python
-from kedro_mlflow.io.metrics import MlflowMetricDataSet
+from kedro_mlflow.io.metrics import MlflowMetricDataset
 
-metric_ds = MlflowMetricDataSet(key="my_metric", run_id="123456789")
+metric_ds = MlflowMetricDataset(key="my_metric", run_id="123456789")
 with mlflow.start_run():
     metric_ds.save(
         0.3
@@ -46,9 +46,9 @@ with mlflow.start_run():
 It is also possible to pass ``load_args`` and ``save_args`` to control which step should be logged (in case you have logged several step for the same metric.) ``save_args`` accepts a ``mode`` key which can be set to ``overwrite`` (mlflow default) or ``append``. In append mode, if no step is specified, saving the metric will "bump" the last existing step to create a linear history. **This is very useful if you have a monitoring pipeline which calculates a metric frequently to check the performance of a deployed model.**
 
 ```python
-from kedro_mlflow.io.metrics import MlflowMetricDataSet
+from kedro_mlflow.io.metrics import MlflowMetricDataset
 
-metric_ds = MlflowMetricDataSet(
+metric_ds = MlflowMetricDataset(
     key="my_metric", load_args={"step": 1}, save_args={"mode": "append"}
 )
 
@@ -64,7 +64,7 @@ Since it is an ``AbstractDataSet``, it can be used with the YAML API in your ``c
 
 ```yaml
 my_model_metric:
-    type: kedro_mlflow.io.metrics.MlflowMetricDataSet
+    type: kedro_mlflow.io.metrics.MlflowMetricDataset
     run_id: 123456 # OPTIONAL, you should likely let it empty to log in the current run
     key: my_awesome_name # OPTIONAL: if not provided, the dataset name will be sued (here "my_model_metric")
     load_args:
@@ -74,17 +74,17 @@ my_model_metric:
         mode: append #  OPTIONAL: likely better than the default "overwrite". Will be ignored if "step" is provided.
 ```
 
-### Saving the evolution of a metric during training with ``MlflowMetricHistoryDataSet``
+### Saving the evolution of a metric during training with ``MlflowMetricHistoryDataset``
 
-The ``MlflowMetricDataSet`` is an ``AbstractDataSet`` which enable to save or load the evolutionf of a metric with various formats. You must specify the ``key`` (i.e. the name to display in mlflow) when creating the dataset. Somes examples follow:
+The ``MlflowMetricDataset`` is an ``AbstractDataSet`` which enable to save or load the evolutionf of a metric with various formats. You must specify the ``key`` (i.e. the name to display in mlflow) when creating the dataset. Somes examples follow:
 
 It enables logging either:
   - a list of int as a metric with incremental step, e.g ``[0.1,0.2,0.3]`` with ``mode=list`` for either ``save_args`` or ``load_args``
 
 ```python
-from kedro_mlflow.io.metrics import MlflowMetricHistoryDataSet
+from kedro_mlflow.io.metrics import MlflowMetricHistoryDataset
 
-metric_history_ds = MlflowMetricDataSet(key="my_metric", save_args={"mode": "list"})
+metric_history_ds = MlflowMetricDataset(key="my_metric", save_args={"mode": "list"})
 
 with mlflow.start_run():
     metric_history_ds.save([0.1, 0.2, 0.3])  # will be logged with incremental steps
@@ -92,9 +92,9 @@ with mlflow.start_run():
   - a dict of {step: value} as a metric:
 
 ```python
-from kedro_mlflow.io.metrics import MlflowMetricHistoryDataSet
+from kedro_mlflow.io.metrics import MlflowMetricHistoryDataset
 
-metric_history_ds = MlflowMetricDataSet(key="my_metric", save_args={"mode": "dict"})
+metric_history_ds = MlflowMetricDataset(key="my_metric", save_args={"mode": "dict"})
 
 with mlflow.start_run():
     metric_history_ds.save(
@@ -105,9 +105,9 @@ with mlflow.start_run():
   - a list of dict [{log_metric_arg: value}] as a metric, e.g:
 
 ```python
-from kedro_mlflow.io.metrics import MlflowMetricHistoryDataSet
+from kedro_mlflow.io.metrics import MlflowMetricHistoryDataset
 
-metric_history_ds = MlflowMetricDataSet(key="my_metric", save_args={"mode": "history"})
+metric_history_ds = MlflowMetricDataset(key="my_metric", save_args={"mode": "history"})
 
 with mlflow.start_run():
     metric_history_ds.save(
@@ -122,9 +122,9 @@ with mlflow.start_run():
 You can combine the different mode for save and load, e.g:
 
 ```python
-from kedro_mlflow.io.metrics import MlflowMetricHistoryDataSet
+from kedro_mlflow.io.metrics import MlflowMetricHistoryDataset
 
-metric_history_ds = MlflowMetricDataSet(
+metric_history_ds = MlflowMetricDataset(
     key="my_metric", save_args={"mode": "dict"}, save_args={"mode": "list"}
 )
 
@@ -139,7 +139,7 @@ As usual, since it is an ``AbstractDataSet``, it can be used with the YAML API i
 
 ```yaml
 my_model_metric:
-    type: kedro_mlflow.io.metrics.MlflowMetricHistoryDataSet
+    type: kedro_mlflow.io.metrics.MlflowMetricHistoryDataset
     run_id: 123456 # OPTIONAL, you should likely let it empty to log in the current run
     key: my_awesome_name # OPTIONAL: if not provided, the dataset name will be used (here "my_model_metric")
     load_args:
@@ -148,13 +148,13 @@ my_model_metric:
         mode: ... # OPTIONAL: "list" by default, one of {"list", "dict", "history"}
 ```
 
-### Saving several metrics with their entire history with ``MlflowMetricsDataSet``
+### Saving several metrics with their entire history with ``MlflowMetricsDataset``
 
 Since it is an ``AbstractDataSet``, it can be used with the YAML API. You can define it in your ``catalog.yml`` as:
 
 ```yaml
 my_model_metrics:
-    type: kedro_mlflow.io.metrics.MlflowMetricsDataSet
+    type: kedro_mlflow.io.metrics.MlflowMetricsDataset
 ```
 
 You can provide a prefix key, which is useful in situations like when you have multiple nodes producing metrics with the same names which you want to distinguish. If you are using the ``v``, it will handle that automatically for you by giving as prefix metrics data set name. In the example above the prefix would be ``my_model_metrics``.
@@ -163,7 +163,7 @@ Let's look at an example with custom prefix:
 
 ```yaml
 my_model_metrics:
-    type: kedro_mlflow.io.metrics.MlflowMetricsDataSet
+    type: kedro_mlflow.io.metrics.MlflowMetricsDataset
     prefix: foo
 ```
 
@@ -179,7 +179,7 @@ def metrics_node() -> Dict[str, Union[float, List[float]]]:
     }
 ```
 
-As you can see above, ``kedro_mlflow.io.metrics.MlflowMetricsDataSet`` can take metrics as:
+As you can see above, ``kedro_mlflow.io.metrics.MlflowMetricsDataset`` can take metrics as:
 
 - ``Dict[str, key]``
 - ``List[Dict[str, key]]``
@@ -188,7 +188,7 @@ To store metrics we need to define metrics dataset in Kedro Catalog:
 
 ```yaml
 my_model_metrics:
-    type: kedro_mlflow.io.metrics.MlflowMetricsDataSet
+    type: kedro_mlflow.io.metrics.MlflowMetricsDataset
 ```
 
 Within a kedro run, the ``MlflowHook`` will automatically prefix the metrics datasets with their name in the catalog. In our example, the metrics will be stored in Mlflow with the following keys: ``my_model_metrics.metric1``, ``my_model_metrics.metric2``.
@@ -197,7 +197,7 @@ It is also prossible to provide a prefix manually:
 
 ```yaml
 my_model_metrics:
-    type: kedro_mlflow.io.metrics.MlflowMetricsDataSet
+    type: kedro_mlflow.io.metrics.MlflowMetricsDataset
     prefix: foo
 ```
 
diff --git a/docs/source/06_interactive_use/01_notebook_use.md b/docs/source/06_interactive_use/01_notebook_use.md
index af99364..2603908 100644
--- a/docs/source/06_interactive_use/01_notebook_use.md
+++ b/docs/source/06_interactive_use/01_notebook_use.md
@@ -39,7 +39,7 @@ If you change your ``mlflow.yml``, reload the kedro extension for the changes to
 
 ## Difference with running through the CLI
 
-- The DataSets `load` and `save` methods works as usual. You can call `catalog.save("my_artifact_dataset", data)` inside a cell, and your data will be logged in mlflow properly (assuming "my_artifact_dataset" is a `kedro_mlflow.io.MlflowArtifactDataSet`).
+- The DataSets `load` and `save` methods works as usual. You can call `catalog.save("my_artifact_dataset", data)` inside a cell, and your data will be logged in mlflow properly (assuming "my_artifact_dataset" is a `kedro_mlflow.io.MlflowArtifactDataset`).
 - The `hooks` which automatically save all parameters/metrics/artifacts in mlflow will work if you run the session interactively, e.g.:
 
 ```python
diff --git a/docs/source/07_python_objects/01_DataSets.md b/docs/source/07_python_objects/01_DataSets.md
index 4a28202..cf1c4a8 100644
--- a/docs/source/07_python_objects/01_DataSets.md
+++ b/docs/source/07_python_objects/01_DataSets.md
@@ -1,12 +1,12 @@
 # New ``DataSet``
 
-## ``MlflowArtifactDataSet``
+## ``MlflowArtifactDataset``
 
-``MlflowArtifactDataSet`` is a wrapper for any ``AbstractDataSet`` which logs the dataset automatically in mlflow as an artifact when its ``save`` method is called. It can be used both with the YAML API:
+``MlflowArtifactDataset`` is a wrapper for any ``AbstractDataSet`` which logs the dataset automatically in mlflow as an artifact when its ``save`` method is called. It can be used both with the YAML API:
 
 ```yaml
 my_dataset_to_version:
-    type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
     data_set:
         type: pandas.CSVDataSet  # or any valid kedro DataSet
         filepath: /path/to/a/local/destination/file.csv
@@ -16,7 +16,7 @@ or with additional parameters:
 
 ```yaml
 my_dataset_to_version:
-    type: kedro_mlflow.io.artifacts.MlflowArtifactDataSet
+    type: kedro_mlflow.io.artifacts.MlflowArtifactDataset
     data_set:
         type: pandas.CSVDataSet  # or any valid kedro DataSet
         filepath: /path/to/a/local/destination/file.csv
@@ -32,10 +32,10 @@ my_dataset_to_version:
 or with the python API:
 
 ```python
-from kedro_mlflow.io.artifacts import MlflowArtifactDataSet
+from kedro_mlflow.io.artifacts import MlflowArtifactDataset
 from kedro.extras.datasets.pandas import CSVDataSet
 
-csv_dataset = MlflowArtifactDataSet(
+csv_dataset = MlflowArtifactDataset(
     data_set={"type": CSVDataSet, "filepath": r"/path/to/a/local/destination/file.csv"}
 )
 csv_dataset.save(data=pd.DataFrame({"a": [1, 2], "b": [3, 4]}))
@@ -43,13 +43,13 @@ csv_dataset.save(data=pd.DataFrame({"a": [1, 2], "b": [3, 4]}))
 
 ## Metrics `DataSets`
 
-### ``MlflowMetricDataSet``
+### ``MlflowMetricDataset``
 
-[The ``MlflowMetricDataSet`` is documented here](https://kedro-mlflow.readthedocs.io/en/latest/source/04_experimentation_tracking/05_version_metrics.html#saving-a-single-float-as-a-metric-with-mlflowmetricdataset).
+[The ``MlflowMetricDataset`` is documented here](https://kedro-mlflow.readthedocs.io/en/latest/source/04_experimentation_tracking/05_version_metrics.html#saving-a-single-float-as-a-metric-with-mlflowmetricdataset).
 
-### ``MlflowMetricHistoryDataSet``
+### ``MlflowMetricHistoryDataset``
 
-[The ``MlflowMetricHistoryDataSet`` is documented here](https://kedro-mlflow.readthedocs.io/en/latest/source/04_experimentation_tracking/05_version_metrics.html#saving-a-single-float-as-a-metric-with-mlflowmetricdataset).
+[The ``MlflowMetricHistoryDataset`` is documented here](https://kedro-mlflow.readthedocs.io/en/latest/source/04_experimentation_tracking/05_version_metrics.html#saving-a-single-float-as-a-metric-with-mlflowmetricdataset).
 
 
 ## Models `DataSets`
@@ -159,9 +159,9 @@ my_model:
     version: <valid-kedro-version>
 ```
 
-### ``MlflowModelRegistryDataSet``
+### ``MlflowModelRegistryDataset``
 
-The ``MlflowModelRegistryDataSet`` accepts the following arguments:
+The ``MlflowModelRegistryDataset`` accepts the following arguments:
 
 - model_name (str): The name of the registered model is the mlflow registry
 - stage_or_version (str): A valid stage (either "staging" or "production") or version number for the registred model.Default to "latest" which fetch the last version and the higher "stage" available.
@@ -189,9 +189,9 @@ with mlflow.start_run():
 You can fetch the model by its name:
 
 ```python
-from kedro_mlflow.io.models import MlflowModelRegistryDataSet
+from kedro_mlflow.io.models import MlflowModelRegistryDataset
 
-mlflow_model_logger = MlflowModelRegistryDataSet(model_name="my_awesome_model")
+mlflow_model_logger = MlflowModelRegistryDataset(model_name="my_awesome_model")
 my_model = mlflow_model_logger.load()
 ```
 
@@ -199,6 +199,6 @@ and with the YAML API in the `catalog.yml` (only for loading an existing model):
 
 ```yaml
 my_model:
-    type: kedro_mlflow.io.models.MlflowModelRegistryDataSet
+    type: kedro_mlflow.io.models.MlflowModelRegistryDataset
     model_name: my_awesome_model
 ```
diff --git a/kedro_mlflow/framework/hooks/mlflow_hook.py b/kedro_mlflow/framework/hooks/mlflow_hook.py
index 0d88096..18497d7 100644
--- a/kedro_mlflow/framework/hooks/mlflow_hook.py
+++ b/kedro_mlflow/framework/hooks/mlflow_hook.py
@@ -25,9 +25,9 @@ from kedro_mlflow.framework.hooks.utils import (
 )
 from kedro_mlflow.io.catalog.switch_catalog_logging import switch_catalog_logging
 from kedro_mlflow.io.metrics import (
-    MlflowMetricDataSet,
-    MlflowMetricHistoryDataSet,
-    MlflowMetricsDataSet,
+    MlflowMetricDataset,
+    MlflowMetricHistoryDataset,
+    MlflowMetricsDataset,
 )
 from kedro_mlflow.mlflow import KedroPipelineModel
 from kedro_mlflow.pipeline.pipeline_ml import PipelineML
@@ -141,39 +141,39 @@ class MlflowHook:
         # we use this hooks to modif "MlflowmetricsDataset" to ensure consistency
         # of the metric name with the catalog name
         for name, dataset in catalog._data_sets.items():
-            if isinstance(dataset, MlflowMetricsDataSet) and dataset._prefix is None:
+            if isinstance(dataset, MlflowMetricsDataset) and dataset._prefix is None:
                 if dataset._run_id is not None:
-                    catalog._data_sets[name] = MlflowMetricsDataSet(
+                    catalog._data_sets[name] = MlflowMetricsDataset(
                         run_id=dataset._run_id, prefix=name
                     )
                 else:
-                    catalog._data_sets[name] = MlflowMetricsDataSet(prefix=name)
+                    catalog._data_sets[name] = MlflowMetricsDataset(prefix=name)
 
-            if isinstance(dataset, MlflowMetricDataSet) and dataset.key is None:
+            if isinstance(dataset, MlflowMetricDataset) and dataset.key is None:
                 if dataset._run_id is not None:
-                    catalog._data_sets[name] = MlflowMetricDataSet(
+                    catalog._data_sets[name] = MlflowMetricDataset(
                         run_id=dataset._run_id,
                         key=name,
                         load_args=dataset._load_args,
                         save_args=dataset._save_args,
                     )
                 else:
-                    catalog._data_sets[name] = MlflowMetricDataSet(
+                    catalog._data_sets[name] = MlflowMetricDataset(
                         key=name,
                         load_args=dataset._load_args,
                         save_args=dataset._save_args,
                     )
 
-            if isinstance(dataset, MlflowMetricHistoryDataSet) and dataset.key is None:
+            if isinstance(dataset, MlflowMetricHistoryDataset) and dataset.key is None:
                 if dataset._run_id is not None:
-                    catalog._data_sets[name] = MlflowMetricHistoryDataSet(
+                    catalog._data_sets[name] = MlflowMetricHistoryDataset(
                         run_id=dataset._run_id,
                         key=name,
                         load_args=dataset._load_args,
                         save_args=dataset._save_args,
                     )
                 else:
-                    catalog._data_sets[name] = MlflowMetricHistoryDataSet(
+                    catalog._data_sets[name] = MlflowMetricHistoryDataset(
                         key=name,
                         load_args=dataset._load_args,
                         save_args=dataset._save_args,
diff --git a/kedro_mlflow/io/artifacts/__init__.py b/kedro_mlflow/io/artifacts/__init__.py
index 0722d98..5482b87 100644
--- a/kedro_mlflow/io/artifacts/__init__.py
+++ b/kedro_mlflow/io/artifacts/__init__.py
@@ -1,1 +1,1 @@
-from .mlflow_artifact_dataset import MlflowArtifactDataSet
+from .mlflow_artifact_dataset import MlflowArtifactDataset
diff --git a/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py b/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
index 98833e4..01f2c0f 100644
--- a/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
+++ b/kedro_mlflow/io/artifacts/mlflow_artifact_dataset.py
@@ -8,7 +8,7 @@ from kedro.io.core import parse_dataset_definition
 from mlflow.tracking import MlflowClient
 
 
-class MlflowArtifactDataSet(AbstractVersionedDataSet):
+class MlflowArtifactDataset(AbstractVersionedDataSet):
     """This class is a wrapper for any kedro AbstractDataSet.
     It decorates their ``save`` method to log the dataset in mlflow when ``save`` is called.
     """
@@ -26,7 +26,7 @@ class MlflowArtifactDataSet(AbstractVersionedDataSet):
         # all dataset (i.e. it should replace AbstractVersionedDataSet)
         # instead and since we can't modify the core package,
         # we create a subclass which inherits dynamically from the data_set class
-        class MlflowArtifactDataSetChildren(data_set):
+        class MlflowArtifactDatasetChildren(data_set):
             def __init__(self, run_id, artifact_path):
                 super().__init__(**data_set_args)
                 self.run_id = run_id
@@ -135,33 +135,33 @@ class MlflowArtifactDataSet(AbstractVersionedDataSet):
 
         # rename the class
         parent_name = data_set.__name__
-        MlflowArtifactDataSetChildren.__name__ = f"Mlflow{parent_name}"
-        MlflowArtifactDataSetChildren.__qualname__ = (
+        MlflowArtifactDatasetChildren.__name__ = f"Mlflow{parent_name}"
+        MlflowArtifactDatasetChildren.__qualname__ = (
             f"{parent_name}.Mlflow{parent_name}"
         )
 
-        mlflow_dataset_instance = MlflowArtifactDataSetChildren(
+        mlflow_dataset_instance = MlflowArtifactDatasetChildren(
             run_id=run_id, artifact_path=artifact_path
         )
         return mlflow_dataset_instance
 
     def _load(self) -> Any:  # pragma: no cover
         """
-        MlflowArtifactDataSet is a factory for DataSet
+        MlflowArtifactDataset is a factory for DataSet
         and consequently does not implements abtracts methods
         """
         pass
 
     def _save(self, data: Any) -> None:  # pragma: no cover
         """
-        MlflowArtifactDataSet is a factory for DataSet
+        MlflowArtifactDataset is a factory for DataSet
         and consequently does not implements abtracts methods
         """
         pass
 
     def _describe(self) -> Dict[str, Any]:  # pragma: no cover
         """
-        MlflowArtifactDataSet is a factory for DataSet
+        MlflowArtifactDataset is a factory for DataSet
         and consequently does not implements abtracts methods
         """
         pass
diff --git a/kedro_mlflow/io/metrics/__init__.py b/kedro_mlflow/io/metrics/__init__.py
index 62e242a..1fde8b0 100644
--- a/kedro_mlflow/io/metrics/__init__.py
+++ b/kedro_mlflow/io/metrics/__init__.py
@@ -1,3 +1,3 @@
-from .mlflow_metric_dataset import MlflowMetricDataSet
-from .mlflow_metric_history_dataset import MlflowMetricHistoryDataSet
-from .mlflow_metrics_dataset import MlflowMetricsDataSet
+from .mlflow_metric_dataset import MlflowMetricDataset
+from .mlflow_metric_history_dataset import MlflowMetricHistoryDataset
+from .mlflow_metrics_dataset import MlflowMetricsDataset
diff --git a/kedro_mlflow/io/metrics/mlflow_abstract_metric_dataset.py b/kedro_mlflow/io/metrics/mlflow_abstract_metric_dataset.py
index 337af99..225f138 100644
--- a/kedro_mlflow/io/metrics/mlflow_abstract_metric_dataset.py
+++ b/kedro_mlflow/io/metrics/mlflow_abstract_metric_dataset.py
@@ -5,7 +5,7 @@ from kedro.io import AbstractDataSet
 from mlflow.tracking import MlflowClient
 
 
-class MlflowAbstractMetricDataSet(AbstractDataSet):
+class MlflowAbstractMetricDataset(AbstractDataSet):
     def __init__(
         self,
         key: str = None,
@@ -13,7 +13,7 @@ class MlflowAbstractMetricDataSet(AbstractDataSet):
         load_args: Dict[str, Any] = None,
         save_args: Dict[str, Any] = None,
     ):
-        """Initialise MlflowMetricsDataSet.
+        """Initialise MlflowMetricsDataset.
 
         Args:
             run_id (str): The ID of the mlflow run where the metric should be logged
diff --git a/kedro_mlflow/io/metrics/mlflow_metric_dataset.py b/kedro_mlflow/io/metrics/mlflow_metric_dataset.py
index b0432e3..58c428f 100644
--- a/kedro_mlflow/io/metrics/mlflow_metric_dataset.py
+++ b/kedro_mlflow/io/metrics/mlflow_metric_dataset.py
@@ -4,11 +4,11 @@ from typing import Any, Dict
 from mlflow.tracking import MlflowClient
 
 from kedro_mlflow.io.metrics.mlflow_abstract_metric_dataset import (
-    MlflowAbstractMetricDataSet,
+    MlflowAbstractMetricDataset,
 )
 
 
-class MlflowMetricDataSet(MlflowAbstractMetricDataSet):
+class MlflowMetricDataset(MlflowAbstractMetricDataset):
     SUPPORTED_SAVE_MODES = {"overwrite", "append"}
     DEFAULT_SAVE_MODE = "overwrite"
 
@@ -19,7 +19,7 @@ class MlflowMetricDataSet(MlflowAbstractMetricDataSet):
         load_args: Dict[str, Any] = None,
         save_args: Dict[str, Any] = None,
     ):
-        """Initialise MlflowMetricDataSet.
+        """Initialise MlflowMetricDataset.
         Args:
             run_id (str): The ID of the mlflow run where the metric should be logged
         """
diff --git a/kedro_mlflow/io/metrics/mlflow_metric_history_dataset.py b/kedro_mlflow/io/metrics/mlflow_metric_history_dataset.py
index 8722916..48fb734 100644
--- a/kedro_mlflow/io/metrics/mlflow_metric_history_dataset.py
+++ b/kedro_mlflow/io/metrics/mlflow_metric_history_dataset.py
@@ -3,11 +3,11 @@ from typing import Any, Dict, List, Union
 from mlflow.tracking import MlflowClient
 
 from kedro_mlflow.io.metrics.mlflow_abstract_metric_dataset import (
-    MlflowAbstractMetricDataSet,
+    MlflowAbstractMetricDataset,
 )
 
 
-class MlflowMetricHistoryDataSet(MlflowAbstractMetricDataSet):
+class MlflowMetricHistoryDataset(MlflowAbstractMetricDataset):
     def __init__(
         self,
         key: str = None,
@@ -15,7 +15,7 @@ class MlflowMetricHistoryDataSet(MlflowAbstractMetricDataSet):
         load_args: Dict[str, Any] = None,
         save_args: Dict[str, Any] = None,
     ):
-        """Initialise MlflowMetricDataSet.
+        """Initialise MlflowMetricDataset.
         Args:
             run_id (str): The ID of the mlflow run where the metric should be logged
         """
diff --git a/kedro_mlflow/io/metrics/mlflow_metrics_dataset.py b/kedro_mlflow/io/metrics/mlflow_metrics_dataset.py
index d9091c2..2c4f228 100644
--- a/kedro_mlflow/io/metrics/mlflow_metrics_dataset.py
+++ b/kedro_mlflow/io/metrics/mlflow_metrics_dataset.py
@@ -11,7 +11,7 @@ MetricTuple = Tuple[str, float, int]
 MetricsDict = Dict[str, MetricItem]
 
 
-class MlflowMetricsDataSet(AbstractDataSet):
+class MlflowMetricsDataset(AbstractDataSet):
     """This class represent MLflow metrics dataset."""
 
     def __init__(
@@ -19,7 +19,7 @@ class MlflowMetricsDataSet(AbstractDataSet):
         run_id: str = None,
         prefix: Optional[str] = None,
     ):
-        """Initialise MlflowMetricsDataSet.
+        """Initialise MlflowMetricsDataset.
 
         Args:
             prefix (Optional[str]): Prefix for metrics logged in MLflow.
diff --git a/kedro_mlflow/io/models/__init__.py b/kedro_mlflow/io/models/__init__.py
index 6d102f8..dcdefe3 100644
--- a/kedro_mlflow/io/models/__init__.py
+++ b/kedro_mlflow/io/models/__init__.py
@@ -1,3 +1,3 @@
 from .mlflow_model_logger_dataset import MlflowModelLoggerDataSet
-from .mlflow_model_registry_dataset import MlflowModelRegistryDataSet
+from .mlflow_model_registry_dataset import MlflowModelRegistryDataset
 from .mlflow_model_saver_dataset import MlflowModelSaverDataSet
diff --git a/kedro_mlflow/io/models/mlflow_abstract_model_dataset.py b/kedro_mlflow/io/models/mlflow_abstract_model_dataset.py
index d850769..7cf1a62 100644
--- a/kedro_mlflow/io/models/mlflow_abstract_model_dataset.py
+++ b/kedro_mlflow/io/models/mlflow_abstract_model_dataset.py
@@ -7,7 +7,7 @@ from kedro.io import AbstractVersionedDataSet, Version
 from kedro.io.core import DataSetError
 
 
-class MlflowAbstractModelDataSet(AbstractVersionedDataSet):
+class MlflowModelRegistryDataset(AbstractVersionedDataSet):
     """
     Absract mother class for model datasets.
     """
diff --git a/kedro_mlflow/io/models/mlflow_model_logger_dataset.py b/kedro_mlflow/io/models/mlflow_model_logger_dataset.py
index 988045c..b218c5d 100644
--- a/kedro_mlflow/io/models/mlflow_model_logger_dataset.py
+++ b/kedro_mlflow/io/models/mlflow_model_logger_dataset.py
@@ -4,11 +4,11 @@ import mlflow
 from kedro.io.core import DataSetError
 
 from kedro_mlflow.io.models.mlflow_abstract_model_dataset import (
-    MlflowAbstractModelDataSet,
+    MlflowModelRegistryDataset,
 )
 
 
-class MlflowModelLoggerDataSet(MlflowAbstractModelDataSet):
+class MlflowModelLoggerDataSet(MlflowModelRegistryDataset):
     """Wrapper for saving, logging and loading for all MLflow model flavor."""
 
     def __init__(
diff --git a/kedro_mlflow/io/models/mlflow_model_registry_dataset.py b/kedro_mlflow/io/models/mlflow_model_registry_dataset.py
index 1354427..1043823 100644
--- a/kedro_mlflow/io/models/mlflow_model_registry_dataset.py
+++ b/kedro_mlflow/io/models/mlflow_model_registry_dataset.py
@@ -1,11 +1,11 @@
 from typing import Any, Dict, Optional, Union
 
 from kedro_mlflow.io.models.mlflow_abstract_model_dataset import (
-    MlflowAbstractModelDataSet,
+    MlflowModelRegistryDataset,
 )
 
 
-class MlflowModelRegistryDataSet(MlflowAbstractModelDataSet):
+class MlflowModelRegistryDataset(MlflowModelRegistryDataset):
     """Wrapper for saving, logging and loading for all MLflow model flavor."""
 
     def __init__(
@@ -16,7 +16,7 @@ class MlflowModelRegistryDataSet(MlflowAbstractModelDataSet):
         pyfunc_workflow: Optional[str] = "python_model",
         load_args: Optional[Dict[str, Any]] = None,
     ) -> None:
-        """Initialize the Kedro MlflowModelRegistryDataSet.
+        """Initialize the Kedro MlflowModelRegistryDataset.
 
         Parameters are passed from the Data Catalog.
 
@@ -66,7 +66,7 @@ class MlflowModelRegistryDataSet(MlflowAbstractModelDataSet):
 
     def _save(self, model: Any) -> None:
         raise NotImplementedError(
-            "The 'save' method is not implemented for MlflowModelRegistryDataSet. You can pass 'registered_model_name' argument in 'MLflowModelLoggerDataSet(..., save_args={registered_model_name='my_model'}' to save and register a model in the same step. "
+            "The 'save' method is not implemented for MlflowModelRegistryDataset. You can pass 'registered_model_name' argument in 'MLflowModelLoggerDataSet(..., save_args={registered_model_name='my_model'}' to save and register a model in the same step. "
         )
 
     def _describe(self) -> Dict[str, Any]:
diff --git a/kedro_mlflow/io/models/mlflow_model_saver_dataset.py b/kedro_mlflow/io/models/mlflow_model_saver_dataset.py
index e1b21fc..985780a 100644
--- a/kedro_mlflow/io/models/mlflow_model_saver_dataset.py
+++ b/kedro_mlflow/io/models/mlflow_model_saver_dataset.py
@@ -5,11 +5,11 @@ from typing import Any, Dict, Optional
 from kedro.io import Version
 
 from kedro_mlflow.io.models.mlflow_abstract_model_dataset import (
-    MlflowAbstractModelDataSet,
+    MlflowModelRegistryDataset,
 )
 
 
-class MlflowModelSaverDataSet(MlflowAbstractModelDataSet):
+class MlflowModelSaverDataSet(MlflowModelRegistryDataset):
     """Wrapper for saving, logging and loading for all MLflow model flavor."""
 
     def __init__(


Test patch:
diff --git a/tests/framework/hooks/test_hook_deactivate_tracking.py b/tests/framework/hooks/test_hook_deactivate_tracking.py
index f8cf3ef..97eac65 100644
--- a/tests/framework/hooks/test_hook_deactivate_tracking.py
+++ b/tests/framework/hooks/test_hook_deactivate_tracking.py
@@ -78,20 +78,20 @@ def catalog_config(kedro_project_path):
     fake_data_filepath = str(kedro_project_path / "fake_data.pkl")
     return {
         "artifact_data": {
-            "type": "kedro_mlflow.io.artifacts.MlflowArtifactDataSet",
+            "type": "kedro_mlflow.io.artifacts.MlflowArtifactDataset",
             "data_set": {
                 "type": "pickle.PickleDataSet",
                 "filepath": fake_data_filepath,
             },
         },
         "metrics_data": {
-            "type": "kedro_mlflow.io.metrics.MlflowMetricsDataSet",
+            "type": "kedro_mlflow.io.metrics.MlflowMetricsDataset",
         },
         "metric_data": {
-            "type": "kedro_mlflow.io.metrics.MlflowMetricDataSet",
+            "type": "kedro_mlflow.io.metrics.MlflowMetricDataset",
         },
         "metric_history_data": {
-            "type": "kedro_mlflow.io.metrics.MlflowMetricHistoryDataSet",
+            "type": "kedro_mlflow.io.metrics.MlflowMetricHistoryDataset",
         },
         "model": {
             "type": "kedro_mlflow.io.models.MlflowModelLoggerDataSet",
diff --git a/tests/framework/hooks/test_hook_log_metrics.py b/tests/framework/hooks/test_hook_log_metrics.py
index 2d0bfc0..04c761a 100644
--- a/tests/framework/hooks/test_hook_log_metrics.py
+++ b/tests/framework/hooks/test_hook_log_metrics.py
@@ -10,9 +10,9 @@ from kedro.runner import SequentialRunner
 
 from kedro_mlflow.framework.hooks.mlflow_hook import MlflowHook
 from kedro_mlflow.io.metrics import (
-    MlflowMetricDataSet,
-    MlflowMetricHistoryDataSet,
-    MlflowMetricsDataSet,
+    MlflowMetricDataset,
+    MlflowMetricHistoryDataset,
+    MlflowMetricsDataset,
 )
 
 
@@ -105,12 +105,12 @@ def dummy_catalog(tmp_path):
             "params:unused_param": MemoryDataSet("blah"),
             "data": MemoryDataSet(),
             "model": PickleDataSet((tmp_path / "model.csv").as_posix()),
-            "my_metrics": MlflowMetricsDataSet(),
-            "another_metrics": MlflowMetricsDataSet(prefix="foo"),
-            "my_metric": MlflowMetricDataSet(),
-            "another_metric": MlflowMetricDataSet(key="foo"),
-            "my_metric_history": MlflowMetricHistoryDataSet(),
-            "another_metric_history": MlflowMetricHistoryDataSet(key="bar"),
+            "my_metrics": MlflowMetricsDataset(),
+            "another_metrics": MlflowMetricsDataset(prefix="foo"),
+            "my_metric": MlflowMetricDataset(),
+            "another_metric": MlflowMetricDataset(key="foo"),
+            "my_metric_history": MlflowMetricHistoryDataset(),
+            "another_metric_history": MlflowMetricHistoryDataset(key="bar"),
         }
     )
     return dummy_catalog
@@ -181,16 +181,16 @@ def test_mlflow_hook_metrics_dataset_with_run_id(
                 "model": PickleDataSet(
                     (kedro_project_with_mlflow_conf / "data" / "model.csv").as_posix()
                 ),
-                "my_metrics": MlflowMetricsDataSet(run_id=existing_run_id),
-                "another_metrics": MlflowMetricsDataSet(
+                "my_metrics": MlflowMetricsDataset(run_id=existing_run_id),
+                "another_metrics": MlflowMetricsDataset(
                     run_id=existing_run_id, prefix="foo"
                 ),
-                "my_metric": MlflowMetricDataSet(run_id=existing_run_id),
-                "another_metric": MlflowMetricDataSet(
+                "my_metric": MlflowMetricDataset(run_id=existing_run_id),
+                "another_metric": MlflowMetricDataset(
                     run_id=existing_run_id, key="foo"
                 ),
-                "my_metric_history": MlflowMetricHistoryDataSet(run_id=existing_run_id),
-                "another_metric_history": MlflowMetricHistoryDataSet(
+                "my_metric_history": MlflowMetricHistoryDataset(run_id=existing_run_id),
+                "another_metric_history": MlflowMetricHistoryDataset(
                     run_id=existing_run_id, key="bar"
                 ),
             }
diff --git a/tests/io/artifacts/test_mlflow_artifact_dataset.py b/tests/io/artifacts/test_mlflow_artifact_dataset.py
index 464279e..802ffcf 100644
--- a/tests/io/artifacts/test_mlflow_artifact_dataset.py
+++ b/tests/io/artifacts/test_mlflow_artifact_dataset.py
@@ -9,7 +9,7 @@ from kedro.io import PartitionedDataSet
 from mlflow.tracking import MlflowClient
 from pytest_lazyfixture import lazy_fixture
 
-from kedro_mlflow.io.artifacts import MlflowArtifactDataSet
+from kedro_mlflow.io.artifacts import MlflowArtifactDataset
 
 
 @pytest.fixture
@@ -52,7 +52,7 @@ def test_mlflow_csv_pkl_dataset_save_reload(
     mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
     filepath = (tmp_path / "data").with_suffix(extension)
 
-    mlflow_dataset = MlflowArtifactDataSet(
+    mlflow_dataset = MlflowArtifactDataset(
         artifact_path=artifact_path,
         data_set=dict(type=dataset, filepath=filepath.as_posix()),
     )
@@ -98,7 +98,7 @@ def test_artifact_dataset_save_with_run_id(
         nb_runs += 1
 
     # then same scenario but the run_id where data is saved is specified
-    mlflow_csv_dataset = MlflowArtifactDataSet(
+    mlflow_csv_dataset = MlflowArtifactDataset(
         data_set=dict(type=CSVDataSet, filepath=(tmp_path / "df1.csv").as_posix()),
         run_id=run_id,
     )
@@ -134,7 +134,7 @@ def test_is_versioned_dataset_logged_correctly_in_mlflow(tmp_path, tracking_uri,
     with mlflow.start_run():
         run_id = mlflow.active_run().info.run_id
 
-        mlflow_csv_dataset = MlflowArtifactDataSet(
+        mlflow_csv_dataset = MlflowArtifactDataset(
             data_set=dict(
                 type=CSVDataSet,
                 filepath=(tmp_path / "df1.csv").as_posix(),
@@ -172,7 +172,7 @@ def test_is_versioned_dataset_logged_correctly_in_mlflow(tmp_path, tracking_uri,
 
 
 def test_artifact_dataset_logging_deactivation(tmp_path, tracking_uri):
-    mlflow_pkl_dataset = MlflowArtifactDataSet(
+    mlflow_pkl_dataset = MlflowArtifactDataset(
         data_set=dict(type=PickleDataSet, filepath=(tmp_path / "df1.csv").as_posix())
     )
 
@@ -199,7 +199,7 @@ def test_artifact_dataset_logging_deactivation(tmp_path, tracking_uri):
 
 
 def test_mlflow_artifact_logging_deactivation_is_bool(tmp_path):
-    mlflow_csv_dataset = MlflowArtifactDataSet(
+    mlflow_csv_dataset = MlflowArtifactDataset(
         data_set=dict(type=CSVDataSet, filepath=(tmp_path / "df1.csv").as_posix())
     )
 
@@ -211,7 +211,7 @@ def test_artifact_dataset_load_with_run_id(tmp_path, tracking_uri, df1, df2):
     mlflow.set_tracking_uri(tracking_uri.as_uri())
 
     # define the logger
-    mlflow_csv_dataset = MlflowArtifactDataSet(
+    mlflow_csv_dataset = MlflowArtifactDataset(
         data_set=dict(type=CSVDataSet, filepath=(tmp_path / "df.csv").as_posix())
     )
 
@@ -239,7 +239,7 @@ def test_artifact_dataset_load_with_run_id_and_artifact_path(
     mlflow.set_tracking_uri(tracking_uri.as_uri())
 
     # save first and retrieve run id
-    mlflow_csv_dataset1 = MlflowArtifactDataSet(
+    mlflow_csv_dataset1 = MlflowArtifactDataset(
         data_set=dict(type=CSVDataSet, filepath=(tmp_path / "df1.csv").as_posix()),
         artifact_path=artifact_path,
     )
@@ -250,7 +250,7 @@ def test_artifact_dataset_load_with_run_id_and_artifact_path(
             tmp_path / "df1.csv"
         ).unlink()  # we need to delete the data, else it is automatically reused instead of downloading
     # same as before, but a closed run_id is specified
-    mlflow_csv_dataset2 = MlflowArtifactDataSet(
+    mlflow_csv_dataset2 = MlflowArtifactDataset(
         data_set=dict(type=CSVDataSet, filepath=(tmp_path / "df1.csv").as_posix()),
         artifact_path=artifact_path,
         run_id=first_run_id,
@@ -269,7 +269,7 @@ def test_partitioned_dataset_save_and_reload(
     mlflow.set_tracking_uri(tracking_uri.as_uri())
     mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
 
-    mlflow_dataset = MlflowArtifactDataSet(
+    mlflow_dataset = MlflowArtifactDataset(
         artifact_path=artifact_path,
         data_set=dict(
             type=PartitionedDataSet,
diff --git a/tests/io/metrics/test_mlflow_metric_dataset.py b/tests/io/metrics/test_mlflow_metric_dataset.py
index 515067e..399a78c 100644
--- a/tests/io/metrics/test_mlflow_metric_dataset.py
+++ b/tests/io/metrics/test_mlflow_metric_dataset.py
@@ -3,7 +3,7 @@ import pytest
 from kedro.io.core import DataSetError
 from mlflow.tracking import MlflowClient
 
-from kedro_mlflow.io.metrics import MlflowMetricDataSet
+from kedro_mlflow.io.metrics import MlflowMetricDataset
 
 
 @pytest.fixture
@@ -21,13 +21,13 @@ def mlflow_client(mlflow_tracking_uri):
 
 def test_mlflow_wrong_save_mode():
     with pytest.raises(DataSetError, match=r"save_args\['mode'\] must be one of"):
-        metric_ds = MlflowMetricDataSet(key="my_metric", save_args={"mode": "bad_mode"})
+        metric_ds = MlflowMetricDataset(key="my_metric", save_args={"mode": "bad_mode"})
         with mlflow.start_run():
             metric_ds.save(0.3)
 
 
 def test_mlflow_metric_dataset_save_without_active_run_or_run_id():
-    metric_ds = MlflowMetricDataSet(key="my_metric")
+    metric_ds = MlflowMetricDataset(key="my_metric")
     with pytest.raises(
         DataSetError,
         match="You must either specify a run_id or have a mlflow active run opened",
@@ -47,7 +47,7 @@ def test_mlflow_metric_dataset_save_without_active_run_or_run_id():
     ],
 )
 def test_mlflow_metric_dataset_save_with_active_run(mlflow_client, save_args):
-    metric_ds = MlflowMetricDataSet(key="my_metric", save_args=save_args)
+    metric_ds = MlflowMetricDataset(key="my_metric", save_args=save_args)
     with mlflow.start_run():
         metric_ds.save(0.3)
         run_id = mlflow.active_run().info.run_id
@@ -72,11 +72,11 @@ def test_mlflow_metric_dataset_save_with_active_run(mlflow_client, save_args):
 )
 def test_mlflow_metric_dataset_save_with_run_id(mlflow_client, save_args):
     # this time, the run is created first and closed
-    # the MlflowMetricDataSet should reopen it to interact
+    # the MlflowMetricDataset should reopen it to interact
     with mlflow.start_run():
         run_id = mlflow.active_run().info.run_id
 
-    metric_ds = MlflowMetricDataSet(run_id=run_id, key="my_metric", save_args=save_args)
+    metric_ds = MlflowMetricDataset(run_id=run_id, key="my_metric", save_args=save_args)
     metric_ds.save(0.3)
     metric_history = mlflow_client.get_metric_history(run_id, "my_metric")
     step = 0 if save_args is None else save_args.get("step", 0)
@@ -90,7 +90,7 @@ def test_mlflow_metric_dataset_save_append_mode(mlflow_client):
     with mlflow.start_run():
         run_id = mlflow.active_run().info.run_id
 
-    metric_ds = MlflowMetricDataSet(
+    metric_ds = MlflowMetricDataset(
         run_id=run_id, key="my_metric", save_args={"mode": "append"}
     )
 
@@ -108,7 +108,7 @@ def test_mlflow_metric_dataset_save_overwrite_mode(mlflow_client):
         run_id = mlflow.active_run().info.run_id
 
     # overwrite is the default mode
-    metric_ds = MlflowMetricDataSet(run_id=run_id, key="my_metric")
+    metric_ds = MlflowMetricDataset(run_id=run_id, key="my_metric")
 
     metric_ds.save(0.3)
     metric_ds.save(1)
@@ -125,7 +125,7 @@ def test_mlflow_metric_dataset_load():
         mlflow.log_metric(key="awesome_metric", value=0.1)
 
     # overwrite is the default mode
-    metric_ds = MlflowMetricDataSet(run_id=run_id, key="awesome_metric")
+    metric_ds = MlflowMetricDataset(run_id=run_id, key="awesome_metric")
 
     assert metric_ds.load() == 0.1
 
@@ -139,7 +139,7 @@ def test_mlflow_metric_dataset_load_last_logged_by_default_if_unordered():
         mlflow.log_metric(key="awesome_metric", value=0.1, step=0)
 
     # overwrite is the default mode
-    metric_ds = MlflowMetricDataSet(run_id=run_id, key="awesome_metric")
+    metric_ds = MlflowMetricDataset(run_id=run_id, key="awesome_metric")
 
     assert (
         metric_ds.load() == 0.1
@@ -155,7 +155,7 @@ def test_mlflow_metric_dataset_load_given_step():
         mlflow.log_metric(key="awesome_metric", value=0.4, step=3)
 
     # overwrite is the default mode
-    metric_ds = MlflowMetricDataSet(
+    metric_ds = MlflowMetricDataset(
         run_id=run_id, key="awesome_metric", load_args={"step": 2}
     )
 
@@ -173,7 +173,7 @@ def test_mlflow_metric_dataset_load_last_given_step_if_duplicated():
         mlflow.log_metric(key="awesome_metric", value=0.4, step=3)
 
     # overwrite is the default mode
-    metric_ds = MlflowMetricDataSet(
+    metric_ds = MlflowMetricDataset(
         run_id=run_id, key="awesome_metric", load_args={"step": 2}
     )
 
@@ -181,7 +181,7 @@ def test_mlflow_metric_dataset_load_last_given_step_if_duplicated():
 
 
 def test_mlflow_metric_dataset_logging_deactivation(mlflow_tracking_uri):
-    metric_ds = MlflowMetricDataSet(key="inactive_metric")
+    metric_ds = MlflowMetricDataset(key="inactive_metric")
     metric_ds._logging_activated = False
     with mlflow.start_run():
         metric_ds.save(1)
@@ -189,7 +189,7 @@ def test_mlflow_metric_dataset_logging_deactivation(mlflow_tracking_uri):
 
 
 def test_mlflow_metric_logging_deactivation_is_bool():
-    mlflow_metric_dataset = MlflowMetricDataSet(key="hello")
+    mlflow_metric_dataset = MlflowMetricDataset(key="hello")
 
     with pytest.raises(ValueError, match="_logging_activated must be a boolean"):
         mlflow_metric_dataset._logging_activated = "hello"
diff --git a/tests/io/metrics/test_mlflow_metric_history_dataset.py b/tests/io/metrics/test_mlflow_metric_history_dataset.py
index 5ac1baa..438d830 100644
--- a/tests/io/metrics/test_mlflow_metric_history_dataset.py
+++ b/tests/io/metrics/test_mlflow_metric_history_dataset.py
@@ -2,7 +2,7 @@ import mlflow
 import pytest
 from mlflow.tracking import MlflowClient
 
-from kedro_mlflow.io.metrics import MlflowMetricHistoryDataSet
+from kedro_mlflow.io.metrics import MlflowMetricHistoryDataset
 
 
 @pytest.fixture
@@ -44,7 +44,7 @@ def test_mlflow_metric_history_dataset_save_load(mlflow_client, save_mode, load_
         "history": metric_as_history,
     }
 
-    metric_ds_saver = MlflowMetricHistoryDataSet(
+    metric_ds_saver = MlflowMetricHistoryDataset(
         key="my_metric", save_args={"mode": save_mode}
     )
     with mlflow.start_run():
@@ -55,7 +55,7 @@ def test_mlflow_metric_history_dataset_save_load(mlflow_client, save_mode, load_
     run = mlflow_client.get_run(run_id)
     assert "my_metric" in run.data.metrics.keys()
 
-    metric_ds_loader = MlflowMetricHistoryDataSet(
+    metric_ds_loader = MlflowMetricHistoryDataset(
         key="my_metric", run_id=run_id, load_args={"mode": load_mode}
     )
 
@@ -63,7 +63,7 @@ def test_mlflow_metric_history_dataset_save_load(mlflow_client, save_mode, load_
 
 
 def test_mlflow_metric_history_dataset_logging_deactivation(mlflow_tracking_uri):
-    metric_ds = MlflowMetricHistoryDataSet(key="inactive_metric")
+    metric_ds = MlflowMetricHistoryDataset(key="inactive_metric")
     metric_ds._logging_activated = False
     with mlflow.start_run():
         metric_ds.save([0.1])
diff --git a/tests/io/metrics/test_mlflow_metrics_dataset.py b/tests/io/metrics/test_mlflow_metrics_dataset.py
index 0125df9..14c568c 100644
--- a/tests/io/metrics/test_mlflow_metrics_dataset.py
+++ b/tests/io/metrics/test_mlflow_metrics_dataset.py
@@ -6,7 +6,7 @@ from kedro.io import DataSetError
 from mlflow.tracking import MlflowClient
 from pytest_lazyfixture import lazy_fixture
 
-from kedro_mlflow.io.metrics import MlflowMetricsDataSet
+from kedro_mlflow.io.metrics import MlflowMetricsDataset
 
 
 def assert_are_metrics_logged(
@@ -81,12 +81,12 @@ def metrics3():
     ],
 )
 def test_mlflow_metrics_dataset_saved_and_logged(tmp_path, tracking_uri, data, prefix):
-    """Check if MlflowMetricsDataSet can be saved in catalog when filepath is given,
+    """Check if MlflowMetricsDataset can be saved in catalog when filepath is given,
     and if logged in mlflow.
     """
     mlflow.set_tracking_uri(tracking_uri.as_uri())
     mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
-    mlflow_metrics_dataset = MlflowMetricsDataSet(prefix=prefix)
+    mlflow_metrics_dataset = MlflowMetricsDataset(prefix=prefix)
 
     with mlflow.start_run():
         run_id = mlflow.active_run().info.run_id
@@ -96,7 +96,7 @@ def test_mlflow_metrics_dataset_saved_and_logged(tmp_path, tracking_uri, data, p
         assert_are_metrics_logged(data, mlflow_client, run_id, prefix)
 
     # Check if metrics are stored in catalog.
-    catalog_metrics = MlflowMetricsDataSet(
+    catalog_metrics = MlflowMetricsDataset(
         prefix=prefix,
         # Run id needs to be provided as there is no active run.
         run_id=run_id,
@@ -109,14 +109,14 @@ def test_mlflow_metrics_dataset_saved_and_logged(tmp_path, tracking_uri, data, p
 
 
 def test_mlflow_metrics_dataset_saved_without_run_id(tmp_path, tracking_uri, metrics3):
-    """Check if MlflowMetricsDataSet can be saved in catalog when filepath is given,
+    """Check if MlflowMetricsDataset can be saved in catalog when filepath is given,
     and if logged in mlflow.
     """
     prefix = "test_metric"
 
     mlflow.set_tracking_uri(tracking_uri.as_uri())
     mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
-    mlflow_metrics_dataset = MlflowMetricsDataSet(prefix=prefix)
+    mlflow_metrics_dataset = MlflowMetricsDataset(prefix=prefix)
 
     # a mlflow run_id is automatically created
     mlflow_metrics_dataset.save(metrics3)
@@ -126,13 +126,13 @@ def test_mlflow_metrics_dataset_saved_without_run_id(tmp_path, tracking_uri, met
 
 
 def test_mlflow_metrics_dataset_exists(tmp_path, tracking_uri, metrics3):
-    """Check if MlflowMetricsDataSet is well identified as
+    """Check if MlflowMetricsDataset is well identified as
     existing if it has already been saved.
     """
     prefix = "test_metric"
 
     mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_metrics_dataset = MlflowMetricsDataSet(prefix=prefix)
+    mlflow_metrics_dataset = MlflowMetricsDataset(prefix=prefix)
 
     # a mlflow run_id is automatically created
     mlflow_metrics_dataset.save(metrics3)
@@ -140,7 +140,7 @@ def test_mlflow_metrics_dataset_exists(tmp_path, tracking_uri, metrics3):
 
 
 def test_mlflow_metrics_dataset_does_not_exist(tmp_path, tracking_uri, metrics3):
-    """Check if MlflowMetricsDataSet is well identified as
+    """Check if MlflowMetricsDataset is well identified as
     not existingif it has never been saved.
     """
 
@@ -148,7 +148,7 @@ def test_mlflow_metrics_dataset_does_not_exist(tmp_path, tracking_uri, metrics3)
     mlflow.start_run()  # starts a run toenable mlflow_metrics_dataset to know where to seacrh
     run_id = mlflow.active_run().info.run_id
     mlflow.end_run()
-    mlflow_metrics_dataset = MlflowMetricsDataSet(prefix="test_metric", run_id=run_id)
+    mlflow_metrics_dataset = MlflowMetricsDataset(prefix="test_metric", run_id=run_id)
     # a mlflow run_id is automatically created
     assert not mlflow_metrics_dataset.exists()
 
@@ -156,12 +156,12 @@ def test_mlflow_metrics_dataset_does_not_exist(tmp_path, tracking_uri, metrics3)
 def test_mlflow_metrics_dataset_fails_with_invalid_metric(
     tmp_path, tracking_uri, metrics3
 ):
-    """Check if MlflowMetricsDataSet is well identified as
+    """Check if MlflowMetricsDataset is well identified as
     not existingif it has never been saved.
     """
 
     mlflow.set_tracking_uri(tracking_uri.as_uri())
-    mlflow_metrics_dataset = MlflowMetricsDataSet(prefix="test_metric")
+    mlflow_metrics_dataset = MlflowMetricsDataset(prefix="test_metric")
 
     with pytest.raises(
         DataSetError, match="Unexpected metric value. Should be of type"
@@ -172,7 +172,7 @@ def test_mlflow_metrics_dataset_fails_with_invalid_metric(
 
 
 def test_mlflow_metrics_logging_deactivation(tracking_uri, metrics):
-    mlflow_metrics_dataset = MlflowMetricsDataSet(prefix="hello")
+    mlflow_metrics_dataset = MlflowMetricsDataset(prefix="hello")
 
     mlflow.set_tracking_uri(tracking_uri.as_uri())
     mlflow_client = MlflowClient(tracking_uri=tracking_uri.as_uri())
@@ -201,7 +201,7 @@ def test_mlflow_metrics_logging_deactivation(tracking_uri, metrics):
 
 
 def test_mlflow_metrics_logging_deactivation_is_bool():
-    mlflow_metrics_dataset = MlflowMetricsDataSet(prefix="hello")
+    mlflow_metrics_dataset = MlflowMetricsDataset(prefix="hello")
 
     with pytest.raises(ValueError, match="_logging_activated must be a boolean"):
         mlflow_metrics_dataset._logging_activated = "hello"
diff --git a/tests/io/models/test_mlflow_model_registry_dataset.py b/tests/io/models/test_mlflow_model_registry_dataset.py
index 2252830..0b2aa5d 100644
--- a/tests/io/models/test_mlflow_model_registry_dataset.py
+++ b/tests/io/models/test_mlflow_model_registry_dataset.py
@@ -4,14 +4,14 @@ from kedro.io.core import DataSetError
 from mlflow import MlflowClient
 from sklearn.tree import DecisionTreeClassifier
 
-from kedro_mlflow.io.models import MlflowModelRegistryDataSet
+from kedro_mlflow.io.models import MlflowModelRegistryDataset
 
 
 def test_mlflow_model_registry_save_not_implemented(tmp_path):
-    ml_ds = MlflowModelRegistryDataSet(model_name="demo_model")
+    ml_ds = MlflowModelRegistryDataset(model_name="demo_model")
     with pytest.raises(
         DataSetError,
-        match=r"The 'save' method is not implemented for MlflowModelRegistryDataSet",
+        match=r"The 'save' method is not implemented for MlflowModelRegistryDataset",
     ):
         ml_ds.save(DecisionTreeClassifier())
 
@@ -40,18 +40,18 @@ def test_mlflow_model_registry_load_given_stage_or_version(tmp_path, monkeypatch
     client.transition_model_version_stage(name="demo_model", version=2, stage="Staging")
 
     # case 1: no version is provided, we take the last one
-    ml_ds = MlflowModelRegistryDataSet(model_name="demo_model")
+    ml_ds = MlflowModelRegistryDataset(model_name="demo_model")
     loaded_model = ml_ds.load()
     assert loaded_model.metadata.run_id == runs[3]
 
     # case 2: a stage is provided, we take the last model with this stage
-    ml_ds = MlflowModelRegistryDataSet(
+    ml_ds = MlflowModelRegistryDataset(
         model_name="demo_model", stage_or_version="staging"
     )
     loaded_model = ml_ds.load()
     assert loaded_model.metadata.run_id == runs[2]
 
     # case 3: a version is provided, we take the associated model
-    ml_ds = MlflowModelRegistryDataSet(model_name="demo_model", stage_or_version="1")
+    ml_ds = MlflowModelRegistryDataset(model_name="demo_model", stage_or_version="1")
     loaded_model = ml_ds.load()
     assert loaded_model.metadata.run_id == runs[1]