from typing import List
from qasem_parser import QasemParser, QasemFrame
arg_parser_path = "cattana/flan-t5-large-qasem-joint-tokenized"
parser = QasemParser.from_pretrained(arg_parser_path)

sentences = [
    "The fox jumped over the fence.",
    "Back in May, a signal consistent with that of a radio beacon was detected in the area, but nothing turned up that helped with the search."
]
# frames is a list of lists, with one sublist per sentence such that len(frames) == len(sentences)  
# frames[i] is a sublist of the semantic frames that occur within sentence[i] 
frames: List[List[QasemFrame]] = parser(sentences)
print(frames[1][0])
# detect-v:  Back in May (when: when was something detected?) 
#            | a signal consistent with that of a radio beacon (R1: what was detected?) 
#            | in the area (where: where was something detected?) 

# The parser also respects original tokenization
# if the input is a batch of tokenized sentences  
pretokenized_sentences = [
    "Unfortunately , extensive property damage is bound to occur even with the best preparation .".split(),
    "Afghanistan to attend the summit after following the election in June , "
    "but the ongoing audit of votes has made this impossible .".split()
]
frames = parser(pretokenized_sentences)

for frames_per_sent in frames:
    # NOTE: frames_per_sent might be empty if no predicate 
    #       is detected in the sentence.
    for frame in frames_per_sent:
        print(frame)
    print()        
# 
# bind-v:  extensive property damage (R0: what is  bound  to do something?) 
#           | occur even with the best preparation (R1: what is something bound  to do?)
# occur-v:  extensive property damage (R0: what is occurring?)
#           | even with the best preparation (how: how is something occurring?)
# damage-n:  extensive property (R1: what is damaged?)
# prepare-n:  extensive property damage (R1: what is prepared?)
#
# call-v:  Plans (R0: what   called  for something?)
#           | the new President , or President-elect , of Afghanistan to attend the summit after following the election in June (R1: what did something call for?)
# attend-v:  the new President , or President-elect , of Afghanistan (R0: who attends something?) 
#           | the summit (R1: what does someone attend?) 
#           | after following the election in June (when: when does someone attend something?)
# follow-v: the election (R1: what was followed?) 
#           | in June (when: when was something followed?)
# make-v:  the ongoing audit of votes (R0: what made something?) | this impossible (R1: what did something make?)
# plan-n:  Plans (R0: what planned something?) 
#           | the new President , or President-elect , of Afghanistan to attend the summit after following the election in June , but the ongoing audit of votes has made this impossible (R1: what did something plan?)
# elect-n:  the new President , or President-elect , of Afghanistan (R1: what was elected?) 
#           | in June (when: when was something elected?)
# audit-n:  votes (R1: what was audited?)