Source code for data_juicer.utils.cache_utils

import os
from functools import wraps

from datasets import disable_caching, enable_caching, is_caching_enabled

# Default cache location
DEFAULT_CACHE_HOME = '~/.cache'
CACHE_HOME = os.getenv('CACHE_HOME', DEFAULT_CACHE_HOME)

# Default data_juicer cache location
DEFAULT_DATA_JUICER_CACHE_HOME = os.path.join(CACHE_HOME, 'data_juicer')
DATA_JUICER_CACHE_HOME = os.path.expanduser(
    os.getenv('DATA_JUICER_CACHE_HOME', DEFAULT_DATA_JUICER_CACHE_HOME))

# Default assets cache location
DEFAULT_DATA_JUICER_ASSETS_CACHE = os.path.join(DATA_JUICER_CACHE_HOME,
                                                'assets')
DATA_JUICER_ASSETS_CACHE = os.getenv('DATA_JUICER_ASSETS_CACHE',
                                     DEFAULT_DATA_JUICER_ASSETS_CACHE)
# Default models cache location
DEFAULT_DATA_JUICER_MODELS_CACHE = os.path.join(DATA_JUICER_CACHE_HOME,
                                                'models')
DATA_JUICER_MODELS_CACHE = os.getenv('DATA_JUICER_MODELS_CACHE',
                                     DEFAULT_DATA_JUICER_MODELS_CACHE)

CACHE_COMPRESS = None




class DatasetCacheControl:
    """Define a range that change the cache state temporarily."""



    def __init__(self, on: bool = False):
        self.on = on


    def __enter__(self):
        """
        Record the original cache state and turn it to the target state.
        """
        self.previous_state = is_caching_enabled()
        if self.on:
            enable_caching()
        else:
            disable_caching()

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        Restore the original cache state.
        """
        if self.previous_state:
            enable_caching()
        else:
            disable_caching()





def dataset_cache_control(on):
    """
    A more easy-to-use decorator for functions that need to control the cache
    state temporarily.
    """

    def dataset_cache_decorator(func):

        @wraps(func)
        def wrapped_function(*args, **kwargs):
            with DatasetCacheControl(on=on):
                return func(*args, **kwargs)

        return wrapped_function

    return dataset_cache_decorator