Source code for data_juicer.format.empty_formatter
from typing import List
import pandas as pd
from datasets import Dataset, Features, Value
from data_juicer.utils.lazy_loader import LazyLoader
from .formatter import FORMATTERS, BaseFormatter
ray = LazyLoader('ray', 'ray')
[docs]@FORMATTERS.register_module()
class EmptyFormatter(BaseFormatter):
"""
The class is used to create empty data.
"""
SUFFIXES = []
[docs] def __init__(self, length, feature_keys: List[str] = [], *args, **kwargs):
"""
Initialization method.
:param length: The empty dataset length.
:param feature_keys: feature key name list.
"""
self.length = length
self.feature_keys = feature_keys
if isinstance(self.feature_keys, str):
self.feature_keys = [self.feature_keys]
@property
def null_value(self):
return None
[docs] def load_dataset(self, *args, **kwargs):
data_dict = {}
features = Features()
for key in self.feature_keys:
features.update({key: Value('string')})
data_dict.update(
{key: [self.null_value for _ in range(self.length)]})
empty_dataset = Dataset.from_dict(data_dict, features=features)
from data_juicer.core.data import NestedDataset
empty_dataset = NestedDataset(empty_dataset)
return empty_dataset
[docs]@FORMATTERS.register_module()
class RayEmptyFormatter(BaseFormatter):
"""
The class is used to create empty data for ray.
"""
SUFFIXES = []
[docs] def __init__(self, length, feature_keys: List[str] = [], *args, **kwargs):
"""
Initialization method.
:param length: The empty dataset length.
:param feature_keys: feature key name list.
"""
self.length = length
self.feature_keys = feature_keys
if isinstance(self.feature_keys, str):
self.feature_keys = [self.feature_keys]
@property
def null_value(self):
return {}
[docs] def load_dataset(self, *args, **kwargs):
if len(self.feature_keys):
df = pd.DataFrame({
col: [self.null_value for _ in range(self.length)]
for col in self.feature_keys
})
else:
df = pd.DataFrame([self.null_value for _ in range(self.length)])
empty_dataset = ray.data.from_pandas(df)
return empty_dataset