
import torch
import argparse
from collections import defaultdict
import os
import json


def load_predictions(input_path):
	pred_list = []
	for file_name in os.listdir(input_path):
		if file_name.endswith('.pt'):
			preds = torch.load(os.path.join(input_path, file_name))
			pred_list.extend(preds)

	question_scores = defaultdict(lambda: defaultdict(dict))
	p_count = 0
	u_count = 0
	for prediction in pred_list:
		doc_pass_id = prediction['id']
		q_p_id = prediction['question_id']
		# score = prediction['pos_score']
		score = prediction['pos_score'] - prediction['neg_score']
		if doc_pass_id not in question_scores or q_p_id not in question_scores[doc_pass_id]:
			p_count += 1
		u_count += 1
		question_scores[doc_pass_id][q_p_id] = score
	print(f'{p_count} unique predictions')
	print(f'{u_count} total predictions')
	return question_scores


def main():
	parser = argparse.ArgumentParser()
	parser.add_argument('-i', '--input_path', required=True)
	parser.add_argument('-o', '--output_path', required=True)
	args = parser.parse_args()

	input_path = args.input_path
	output_path = args.output_path

	question_scores = load_predictions(input_path)
	with open(output_path, 'w') as f:
		json.dump(question_scores, f)


if __name__ == '__main__':
	main()

[((11, 18, 11, 40), 'os.listdir', 'os.listdir', ({(11, 29, 11, 39): 'input_path'}, {}), '(input_path)', False, 'import os\n'), ((34, 10, 34, 35), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((44, 2, 44, 31), 'json.dump', 'json.dump', ({(44, 12, 44, 27): 'question_scores', (44, 29, 44, 30): 'f'}, {}), '(question_scores, f)', False, 'import json\n'), ((16, 39, 16, 56), 'collections.defaultdict', 'defaultdict', ({(16, 51, 16, 55): 'dict'}, {}), '(dict)', False, 'from collections import defaultdict\n'), ((13, 22, 13, 57), 'os.path.join', 'os.path.join', ({(13, 35, 13, 45): 'input_path', (13, 47, 13, 56): 'file_name'}, {}), '(input_path, file_name)', False, 'import os\n')]