File size: 989 Bytes
ab5dfc2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
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
|