src/utils/preprocessing.py

from typing import Dict, List


def result_to_reader_input(result: Dict[str, List[str]]) \
        -> Dict[str, List[str]]:
    """Takes the output of the retriever and turns it into a format the reader
    understands.

    Args:
        result (Dict[str, List[str]]): The result from the retriever
    """

    # Take the number of valeus of an arbitrary item as the number of entries
    # (This should always be valid)
    num_entries = len(result['n_chapter'])

    # Prepare result
    reader_result = {
        'titles': [],
        'texts': []
    }

    for n in range(num_entries):
        # Get the most specific title
        if result['subsection'][n] != 'nan':
            title = result['subsection'][n]
        elif result['section'][n] != 'nan':
            title = result['section'][n]
        else:
            title = result['chapter'][n]

        reader_result['titles'].append(title)
        reader_result['texts'].append(result['text'][n])

    return reader_result