# coding=utf-8

# Copyright 2023 The TensorFlow Datasets Authors.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# 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.

"""DatasetBuilder that stores data as a TFDS dataset.

Adhoc builders can be used to easily create a new TFDS dataset without having to

define a new `DatasetBuilder` class. This can be handy when working in a

notebook. For example, if you are in a notebook and you want to transform a

`tf.data.Dataset` into a TFDS dataset, then you can do so as follows:

```

import numpy as np

import tensorflow as tf

import tensorflow_datasets.public_api as tfds

my_ds_train = tf.data.Dataset.from_tensor_slices({"number": [1, 2, 3]})

my_ds_test = tf.data.Dataset.from_tensor_slices({"number": [4, 5]})

# Optionally define a custom `data_dir`. If None, then the default data dir is

# used.

custom_data_dir = "/my/folder"

builder = tfds.dataset_builders.store_as_tfds_dataset(

name="my_dataset",

config="single_number",

version="1.0.0",

data_dir=custom_data_dir,

split_datasets={

"train": my_ds_train,

"test": my_ds_test,

},

features=tfds.features.FeaturesDict({

"number": tfds.features.Scalar(dtype=np.int64),

}),

description="My dataset with a single number.",

release_notes={

"1.0.0": "Initial release with numbers up to 5!",

}

)

```

The `config` argument is optional and can be useful if you want to store

different configs under the same TFDS dataset.

The `data_dir` argument can be used to store the generated TFDS dataset in a

different folder, for example in your own sandbox if you don't want to share

this with others (yet). Note that when doing this, you also need to pass the

`data_dir` to `tfds.load`. If the `data_dir` argument is not specified, then

the default TFDS data dir will be used.

After the TFDS dataset has been stored, it can be loaded from other scripts:

```

# If no custom data dir was specified:

ds_test = tfds.load("my_dataset/single_number", split="test")

# When there are multiple versions, you can also specify the version.

ds_test = tfds.load("my_dataset/single_number:1.0.0", split="test")

# If the TFDS was stored in a custom folder, then it can be loaded as follows:

custom_data_dir = "/my/folder"

ds_test = tfds.load("my_dataset/single_number:1.0.0", split="test",

data_dir=custom_data_dir)

```

"""

from __future__ import annotations

import sys

import typing

from typing import Any, Dict, Iterable, Mapping, Optional, Union

from absl import logging

from etils import epath

from tensorflow_datasets.core import dataset_builder

from tensorflow_datasets.core import dataset_info

from tensorflow_datasets.core import dataset_utils

from tensorflow_datasets.core import download

from tensorflow_datasets.core import features as feature_lib

from tensorflow_datasets.core import file_adapters

from tensorflow_datasets.core import split_builder as split_builder_lib

from tensorflow_datasets.core import splits as splits_lib

from tensorflow_datasets.core import utils

from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf

KeyExample = split_builder_lib.KeyExample

if typing.TYPE_CHECKING:

import tensorflow as tf # pytype: disable=import-error

import apache_beam as beam # pytype: disable=import-error

BeamInput = Union[beam.PTransform, beam.PCollection[KeyExample]]

InputData = Union[

tf.data.Dataset,

beam.PTransform,

beam.PCollection[KeyExample],

Iterable[KeyExample],

]

else:

BeamInput = Union[

"beam.PTransform", "beam.PCollection[split_builder_lib.KeyExample]"

]

InputData = Union[

"tf.data.Dataset",

"beam.PTransform",

"beam.PCollection[split_builder_lib.KeyExample]",

"Iterable[KeyExample]",

]

class AdhocBuilder(

raise NotImplementedError()

def store_as_tfds_dataset(

return builder

class TfDataBuilder(AdhocBuilder, skip_registration=True):

super().__init__(split_datasets=split_datasets, **kwargs)