<commit_before>#!/usr/bin/env python
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# -*- coding: utf-8 -*-

import itertools
import time
import tensorflow as tf

from benchmark_autotuner import auto_tf_func_tuner
from benchmark_logger import logging
from benchmark_utils import force_gpu_resync


def SyntheticDataset(dataset, device):
    data_batch = next(iter(dataset))

    def copy_on_device(t):

        if t.dtype != tf.int32:
            with tf.device(device):
                return tf.identity(t)

        return t

    if isinstance(data_batch, (tuple, list)):
        data_batch = [copy_on_device(t) for t in data_batch]
    elif isinstance(data_batch, dict):
        data_batch = {k: copy_on_device(t) for k, t in data_batch.items()}
    else:
        data_batch = copy_on_device(data_batch)

    return itertools.repeat(data_batch)


def ensure_dataset_on_gpu(dataset, device):

    # ensuring no tensor dtype == int32
    input_batch = next(iter(dataset))
    if isinstance(input_batch, dict):
        input_batch = input_batch.values()
    elif not isinstance(input_batch, (tuple, list)):
        input_batch = [input_batch]

    if any([t.dtype == tf.int32 for t in input_batch]):
        logging.warning(<commit_msg>[Benchmarking-Py] Fix Synthetic Dataset with DS returning `x, y`<commit_after>#!/usr/bin/env python
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# -*- coding: utf-8 -*-

import itertools
import time
import tensorflow as tf

from benchmark_autotuner import auto_tf_func_tuner
from benchmark_logger import logging
from benchmark_utils import force_gpu_resync


def SyntheticDataset(dataset, device):
    data_batch = next(iter(dataset))

    def copy_on_device(data):
        if isinstance(data, (tuple, list)):
            return [copy_on_device(t) for t in data]
        elif isinstance(data, dict):
            return {k: copy_on_device(t) for k, t in data.items()}
        else:
            try:
                if data.dtype != tf.int32:
                    with tf.device(device):
                        return tf.identity(data)
            except AttributeError:
                pass

            return data

    data_batch = copy_on_device(data_batch)

    return itertools.repeat(data_batch)


def ensure_dataset_on_gpu(dataset, device):

    # ensuring no tensor dtype == int32
    input_batch = next(iter(dataset))
    if isinstance(input_batch, dict):
        input_batch = input_batch.values()
    elif not isinstance(input_batch, (tuple, list)):
        input_batch = [input_batch]

    if any([t.dtype == tf.int32 for t in input_batch]):
        logging.warning(
{'fraction_non_alphanumeric': 0.09352787130564909, 'fraction_numerical': 0.00823045267489712, 'mean_word_length': 2.973254086181278, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2486382', 'n_tokens_mistral': 853, 'n_tokens_neox': 809, 'n_words': 221}