from clean import clean_text

from datasets import load_dataset

dataset_v0 = load_dataset('oscar', "unshuffled_deduplicated_nl", split='train', streaming=True)


def f(obj):
    obj["text"] = clean_text(obj["text"])
    return obj


dataset_v1 = dataset_v0.map(f)
dataset_v2 = dataset_v1.filter(lambda obj: obj['text'] is not None)

it = iter(dataset_v0)
print(next(it))
print(next(it))
print(next(it))

it = iter(dataset_v1)
print(next(it))
print(next(it))
print(next(it))

it = iter(dataset_v2)
print(next(it))
print(next(it))
print(next(it))