GritLM/GritLM-7B · Assembling cache

I'd like to save the documents' caches to use them during generation. The following function extracts a document's cache from a batch of caches (as explained in the documentation):

def cache_by_id(caches, idx):
    cache = tuple(
        [
            (caches[i][0][idx : idx + 1], caches[i][1][idx : idx + 1])
            for i, c in enumerate(caches)
        ]
    )
    return cache

I'm not enough familiar with torch to be able to write the inverse function, which would assemble a cache from the top-k document caches. Can anybody help? Thanks!