maximize_i_j.py

text = r"""
George Washington (February 22, 1732[b] – December 14, 1799) was an American political leader, military general, statesman, and Founding Father who served as the first president of the United States from 1789 to 1797. Previously, he led Patriot forces to victory in the nation's War for Independence. He presided at the Constitutional Convention of 1787, which established the U.S. Constitution and a federal government. Washington has been called the "Father of His Country" for his manifold leadership in the formative days of the new nation.
"""

question = "Who was the first president?"

inputs = tokenizer.encode_plus(question, text, add_special_tokens=True, return_tensors="pt")
input_ids = inputs["input_ids"].tolist()[0]

text_tokens = tokenizer.convert_ids_to_tokens(input_ids)

pred = model(**inputs)
answer_start_scores, answer_end_scores = pred['start_logits'][0] ,pred['end_logits'][0] 

#get the index of first highest 20 tokens only
start_indexes = list(np.argsort(answer_start_scores.detach().numpy()))[::-1][:20]
end_indexes = list(np.argsort(answer_end_scores.detach().numpy()))[::-1][:20]

valid_answers = []
for start_index in start_indexes:
    for end_index in end_indexes:
        if start_index <= end_index:
            valid_answers.append(
                {
                    "score": answer_start_scores[start_index] + answer_end_scores[end_index],
                    "text": tokenizer.convert_tokens_to_string(
                        tokenizer.convert_ids_to_tokens(input_ids[start_index:end_index+1]))
                }
            )
    
valid_answers = sorted(valid_answers, key=lambda x: x["score"], reverse=True)