We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Revisit TPU training
## 🚀 Feature

just curious, have you used TPU on Kaggle recently?
https://www.kaggle.com/code/jirkaborovec/demo-flash-image-classification-on-tpu

### Motivation

Kaggle has offered 20 hours of training which is mostly used with TF of Keras as any other lib is rather difficult, so there is huge potential to take this space by Flash...
Also With Flash, it shall be trivial to start your work with CPU for exploration and then alternate between GPU and TPU offering to finish with the best-trained model :rabbit: 

### Pitch

<!-- A clear and concise description of what you want to happen. -->

### Alternatives

Reason: https://pytorch-lightning.readthedocs.io/en/latest/accelerators/tpu_faq.html#unsupported-datatype-transfer-to-tpus

Unsupported data type transfer to TPUs?
```
File "/usr/local/lib/python3.8/dist-packages/torch_xla/utils/utils.py", line 205, in _for_each_instance_rewrite
    v = _for_each_instance_rewrite(result.__dict__[k], select_fn, fn, rwmap)
File "/usr/local/lib/python3.8/dist-packages/torch_xla/utils/utils.py", line 206, in _for_each_instance_rewrite
    result.__dict__[k] = v
TypeError: 'mappingproxy' object does not support item assignment
```

### Additional context

<!-- Add any other context or screenshots about the feature request here. -->


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flash/core/data/io/input.py`
Content:
```
1 # Copyright The PyTorch Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 import functools
15 import os
16 import sys
17 from copy import deepcopy
18 from typing import Any, cast, Dict, Iterable, List, Sequence, Tuple, Union
19 
20 from pytorch_lightning.utilities.enums import LightningEnum
21 from pytorch_lightning.utilities.exceptions import MisconfigurationException
22 from torch.utils.data import Dataset
23 
24 from flash.core.data.properties import Properties
25 from flash.core.data.utils import _STAGES_PREFIX
26 from flash.core.utilities.stages import RunningStage
27 
28 if sys.version_info < (3, 7):
29     from typing import GenericMeta
30 else:
31     GenericMeta = type
32 
33 
34 if not os.environ.get("READTHEDOCS", False):
35     from torch.utils.data import IterableDataset
36 else:
37     # ReadTheDocs mocks the `IterableDataset` import so it's type cannot be used as a base for a metaclass, so we
38     # replace it here.
39     IterableDataset = object
40 
41 
42 class InputFormat(LightningEnum):
43     """The ``InputFormat`` enum contains the data source names used by all of the default ``from_*`` methods in
44     :class:`~flash.core.data.data_module.DataModule`."""
45 
46     FOLDERS = "folders"
47     FILES = "files"
48     NUMPY = "numpy"
49     TENSORS = "tensors"
50     CSV = "csv"
51     JSON = "json"
52     PARQUET = "parquet"
53     DATASETS = "datasets"
54     HUGGINGFACE_DATASET = "hf_datasets"
55     FIFTYONE = "fiftyone"
56     DATAFRAME = "data_frame"
57     LISTS = "lists"
58     LABELSTUDIO = "labelstudio"
59 
60     # TODO: Create a FlashEnum class???
61     def __hash__(self) -> int:
62         return hash(self.value)
63 
64 
65 class DataKeys(LightningEnum):
66     """The ``DataKeys`` enum contains the keys that are used by built-in data sources to refer to inputs and
67     targets."""
68 
69     INPUT = "input"
70     PREDS = "preds"
71     TARGET = "target"
72     METADATA = "metadata"
73 
74     # TODO: Create a FlashEnum class???
75     def __hash__(self) -> int:
76         return hash(self.value)
77 
78 
79 class BaseDataFormat(LightningEnum):
80     """The base class for creating ``data_format`` for :class:`~flash.core.data.io.input.Input`."""
81 
82     def __hash__(self) -> int:
83         return hash(self.value)
84 
85 
86 def _has_len(data: Union[Sequence, Iterable]) -> bool:
87     """Duck typing check to see if the argument supports getting the length.
88 
89     Args:
90         data: The object to check for length support.
91     """
92     try:
93         len(data)
94         return True
95     except (TypeError, NotImplementedError):
96         return False
97 
98 
99 def _validate_input(input: "InputBase") -> None:
100     """Helper function to validate that the type of an ``InputBase.data`` is appropriate for the type of
101     ``InputBase`` being used.
102 
103     Args:
104         input: The ``InputBase`` instance to validate.
105 
106     Raises:
107         RuntimeError: If the ``input`` is of type ``Input`` and it's ``data`` attribute does not support ``len``.
108         RuntimeError: If the ``input`` is of type ``IterableInput`` and it's ``data`` attribute does support ``len``.
109     """
110     if input.data is not None:
111         if isinstance(input, Input) and not _has_len(input.data):
112             raise RuntimeError("`Input.data` is not a sequence with a defined length. Use `IterableInput` instead.")
113         elif isinstance(input, IterableInput) and _has_len(input.data):
114             raise RuntimeError("`IterableInput.data` is a sequence with a defined length. Use `Input` instead.")
115 
116 
117 def _wrap_init(class_dict: Dict[str, Any]) -> None:
118     """Helper function to wrap the ``__init__`` (if present) from a class construction dict to apply the
119     ``_validate_input`` function after instantiation. Modifies the dict inplace.
120 
121     Args:
122         class_dict: The class construction dict, optionally containing an init to wrap.
123     """
124     if "__init__" in class_dict:
125         fn = class_dict["__init__"]
126 
127         @functools.wraps(fn)
128         def wrapper(self, *args, **kwargs):
129             fn(self, *args, **kwargs)
130             _validate_input(self)
131 
132         class_dict["__init__"] = wrapper
133 
134 
135 class _InputMeta(GenericMeta):
136     """Metaclass for the ``InputBase`` which wraps any init defined in a subclass with the ``_validate_input``
137     helper."""
138 
139     def __new__(mcs, name: str, bases: Tuple, class_dict: Dict[str, Any]) -> "_InputMeta":
140         _wrap_init(class_dict)
141         return cast(_InputMeta, super().__new__(mcs, name, bases, class_dict))
142 
143 
144 class _IterableInputMeta(_InputMeta, type(IterableDataset)):
145     """Metaclass for the ``IterableInput`` which extends ``_InputMeta`` and avoids metaclass conflict with
146     ``IterableDataset``."""
147 
148     def __new__(mcs, name: str, bases: Tuple, class_dict: Dict[str, Any]) -> "_IterableInputMeta":
149         return cast(_IterableInputMeta, super().__new__(mcs, name, bases, class_dict))
150 
151 
152 class InputBase(Properties, metaclass=_InputMeta):
153     """``InputBase`` is the base class for the :class:`~flash.core.data.io.input.Input` and
154     :class:`~flash.core.data.io.input.IterableInput` dataset implementations in Flash. These datasets are
155     constructed via the ``load_data`` and ``load_sample`` hooks, which allow a single dataset object to include custom
156     loading logic according to the running stage (e.g. train, validate, test, predict).
157 
158     Args:
159         running_stage: The running stage for which the input will be used.
160         *args: Any arguments that are to be passed to the ``load_data`` hook.
161         **kwargs: Any additional keyword arguments to pass to the ``load_data`` hook.
162     """
163 
164     def __init__(self, running_stage: RunningStage, *args: Any, **kwargs: Any) -> None:
165 
166         super().__init__(running_stage=running_stage)
167 
168         self.data = None
169         if len(args) >= 1 and args[0] is not None:
170             self.data = getattr(self, f"{_STAGES_PREFIX[running_stage]}_load_data")(*args, **kwargs)
171 
172     def _call_load_sample(self, sample: Any) -> Any:
173         # Deepcopy the sample to avoid leaks with complex data structures
174         return getattr(self, f"{_STAGES_PREFIX[self.running_stage]}_load_sample")(deepcopy(sample))
175 
176     @staticmethod
177     def load_data(*args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:
178         """The ``load_data`` hook should return a collection of samples. To reduce the memory footprint, these
179         samples should typically not have been loaded. For example, an input which loads images from disk would
180         only return the list of filenames here rather than the loaded images.
181 
182         Args:
183             *args: Any arguments that the input requires.
184             **kwargs: Any additional keyword arguments that the input requires.
185         """
186         return args[0]
187 
188     def train_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:
189         """Override the ``train_load_data`` hook with data loading logic that is only required during training.
190 
191         Args:
192             *args: Any arguments that the input requires.
193             **kwargs: Any additional keyword arguments that the input requires.
194         """
195         return self.load_data(*args, **kwargs)
196 
197     def val_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:
198         """Override the ``val_load_data`` hook with data loading logic that is only required during validating.
199 
200         Args:
201             *args: Any arguments that the input requires.
202             **kwargs: Any additional keyword arguments that the input requires.
203         """
204         return self.load_data(*args, **kwargs)
205 
206     def test_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:
207         """Override the ``test_load_data`` hook with data loading logic that is only required during testing.
208 
209         Args:
210             *args: Any arguments that the input requires.
211             **kwargs: Any additional keyword arguments that the input requires.
212         """
213         return self.load_data(*args, **kwargs)
214 
215     def predict_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:
216         """Override the ``predict_load_data`` hook with data loading logic that is only required during predicting.
217 
218         Args:
219             *args: Any arguments that the input requires.
220             **kwargs: Any additional keyword arguments that the input requires.
221         """
222         return self.load_data(*args, **kwargs)
223 
224     @staticmethod
225     def load_sample(sample: Dict[str, Any]) -> Any:
226         """The ``load_sample`` hook is called for each ``__getitem__`` or ``__next__`` call to the dataset with a
227         single sample from the output of the ``load_data`` hook as input.
228 
229         Args:
230             sample: A single sample from the output of the ``load_data`` hook.
231         """
232         return sample
233 
234     def train_load_sample(self, sample: Dict[str, Any]) -> Any:
235         """Override the ``train_load_sample`` hook with data loading logic that is only required during training.
236 
237         Args:
238             sample: A single sample from the output of the ``load_data`` hook.
239         """
240         return self.load_sample(sample)
241 
242     def val_load_sample(self, sample: Dict[str, Any]) -> Any:
243         """Override the ``val_load_sample`` hook with data loading logic that is only required during validating.
244 
245         Args:
246             sample: A single sample from the output of the ``load_data`` hook.
247         """
248         return self.load_sample(sample)
249 
250     def test_load_sample(self, sample: Dict[str, Any]) -> Any:
251         """Override the ``test_load_sample`` hook with data loading logic that is only required during testing.
252 
253         Args:
254             sample: A single sample from the output of the ``load_data`` hook.
255         """
256         return self.load_sample(sample)
257 
258     def predict_load_sample(self, sample: Dict[str, Any]) -> Any:
259         """Override the ``predict_load_sample`` hook with data loading logic that is only required during
260         predicting.
261 
262         Args:
263             sample: A single sample from the output of the ``load_data`` hook.
264         """
265         return self.load_sample(sample)
266 
267     def __bool__(self):
268         """If ``self.data`` is ``None`` then the ``InputBase`` is considered falsey.
269 
270         This allows for quickly checking whether or not the ``InputBase`` is populated with data.
271         """
272         return self.data is not None
273 
274 
275 class Input(InputBase, Dataset):
276     def __getitem__(self, index: int) -> Any:
277         return self._call_load_sample(self.data[index])
278 
279     def __len__(self) -> int:
280         return len(self.data) if self.data is not None else 0
281 
282 
283 class IterableInput(InputBase, IterableDataset, metaclass=_IterableInputMeta):
284     def __iter__(self):
285         self.data_iter = iter(self.data)
286         return self
287 
288     def __next__(self) -> Any:
289         return self._call_load_sample(next(self.data_iter))
290 
291 
292 class ServeInput(Input):
293     def __init__(self) -> None:
294         if hasattr(self, "serve_load_data"):
295             raise MisconfigurationException("`serve_load_data` shouldn't be implemented.")
296 
297         super().__init__(RunningStage.SERVING)
298 
299     def serve_load_sample(self, sample: Any) -> List[Any]:
300         raise NotImplementedError
301 
302     def example_input(self) -> str:
303         raise NotImplementedError
304 
305     def __bool__(self):
306         return True
307
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/flash/core/data/io/input.py b/flash/core/data/io/input.py
--- a/flash/core/data/io/input.py
+++ b/flash/core/data/io/input.py
@@ -15,6 +15,7 @@
 import os
 import sys
 from copy import deepcopy
+from enum import Enum
 from typing import Any, cast, Dict, Iterable, List, Sequence, Tuple, Union
 
 from pytorch_lightning.utilities.enums import LightningEnum
@@ -171,7 +172,18 @@
 
     def _call_load_sample(self, sample: Any) -> Any:
         # Deepcopy the sample to avoid leaks with complex data structures
-        return getattr(self, f"{_STAGES_PREFIX[self.running_stage]}_load_sample")(deepcopy(sample))
+        sample_output = getattr(self, f"{_STAGES_PREFIX[self.running_stage]}_load_sample")(deepcopy(sample))
+
+        # Change DataKeys Enum to strings
+        if isinstance(sample_output, dict):
+            output_dict = {}
+            for key, val in sample_output.items():
+                if isinstance(key, Enum) and hasattr(key, "value"):
+                    output_dict[key.value] = val
+                else:
+                    output_dict[key] = val
+            return output_dict
+        return sample_output
 
     @staticmethod
     def load_data(*args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:

{"golden_diff": "diff --git a/flash/core/data/io/input.py b/flash/core/data/io/input.py\n--- a/flash/core/data/io/input.py\n+++ b/flash/core/data/io/input.py\n@@ -15,6 +15,7 @@\n import os\n import sys\n from copy import deepcopy\n+from enum import Enum\n from typing import Any, cast, Dict, Iterable, List, Sequence, Tuple, Union\n \n from pytorch_lightning.utilities.enums import LightningEnum\n@@ -171,7 +172,18 @@\n \n     def _call_load_sample(self, sample: Any) -> Any:\n         # Deepcopy the sample to avoid leaks with complex data structures\n-        return getattr(self, f\"{_STAGES_PREFIX[self.running_stage]}_load_sample\")(deepcopy(sample))\n+        sample_output = getattr(self, f\"{_STAGES_PREFIX[self.running_stage]}_load_sample\")(deepcopy(sample))\n+\n+        # Change DataKeys Enum to strings\n+        if isinstance(sample_output, dict):\n+            output_dict = {}\n+            for key, val in sample_output.items():\n+                if isinstance(key, Enum) and hasattr(key, \"value\"):\n+                    output_dict[key.value] = val\n+                else:\n+                    output_dict[key] = val\n+            return output_dict\n+        return sample_output\n \n     @staticmethod\n     def load_data(*args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n", "issue": "Revisit TPU training\n## \ud83d\ude80 Feature\r\n\r\njust curious, have you used TPU on Kaggle recently?\r\nhttps://www.kaggle.com/code/jirkaborovec/demo-flash-image-classification-on-tpu\r\n\r\n### Motivation\r\n\r\nKaggle has offered 20 hours of training which is mostly used with TF of Keras as any other lib is rather difficult, so there is huge potential to take this space by Flash...\r\nAlso With Flash, it shall be trivial to start your work with CPU for exploration and then alternate between GPU and TPU offering to finish with the best-trained model :rabbit: \r\n\r\n### Pitch\r\n\r\n<!-- A clear and concise description of what you want to happen. -->\r\n\r\n### Alternatives\r\n\r\nReason: https://pytorch-lightning.readthedocs.io/en/latest/accelerators/tpu_faq.html#unsupported-datatype-transfer-to-tpus\r\n\r\nUnsupported data type transfer to TPUs?\r\n```\r\nFile \"/usr/local/lib/python3.8/dist-packages/torch_xla/utils/utils.py\", line 205, in _for_each_instance_rewrite\r\n    v = _for_each_instance_rewrite(result.__dict__[k], select_fn, fn, rwmap)\r\nFile \"/usr/local/lib/python3.8/dist-packages/torch_xla/utils/utils.py\", line 206, in _for_each_instance_rewrite\r\n    result.__dict__[k] = v\r\nTypeError: 'mappingproxy' object does not support item assignment\r\n```\r\n\r\n### Additional context\r\n\r\n<!-- Add any other context or screenshots about the feature request here. -->\r\n\n", "before_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport functools\nimport os\nimport sys\nfrom copy import deepcopy\nfrom typing import Any, cast, Dict, Iterable, List, Sequence, Tuple, Union\n\nfrom pytorch_lightning.utilities.enums import LightningEnum\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch.utils.data import Dataset\n\nfrom flash.core.data.properties import Properties\nfrom flash.core.data.utils import _STAGES_PREFIX\nfrom flash.core.utilities.stages import RunningStage\n\nif sys.version_info < (3, 7):\n    from typing import GenericMeta\nelse:\n    GenericMeta = type\n\n\nif not os.environ.get(\"READTHEDOCS\", False):\n    from torch.utils.data import IterableDataset\nelse:\n    # ReadTheDocs mocks the `IterableDataset` import so it's type cannot be used as a base for a metaclass, so we\n    # replace it here.\n    IterableDataset = object\n\n\nclass InputFormat(LightningEnum):\n    \"\"\"The ``InputFormat`` enum contains the data source names used by all of the default ``from_*`` methods in\n    :class:`~flash.core.data.data_module.DataModule`.\"\"\"\n\n    FOLDERS = \"folders\"\n    FILES = \"files\"\n    NUMPY = \"numpy\"\n    TENSORS = \"tensors\"\n    CSV = \"csv\"\n    JSON = \"json\"\n    PARQUET = \"parquet\"\n    DATASETS = \"datasets\"\n    HUGGINGFACE_DATASET = \"hf_datasets\"\n    FIFTYONE = \"fiftyone\"\n    DATAFRAME = \"data_frame\"\n    LISTS = \"lists\"\n    LABELSTUDIO = \"labelstudio\"\n\n    # TODO: Create a FlashEnum class???\n    def __hash__(self) -> int:\n        return hash(self.value)\n\n\nclass DataKeys(LightningEnum):\n    \"\"\"The ``DataKeys`` enum contains the keys that are used by built-in data sources to refer to inputs and\n    targets.\"\"\"\n\n    INPUT = \"input\"\n    PREDS = \"preds\"\n    TARGET = \"target\"\n    METADATA = \"metadata\"\n\n    # TODO: Create a FlashEnum class???\n    def __hash__(self) -> int:\n        return hash(self.value)\n\n\nclass BaseDataFormat(LightningEnum):\n    \"\"\"The base class for creating ``data_format`` for :class:`~flash.core.data.io.input.Input`.\"\"\"\n\n    def __hash__(self) -> int:\n        return hash(self.value)\n\n\ndef _has_len(data: Union[Sequence, Iterable]) -> bool:\n    \"\"\"Duck typing check to see if the argument supports getting the length.\n\n    Args:\n        data: The object to check for length support.\n    \"\"\"\n    try:\n        len(data)\n        return True\n    except (TypeError, NotImplementedError):\n        return False\n\n\ndef _validate_input(input: \"InputBase\") -> None:\n    \"\"\"Helper function to validate that the type of an ``InputBase.data`` is appropriate for the type of\n    ``InputBase`` being used.\n\n    Args:\n        input: The ``InputBase`` instance to validate.\n\n    Raises:\n        RuntimeError: If the ``input`` is of type ``Input`` and it's ``data`` attribute does not support ``len``.\n        RuntimeError: If the ``input`` is of type ``IterableInput`` and it's ``data`` attribute does support ``len``.\n    \"\"\"\n    if input.data is not None:\n        if isinstance(input, Input) and not _has_len(input.data):\n            raise RuntimeError(\"`Input.data` is not a sequence with a defined length. Use `IterableInput` instead.\")\n        elif isinstance(input, IterableInput) and _has_len(input.data):\n            raise RuntimeError(\"`IterableInput.data` is a sequence with a defined length. Use `Input` instead.\")\n\n\ndef _wrap_init(class_dict: Dict[str, Any]) -> None:\n    \"\"\"Helper function to wrap the ``__init__`` (if present) from a class construction dict to apply the\n    ``_validate_input`` function after instantiation. Modifies the dict inplace.\n\n    Args:\n        class_dict: The class construction dict, optionally containing an init to wrap.\n    \"\"\"\n    if \"__init__\" in class_dict:\n        fn = class_dict[\"__init__\"]\n\n        @functools.wraps(fn)\n        def wrapper(self, *args, **kwargs):\n            fn(self, *args, **kwargs)\n            _validate_input(self)\n\n        class_dict[\"__init__\"] = wrapper\n\n\nclass _InputMeta(GenericMeta):\n    \"\"\"Metaclass for the ``InputBase`` which wraps any init defined in a subclass with the ``_validate_input``\n    helper.\"\"\"\n\n    def __new__(mcs, name: str, bases: Tuple, class_dict: Dict[str, Any]) -> \"_InputMeta\":\n        _wrap_init(class_dict)\n        return cast(_InputMeta, super().__new__(mcs, name, bases, class_dict))\n\n\nclass _IterableInputMeta(_InputMeta, type(IterableDataset)):\n    \"\"\"Metaclass for the ``IterableInput`` which extends ``_InputMeta`` and avoids metaclass conflict with\n    ``IterableDataset``.\"\"\"\n\n    def __new__(mcs, name: str, bases: Tuple, class_dict: Dict[str, Any]) -> \"_IterableInputMeta\":\n        return cast(_IterableInputMeta, super().__new__(mcs, name, bases, class_dict))\n\n\nclass InputBase(Properties, metaclass=_InputMeta):\n    \"\"\"``InputBase`` is the base class for the :class:`~flash.core.data.io.input.Input` and\n    :class:`~flash.core.data.io.input.IterableInput` dataset implementations in Flash. These datasets are\n    constructed via the ``load_data`` and ``load_sample`` hooks, which allow a single dataset object to include custom\n    loading logic according to the running stage (e.g. train, validate, test, predict).\n\n    Args:\n        running_stage: The running stage for which the input will be used.\n        *args: Any arguments that are to be passed to the ``load_data`` hook.\n        **kwargs: Any additional keyword arguments to pass to the ``load_data`` hook.\n    \"\"\"\n\n    def __init__(self, running_stage: RunningStage, *args: Any, **kwargs: Any) -> None:\n\n        super().__init__(running_stage=running_stage)\n\n        self.data = None\n        if len(args) >= 1 and args[0] is not None:\n            self.data = getattr(self, f\"{_STAGES_PREFIX[running_stage]}_load_data\")(*args, **kwargs)\n\n    def _call_load_sample(self, sample: Any) -> Any:\n        # Deepcopy the sample to avoid leaks with complex data structures\n        return getattr(self, f\"{_STAGES_PREFIX[self.running_stage]}_load_sample\")(deepcopy(sample))\n\n    @staticmethod\n    def load_data(*args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"The ``load_data`` hook should return a collection of samples. To reduce the memory footprint, these\n        samples should typically not have been loaded. For example, an input which loads images from disk would\n        only return the list of filenames here rather than the loaded images.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return args[0]\n\n    def train_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"Override the ``train_load_data`` hook with data loading logic that is only required during training.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return self.load_data(*args, **kwargs)\n\n    def val_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"Override the ``val_load_data`` hook with data loading logic that is only required during validating.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return self.load_data(*args, **kwargs)\n\n    def test_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"Override the ``test_load_data`` hook with data loading logic that is only required during testing.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return self.load_data(*args, **kwargs)\n\n    def predict_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"Override the ``predict_load_data`` hook with data loading logic that is only required during predicting.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return self.load_data(*args, **kwargs)\n\n    @staticmethod\n    def load_sample(sample: Dict[str, Any]) -> Any:\n        \"\"\"The ``load_sample`` hook is called for each ``__getitem__`` or ``__next__`` call to the dataset with a\n        single sample from the output of the ``load_data`` hook as input.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return sample\n\n    def train_load_sample(self, sample: Dict[str, Any]) -> Any:\n        \"\"\"Override the ``train_load_sample`` hook with data loading logic that is only required during training.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return self.load_sample(sample)\n\n    def val_load_sample(self, sample: Dict[str, Any]) -> Any:\n        \"\"\"Override the ``val_load_sample`` hook with data loading logic that is only required during validating.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return self.load_sample(sample)\n\n    def test_load_sample(self, sample: Dict[str, Any]) -> Any:\n        \"\"\"Override the ``test_load_sample`` hook with data loading logic that is only required during testing.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return self.load_sample(sample)\n\n    def predict_load_sample(self, sample: Dict[str, Any]) -> Any:\n        \"\"\"Override the ``predict_load_sample`` hook with data loading logic that is only required during\n        predicting.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return self.load_sample(sample)\n\n    def __bool__(self):\n        \"\"\"If ``self.data`` is ``None`` then the ``InputBase`` is considered falsey.\n\n        This allows for quickly checking whether or not the ``InputBase`` is populated with data.\n        \"\"\"\n        return self.data is not None\n\n\nclass Input(InputBase, Dataset):\n    def __getitem__(self, index: int) -> Any:\n        return self._call_load_sample(self.data[index])\n\n    def __len__(self) -> int:\n        return len(self.data) if self.data is not None else 0\n\n\nclass IterableInput(InputBase, IterableDataset, metaclass=_IterableInputMeta):\n    def __iter__(self):\n        self.data_iter = iter(self.data)\n        return self\n\n    def __next__(self) -> Any:\n        return self._call_load_sample(next(self.data_iter))\n\n\nclass ServeInput(Input):\n    def __init__(self) -> None:\n        if hasattr(self, \"serve_load_data\"):\n            raise MisconfigurationException(\"`serve_load_data` shouldn't be implemented.\")\n\n        super().__init__(RunningStage.SERVING)\n\n    def serve_load_sample(self, sample: Any) -> List[Any]:\n        raise NotImplementedError\n\n    def example_input(self) -> str:\n        raise NotImplementedError\n\n    def __bool__(self):\n        return True\n", "path": "flash/core/data/io/input.py"}], "after_files": [{"content": "# Copyright The PyTorch Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport functools\nimport os\nimport sys\nfrom copy import deepcopy\nfrom enum import Enum\nfrom typing import Any, cast, Dict, Iterable, List, Sequence, Tuple, Union\n\nfrom pytorch_lightning.utilities.enums import LightningEnum\nfrom pytorch_lightning.utilities.exceptions import MisconfigurationException\nfrom torch.utils.data import Dataset\n\nfrom flash.core.data.properties import Properties\nfrom flash.core.data.utils import _STAGES_PREFIX\nfrom flash.core.utilities.stages import RunningStage\n\nif sys.version_info < (3, 7):\n    from typing import GenericMeta\nelse:\n    GenericMeta = type\n\n\nif not os.environ.get(\"READTHEDOCS\", False):\n    from torch.utils.data import IterableDataset\nelse:\n    # ReadTheDocs mocks the `IterableDataset` import so it's type cannot be used as a base for a metaclass, so we\n    # replace it here.\n    IterableDataset = object\n\n\nclass InputFormat(LightningEnum):\n    \"\"\"The ``InputFormat`` enum contains the data source names used by all of the default ``from_*`` methods in\n    :class:`~flash.core.data.data_module.DataModule`.\"\"\"\n\n    FOLDERS = \"folders\"\n    FILES = \"files\"\n    NUMPY = \"numpy\"\n    TENSORS = \"tensors\"\n    CSV = \"csv\"\n    JSON = \"json\"\n    PARQUET = \"parquet\"\n    DATASETS = \"datasets\"\n    HUGGINGFACE_DATASET = \"hf_datasets\"\n    FIFTYONE = \"fiftyone\"\n    DATAFRAME = \"data_frame\"\n    LISTS = \"lists\"\n    LABELSTUDIO = \"labelstudio\"\n\n    # TODO: Create a FlashEnum class???\n    def __hash__(self) -> int:\n        return hash(self.value)\n\n\nclass DataKeys(LightningEnum):\n    \"\"\"The ``DataKeys`` enum contains the keys that are used by built-in data sources to refer to inputs and\n    targets.\"\"\"\n\n    INPUT = \"input\"\n    PREDS = \"preds\"\n    TARGET = \"target\"\n    METADATA = \"metadata\"\n\n    # TODO: Create a FlashEnum class???\n    def __hash__(self) -> int:\n        return hash(self.value)\n\n\nclass BaseDataFormat(LightningEnum):\n    \"\"\"The base class for creating ``data_format`` for :class:`~flash.core.data.io.input.Input`.\"\"\"\n\n    def __hash__(self) -> int:\n        return hash(self.value)\n\n\ndef _has_len(data: Union[Sequence, Iterable]) -> bool:\n    \"\"\"Duck typing check to see if the argument supports getting the length.\n\n    Args:\n        data: The object to check for length support.\n    \"\"\"\n    try:\n        len(data)\n        return True\n    except (TypeError, NotImplementedError):\n        return False\n\n\ndef _validate_input(input: \"InputBase\") -> None:\n    \"\"\"Helper function to validate that the type of an ``InputBase.data`` is appropriate for the type of\n    ``InputBase`` being used.\n\n    Args:\n        input: The ``InputBase`` instance to validate.\n\n    Raises:\n        RuntimeError: If the ``input`` is of type ``Input`` and it's ``data`` attribute does not support ``len``.\n        RuntimeError: If the ``input`` is of type ``IterableInput`` and it's ``data`` attribute does support ``len``.\n    \"\"\"\n    if input.data is not None:\n        if isinstance(input, Input) and not _has_len(input.data):\n            raise RuntimeError(\"`Input.data` is not a sequence with a defined length. Use `IterableInput` instead.\")\n        elif isinstance(input, IterableInput) and _has_len(input.data):\n            raise RuntimeError(\"`IterableInput.data` is a sequence with a defined length. Use `Input` instead.\")\n\n\ndef _wrap_init(class_dict: Dict[str, Any]) -> None:\n    \"\"\"Helper function to wrap the ``__init__`` (if present) from a class construction dict to apply the\n    ``_validate_input`` function after instantiation. Modifies the dict inplace.\n\n    Args:\n        class_dict: The class construction dict, optionally containing an init to wrap.\n    \"\"\"\n    if \"__init__\" in class_dict:\n        fn = class_dict[\"__init__\"]\n\n        @functools.wraps(fn)\n        def wrapper(self, *args, **kwargs):\n            fn(self, *args, **kwargs)\n            _validate_input(self)\n\n        class_dict[\"__init__\"] = wrapper\n\n\nclass _InputMeta(GenericMeta):\n    \"\"\"Metaclass for the ``InputBase`` which wraps any init defined in a subclass with the ``_validate_input``\n    helper.\"\"\"\n\n    def __new__(mcs, name: str, bases: Tuple, class_dict: Dict[str, Any]) -> \"_InputMeta\":\n        _wrap_init(class_dict)\n        return cast(_InputMeta, super().__new__(mcs, name, bases, class_dict))\n\n\nclass _IterableInputMeta(_InputMeta, type(IterableDataset)):\n    \"\"\"Metaclass for the ``IterableInput`` which extends ``_InputMeta`` and avoids metaclass conflict with\n    ``IterableDataset``.\"\"\"\n\n    def __new__(mcs, name: str, bases: Tuple, class_dict: Dict[str, Any]) -> \"_IterableInputMeta\":\n        return cast(_IterableInputMeta, super().__new__(mcs, name, bases, class_dict))\n\n\nclass InputBase(Properties, metaclass=_InputMeta):\n    \"\"\"``InputBase`` is the base class for the :class:`~flash.core.data.io.input.Input` and\n    :class:`~flash.core.data.io.input.IterableInput` dataset implementations in Flash. These datasets are\n    constructed via the ``load_data`` and ``load_sample`` hooks, which allow a single dataset object to include custom\n    loading logic according to the running stage (e.g. train, validate, test, predict).\n\n    Args:\n        running_stage: The running stage for which the input will be used.\n        *args: Any arguments that are to be passed to the ``load_data`` hook.\n        **kwargs: Any additional keyword arguments to pass to the ``load_data`` hook.\n    \"\"\"\n\n    def __init__(self, running_stage: RunningStage, *args: Any, **kwargs: Any) -> None:\n\n        super().__init__(running_stage=running_stage)\n\n        self.data = None\n        if len(args) >= 1 and args[0] is not None:\n            self.data = getattr(self, f\"{_STAGES_PREFIX[running_stage]}_load_data\")(*args, **kwargs)\n\n    def _call_load_sample(self, sample: Any) -> Any:\n        # Deepcopy the sample to avoid leaks with complex data structures\n        sample_output = getattr(self, f\"{_STAGES_PREFIX[self.running_stage]}_load_sample\")(deepcopy(sample))\n\n        # Change DataKeys Enum to strings\n        if isinstance(sample_output, dict):\n            output_dict = {}\n            for key, val in sample_output.items():\n                if isinstance(key, Enum) and hasattr(key, \"value\"):\n                    output_dict[key.value] = val\n                else:\n                    output_dict[key] = val\n            return output_dict\n        return sample_output\n\n    @staticmethod\n    def load_data(*args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"The ``load_data`` hook should return a collection of samples. To reduce the memory footprint, these\n        samples should typically not have been loaded. For example, an input which loads images from disk would\n        only return the list of filenames here rather than the loaded images.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return args[0]\n\n    def train_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"Override the ``train_load_data`` hook with data loading logic that is only required during training.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return self.load_data(*args, **kwargs)\n\n    def val_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"Override the ``val_load_data`` hook with data loading logic that is only required during validating.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return self.load_data(*args, **kwargs)\n\n    def test_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"Override the ``test_load_data`` hook with data loading logic that is only required during testing.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return self.load_data(*args, **kwargs)\n\n    def predict_load_data(self, *args: Any, **kwargs: Any) -> Union[Sequence, Iterable]:\n        \"\"\"Override the ``predict_load_data`` hook with data loading logic that is only required during predicting.\n\n        Args:\n            *args: Any arguments that the input requires.\n            **kwargs: Any additional keyword arguments that the input requires.\n        \"\"\"\n        return self.load_data(*args, **kwargs)\n\n    @staticmethod\n    def load_sample(sample: Dict[str, Any]) -> Any:\n        \"\"\"The ``load_sample`` hook is called for each ``__getitem__`` or ``__next__`` call to the dataset with a\n        single sample from the output of the ``load_data`` hook as input.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return sample\n\n    def train_load_sample(self, sample: Dict[str, Any]) -> Any:\n        \"\"\"Override the ``train_load_sample`` hook with data loading logic that is only required during training.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return self.load_sample(sample)\n\n    def val_load_sample(self, sample: Dict[str, Any]) -> Any:\n        \"\"\"Override the ``val_load_sample`` hook with data loading logic that is only required during validating.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return self.load_sample(sample)\n\n    def test_load_sample(self, sample: Dict[str, Any]) -> Any:\n        \"\"\"Override the ``test_load_sample`` hook with data loading logic that is only required during testing.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return self.load_sample(sample)\n\n    def predict_load_sample(self, sample: Dict[str, Any]) -> Any:\n        \"\"\"Override the ``predict_load_sample`` hook with data loading logic that is only required during\n        predicting.\n\n        Args:\n            sample: A single sample from the output of the ``load_data`` hook.\n        \"\"\"\n        return self.load_sample(sample)\n\n    def __bool__(self):\n        \"\"\"If ``self.data`` is ``None`` then the ``InputBase`` is considered falsey.\n\n        This allows for quickly checking whether or not the ``InputBase`` is populated with data.\n        \"\"\"\n        return self.data is not None\n\n\nclass Input(InputBase, Dataset):\n    def __getitem__(self, index: int) -> Any:\n        return self._call_load_sample(self.data[index])\n\n    def __len__(self) -> int:\n        return len(self.data) if self.data is not None else 0\n\n\nclass IterableInput(InputBase, IterableDataset, metaclass=_IterableInputMeta):\n    def __iter__(self):\n        self.data_iter = iter(self.data)\n        return self\n\n    def __next__(self) -> Any:\n        return self._call_load_sample(next(self.data_iter))\n\n\nclass ServeInput(Input):\n    def __init__(self) -> None:\n        if hasattr(self, \"serve_load_data\"):\n            raise MisconfigurationException(\"`serve_load_data` shouldn't be implemented.\")\n\n        super().__init__(RunningStage.SERVING)\n\n    def serve_load_sample(self, sample: Any) -> List[Any]:\n        raise NotImplementedError\n\n    def example_input(self) -> str:\n        raise NotImplementedError\n\n    def __bool__(self):\n        return True\n", "path": "flash/core/data/io/input.py"}]}