Function _text_ds

axlearn/common/input_tf_data_test.py:298–310

(texts: list[str], *, repeat=1)

def _text_ds(texts: list[str], *, repeat=1) -> tf.data.Dataset:
    dataset = {
        "text": [],
        "index": [],
        "is_valid": [],
    }

    for index, text in enumerate(texts):
        dataset["index"].append(tf.constant(index, dtype=tf.int32))
        dataset["text"].append(tf.constant(text, dtype=tf.string))
        dataset["is_valid"].append(tf.constant(True, dtype=tf.bool))

    return tf.data.Dataset.from_tensor_slices(dataset).repeat(repeat)


def _tokens_ds(tokens: list[list[int]], *, repeat=1) -> tf.data.Dataset: