_enforce_nested_string_type() takes too much time for Sequence(Value(...))

[NouamaneTazi]

As you can see in the figure below, most of time spent adding a batch is consumed by _enforce_nested_string_type(). Is that necessary?

import logging

import torch
from datasets import Features, Sequence, Value, Array2D

from evaluate.module import EvaluationModule, EvaluationModuleInfo
import time

logging.basicConfig(level=logging.INFO)


class DummyMetric(EvaluationModule):
    def _info(self):
        return EvaluationModuleInfo(
            description="dummy metric for tests",
            citation="insert citation here",
            features=Features(
                {
                    "tensor": Sequence(Value("int64")),
                }
            ),
        )


metric = DummyMetric()
start_time = time.time()
metric.add_batch(
    tensor=torch.randint(0, 10, (1000, 7000)),
)
print(f"time={round(time.time()-start_time, 2)}s") # outputs time=32.41s

Related to: #33

[NouamaneTazi]

For my use case, using Array2D yielded a much better performance

import logging

import torch
from datasets import Features, Sequence, Value, Array2D

from evaluate.module import EvaluationModule, EvaluationModuleInfo
import time

logging.basicConfig(level=logging.INFO)


class DummyMetric(EvaluationModule):
    def _info(self):
        return EvaluationModuleInfo(
            description="dummy metric for tests",
            citation="insert citation here",
            features=Features(
                {
                    "tensor": Array2D((None, 7000), "int64"),
                }
            ),
        )


metric = DummyMetric()
start_time = time.time()
metric.add_batch(
    tensor=torch.randint(0, 10, (1, 1000, 7000)),
)
print(f"time={round(time.time()-start_time, 2)}s") # outputs time=0.38s

[lvwerra]

Indeed, that is very slow. We could probably just check it for the first element of the batch instead of all, but as your analysis points out even if the _enforce_nested_string_type is much faster you still spend a significant amount of time on encode_batch.