Package: tfdsio
TensorFlow Datasets IO (tfdsio)Dynamic TensorFlow Datasets with Pytorch Support + MoreFeaturestfdsioallows for creation of tensorflow_datasets dynamically defined by a config (json/dict)
without requiring writing custom python classes, which can lead to library bloat. Additionally,
you can also define a preprocessor to handle the data transformation prior to being written as a tfrecord.tfdsioenables you to create multiple version-controlled dataset variations, so that the final tfrecords
contain only what you need to train your models, ensuring a more efficient data pipeline.Single function call to return a custom tensorflow_dataset/tf.data.Dataset object (or as numpy iterator/pandas dataframe)Efficiently readtext,csv,jsonlfiles.Support for loadingtfdsiodatasets in other formats through Custom Adapterstext-to-text-transformerstorch(WIP)Installationtfdsiois available onpypiand can be
installed by running:# Install from command line$pipinstalltfdsio# Install from source$pipinstall--upgradehttps://github.com/trisongz/tfdsioTFDSIO Dataset ConfigurationThe full config spec is shown below. The values below will be overwritten after loading
the dataset configuration passed.name:Optional[str]=None# Dataset Name i.e. 'my_dataset'classifier:Optional[str]=None# Dataset variation, such as 'lm' or 'qa'version:Optional[VersionOrStr]='1.0.0'# Dataset version, i.e. '0.0.1'release_notes:Optional[Dict[str,str]]=None# Release notes included in the dataset metadatasupported_versions:List[str]=None# Optional list of versions ['0.0.1', '0.0.2']description:Optional[str]=None# Description of dataset included in dataset metadatadataset_urls:Optional[Any]=None# Defines your dataset urls, expected to be a dictdataset_format:Optional[str]='jsonlines'# ['jsonlines', 'text', 'csv']features:Optional[Any]=None# Maps the dataset output dict keys to tf.features, supports ['text', 'audio', 'image']datamap:Optional[Any]=None# Maps your dataset input dict keys to the dataset output keyssupervised_keys:Optional[Any]=None# Optionalhomepage:Optional[str]=''# homepage for datasetcitation:Optional[Any]=None# citation for datasetmetadata:Optional[Any]=None# metadata for datasetredistribution_info:Optional[Any]=None# redistribution info for datasetdata_dir:Optional[str]=None# [IMPORTANT]: This should be your GCS Directory or local drive that stores your dataset.process_args:Optional[Any]=None# Args passed to your preprocessor functionExamplestfdsioaims to make it simple to turn your custom dataset into a training ready data pipeline.
Built intotfdsioinclude useful utilities to allow working with GCS storage/Files much easierMinimal Examplefromtfdsioimporttfds_dataset,tfds_sample,set_adcfromtfdsioimporttfds# import tensorflow_datasets already initialized# Remember if you are reading from a private bucket to ensure ADC is setset_adc('/path/to/adc.json')dataset_config={'name':'my_dataset','classifier':'qa','version':'1.0.0','features':{'input_text':'text','target_text':'text'},'datamap':{'question':'input_text','text':'target_text',},'dataset_urls':{'train':'gs://your-storage-bucket/datasets/custom_dataset.jsonl'},'dataset_format':'jsonlines','homepage':'https://growthengineai.com','data_dir':'gs://your-storage-bucket/datasets/cached','description':'My Custom Question Answering Dataset'}# As long as the above configuration matches, the next time it's called, it will load from pre-built datasetdataset=tfds_dataset(dataset_config,preprocessor=None,build=True,as_tfdataset=True,as_numpy=False,as_df=False)# If it wasn't already created with as_numpy=Truesamples=tfds_sample(dataset,num_samples=5,return_samples=True)# or the standard methodforexintfds.as_numpy(dataset.take(5)):print(ex)Create Dataset Variations EasilyUsing the same config as above, you can modify your config and define a different feature/datamap
and have a different dataset variation.A dataset's identifier is comprised of(<dataset_name>/<dataset_classifier>/<dataset_version>),
correlating toname,classifier, andversion.dataset_config_2=dataset_config.copy()dataset_config_2['classifier']='lm'dataset_config_2['datamap']={'context':'input_text','answer':'target_text'}dataset2=tfds_dataset(dataset_config_2,preprocessor=None,build=True,as_tfdataset=True,as_numpy=False,as_df=False)# Your new dataset variationsamples=tfds_sample(dataset2,num_samples=5,return_samples=True)Remove a DatasetIf you something messed up during the dataset creation process, you can usetfds_removeto remove the directoryfromtfdsioimporttfds_removetfds_remove(dataset_config,prompt=True)# prompt = False will delete without askingUsing PreprocessorsPreprocessors can be any function, but should expect the minimum of the following args
kwargs will contain any args that were passed from config.process_args set in the original
configdefpreprocessor(idx:int,data:dict,extra:Optional[Filepath]=None,**kwargs):# if return_data is a list, will create an example per item in listreturn_data=[]# do stuffreturnreturn_data# if return_data is a dict, only one example will be createdreturn_data={}# do stuffreturnreturn_datadataset=tfds_dataset(dataset_config,preprocessor=preprocessor)Using AdaptersCurrently, tfdsio has support fortext-to-text-transformersort5.importtensorflowastfimportt5fromt5importseqiofromtfdsioimporttfds,tfds_samplefromtfdsio.adapterimportT5DataSourcevocab='/path/to/vocab/sentencepiece.model'seqio.TaskRegistry.add("my_dataset",source=T5DataSource(config_or_file=dataset_config,splits={"train":"train[:90%]","validation":"train[90%:]","test":"validation"}),preprocessors=[seqio.preprocessors.tokenize,seqio.preprocessors.append_eos,],output_features={"inputs":seqio.Feature(seqio.SentencePieceVocabulary(vocab),add_eos=False,dtype=tf.int32),"targets":seqio.Feature(seqio.SentencePieceVocabulary(vocab),add_eos=True,dtype=tf.int32),},metric_fns=[])my_task=t5.data.TaskRegistry.get("my_dataset")ds=my_task.get_dataset(split="validation",sequence_length={"inputs":128,"targets":128})print("A few preprocessed validation examples...")samples=tfds_sample(ds,num_samples=5,return_samples=True)LimitationsWhiletfdsiohas many useful utilities that allows extendingtensorflow_datasetsbeyond the
base library, there are still challenges, including not being able to simply usetfds.load()for any custom dataset, which can be limiting when working with many datasets in a single run.
However, taking fromt5, you can develop your own adapters to get around this limitation.Contributions are welcome!Motivation & AboutI have worked extensively with all dataset formats and came to really like tfds when working with TPUs
in TensorFlow. However, I found that continually expanding the dataset came with lots of bloat and time
to dig into each dataset to fix a bug, rather than having a change that affected all the datasets to be
updated.The long term roadmap is to expandtfdsioto enable cross-compatibility with any framework, supporting all
major data types, backed with the high-performanttf.data.Datasetbackend.I lead the ML Team atGrowth Engine AIworking with large scale NLP models in EdTech.
If you find working on NLP and MLOps challenges exciting, and would want to join our team, shoot me an
email: ts at growthengineai.comAcknowledgementsDevelopment oftfdsiorelied on contributions made from the following projects,
and I recommend checking them out as well!tensorflow-datasetstext-to-text-transformersrwightmanTPU Research Cloud