def reduce_docs(docs, max_budget, offset, sort_docs=False):
    """Reduces the number of `docs` (subwords) to fit maximum budget by
    optionally sorting them by length, and then removing the shortest ones
    until `max_budget` is exceeded. Otherwise, will go left-to-right to preserve
    documents.

    Args:
         docs (list/array): tokenized documents.
         max_budget (int): the allowed budget in words/subwords.
         offset (int): additional offset for each doc.
         sort_docs (bool): if set, will sort documents by length and will
            attempt to remove the short ones.

    Returns:
        list of selected docs (preserved order).
    """
    if sort_docs:
        doc_lens = [len(d) for d in docs]
        order_indxs = sorted(range(len(doc_lens)), key=lambda k: doc_lens[k],
                             reverse=True)
    else:
        order_indxs = range(len(docs))
    curr_len = 0
    indx_coll = []
    for i, indx in enumerate(order_indxs):
        doc = docs[indx]
        if isinstance(doc, str):
            raise ValueError("Documents must be tokenized.")
        doc_len_w_offset = 0
        # adding the separator symbol offset to all but the first doc
        if i > 0:
            doc_len_w_offset += offset
        doc_len_w_offset += len(doc)
        if doc_len_w_offset <= max_budget - curr_len:
            curr_len += doc_len_w_offset
            indx_coll.append(indx)
        else:
            break
    indx_coll = sorted(indx_coll)
    docs = [docs[j] for j in indx_coll]
    return docs, indx_coll