Repository: Lightning-Universe/lightning-flash

Issue:
Add FlashDataset
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
 Please also include relevant motivation and context.
 List any dependencies that are required for this change.
-->

This PR is part of an internal refactor within Flash. The FlashDataset is meant to replace entirely the DataSource / AutoDataset components.

Fixes # (issue)

## Before submitting
- [ ] Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
- [ ] Did you read the [contributor guideline](https://github.com/PyTorchLightning/lightning-flash/tree/master/.github/CONTRIBUTING.md), Pull Request section?
- [ ] Did you make sure your PR does only one thing, instead of bundling different changes together?
- [ ] Did you make sure to update the documentation with your changes?
- [ ] Did you write any new necessary tests? [not needed for typos/docs]
- [ ] Did you verify new and existing tests pass locally with your changes?
- [ ] If you made a notable change (that affects users), did you update the [CHANGELOG](https://github.com/PyTorchLightning/lightning-flash/blob/master/CHANGELOG.md)?

<!-- For CHANGELOG separate each item in unreleased section by a blank line to reduce collisions -->

## PR review
 - [ ] Is this pull request ready for review? (if not, please submit in draft mode)

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

## Did you have fun?
Make sure you had fun coding 🙃



Discussion:
# [Codecov](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) Report
> Merging [#851](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) (34af344) into [master](https://codecov.io/gh/PyTorchLightning/lightning-flash/commit/fcb0a900c168dcf1c4e3564b8dbbcd83ab8fa9ee?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) (fcb0a90) will **increase** coverage by `0.02%`.
> The diff coverage is `92.75%`.

[![Impacted file tree graph](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851/graphs/tree.svg?width=650&height=150&src=pr&token=oLuUr9q1vt&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning)](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning)

```diff
@@            Coverage Diff             @@
##           master     #851      +/-   ##
==========================================
+ Coverage   89.00%   89.02%   +0.02%     
==========================================
  Files         224      225       +1     
  Lines       12117    12186      +69     
==========================================
+ Hits        10785    10849      +64     
- Misses       1332     1337       +5     
```

| Flag | Coverage Δ | |
|---|---|---|
| unittests | `89.02% <92.75%> (+0.02%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning) | Coverage Δ | |
|---|---|---|
| [flash/core/data/datasets.py](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning#diff-Zmxhc2gvY29yZS9kYXRhL2RhdGFzZXRzLnB5) | `92.75% <92.75%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning). Last update [fcb0a90...34af344](https://codecov.io/gh/PyTorchLightning/lightning-flash/pull/851?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PyTorchLightning).


Resolving patch:
diff --git a/flash/core/data/datasets.py b/flash/core/data/datasets.py
--- a/flash/core/data/datasets.py
+++ b/flash/core/data/datasets.py
@@ -11,7 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Generic, Iterable, Optional, Sequence, Type, TypeVar
+from abc import abstractmethod
+from typing import Any, Callable, Iterable, Mapping, Optional, Type, Union
 
 from pytorch_lightning.trainer.states import RunningStage
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
@@ -20,24 +21,26 @@
 from flash.core.data.properties import Properties
 from flash.core.registry import FlashRegistry
 
-DATA_TYPE = TypeVar("DATA_TYPE")
 
-
-class BaseDataset(Generic[DATA_TYPE], Properties):
+class BaseDataset(Properties):
 
     DATASET_KEY = "dataset"
 
     transform_registry: Optional[FlashRegistry] = None
 
-    def load_data(self, data: Any) -> Any:
-        return data
+    @abstractmethod
+    def load_data(self, data: Any) -> Union[Iterable, Mapping]:
+        """The `load_data` hook should return either a Mapping or an Iterable.
+
+        Override to add your dataset logic creation logic.
+        """
 
+    @abstractmethod
     def load_sample(self, data: Any) -> Any:
-        return data
+        """The `load_sample` hook contains the logic to load a single sample."""
 
     def __init__(self, running_stage: RunningStage) -> None:
         super().__init__()
-
         self.running_stage = running_stage
 
     def pass_args_to_load_data(
@@ -60,7 +63,7 @@ def running_stage(self, running_stage: RunningStage) -> None:
     def _resolve_functions(self, func_name: str, cls: Type["BaseDataset"]) -> None:
         from flash.core.data.data_pipeline import DataPipeline  # noqa F811
 
-        function: Callable[[DATA_TYPE, Optional[Any]], Any] = getattr(
+        function: Callable[[Any, Optional[Any]], Any] = getattr(
             self,
             DataPipeline._resolve_function_hierarchy(
                 func_name,
@@ -80,7 +83,7 @@ def _call_load_sample(self, sample: Any) -> Any:
     def from_data(
         cls,
         *load_data_args,
-        running_stage: RunningStage = None,
+        running_stage: Optional[RunningStage] = None,
         **dataset_kwargs: Any,
     ) -> "BaseDataset":
         if not running_stage:
@@ -95,20 +98,30 @@ def from_data(
     def resolve_functions(self):
         raise NotImplementedError
 
-    _load_data = load_data
-    _load_sample = load_sample
+    # Set to None as they are dymically resolved when the dataset is made stage aware
+    # c.f running_stage is set in `__init__` function.
+    _load_data = None
+    _load_sample = None
 
 
-class FlashDataset(BaseDataset[Sequence], Dataset):
+class FlashDataset(Dataset, BaseDataset):
     """The ``FlashDataset`` is a ``BaseDataset`` and a :class:`~torch.utils.data.Dataset`.
 
-    The `data` argument must be a ``Sequence`` (it must have a length).
+    The `data` argument must be a ``Mapping`` (it must have a length).
     """
 
-    def load_data(self, data: Any) -> Any:
+    def load_data(self, data: Any) -> Mapping:
+        """By default, the `load_data` perform an identity operation.
+
+        Override to add your own logic to load the data.
+        """
         return data
 
     def load_sample(self, data: Any) -> Any:
+        """By default, the `load_sample` perform an identity operation.
+
+        Override to add your own logic to load a single sample.
+        """
         return data
 
     def resolve_functions(self):
@@ -122,16 +135,24 @@ def __len__(self) -> int:
         return len(self.data)
 
 
-class FlashIterableDataset(BaseDataset[Iterable], IterableDataset):
+class FlashIterableDataset(IterableDataset, BaseDataset):
     """The ``IterableAutoDataset`` is a ``BaseDataset`` and a :class:`~torch.utils.data.IterableDataset`.
 
     The `data` argument must be an ``Iterable``.
     """
 
-    def load_data(self, data: Any) -> Any:
+    def load_data(self, data: Any) -> Iterable:
+        """By default, the `load_data` perform an identity operation.
+
+        Override to add your own logic to load the data.
+        """
         return data
 
     def load_sample(self, data: Any) -> Any:
+        """By default, the `load_sample` perform an identity operation.
+
+        Override to add your own logic to load a single sample.
+        """
         return data
 
     def resolve_functions(self):


Test patch: