import torch


def _get_data_sorce_start_position_and_cache_sizes(
    cache, embedding_key, sorted_data_sources
):
    """Gets the first index and size per data sources in the concatenated data."""
    curr_position = torch.tensor(0, dtype=torch.int64).detach()
    start_positions = {}
    cache_sizes = {}
    for data_source in sorted_data_sources:
        start_positions[data_source] = curr_position
        cache_sizes[data_source] = cache[data_source].data[embedding_key].size()[0]
        curr_position = curr_position + cache_sizes[data_source]
    return start_positions, cache_sizes