utils.py

"""⭐ Text Classification with Optimum and ONNXRuntime

Utils functions.

Author:
    - @ChainYo - https://github.com/ChainYo
"""

from contextlib import contextmanager
from time import time
from typing import List


@contextmanager
def calculate_inference_time(buffer: List[int]):
    """Calculate inference time.

    Args:
        buffer (list): List of inference times.

    Returns:
        float: Average inference time.
    """
    start = time()
    yield
    end = time()
    buffer.append(end - start)