UCDP-AEC

Models and data for the UCDP Abstractive Event analysis Corpus paper
git clone https://esimon.eu/repos/UCDP-AEC.git
Log | Files | Refs | README | LICENSE

label_tree.py (2410B)

      1 from typing import Dict
      2 
      3 from transformers import AutoTokenizer
      4 
      5 
      6 def list_dictionary(d, n_tab=-1):
      7     if isinstance(d, list):
      8         for i in d:
      9             list_dictionary(i, n_tab)
     10     elif isinstance(d, dict):
     11         n_tab += 1
     12         for key, value in d.items():
     13             if key == '<end>':
     14                 print("{}{}".format(" " * n_tab, key))
     15             else:
     16                 print("{}{}".format(" " * n_tab, key))
     17                 list_dictionary(value, n_tab)
     18     else:
     19         print("{}{}".format("\t" * n_tab, d))
     20 
     21 
     22 def print_tree(tree):
     23     list_dictionary(tree)
     24 
     25 
     26 def get_label_name_tree(label_name_list, tokenizer, end_symbol='<end>'):
     27     sub_token_tree = dict()
     28 
     29     label_tree = dict()
     30     for typename in label_name_list:
     31         after_tokenized = tokenizer.encode(typename, add_special_tokens=False)
     32         label_tree[typename] = after_tokenized
     33 
     34     for _, sub_label_seq in label_tree.items():
     35         parent = sub_token_tree
     36         for value in sub_label_seq:
     37             if value not in parent:
     38                 parent[value] = dict()
     39             parent = parent[value]
     40 
     41         parent[end_symbol] = None
     42 
     43     return sub_token_tree
     44 
     45 
     46 class PrefixTree:
     47     def __init__(self, label_name_list, tokenizer, end_symbol='<end>'):
     48         self.label_name_list = label_name_list
     49         self._tokenizer = tokenizer
     50         self.label_name_tree = get_label_name_tree(
     51             label_name_list, tokenizer, end_symbol)
     52         self._end_symbol = end_symbol
     53 
     54     def is_end_of_tree(self, tree: Dict):
     55         return len(tree) == 1 and self._end_symbol in tree
     56 
     57 
     58 if __name__ == "__main__":
     59     event_subtype_name = [line.strip() for line in """Die
     60     Marry
     61     Divorce
     62     Injure
     63     Transfer-Ownership
     64     Transfer-Money
     65     Transport
     66     Start-Org
     67     Be-Born
     68     End-Org
     69     Declare-Bankruptcy
     70     Merge-Org
     71     Attack
     72     Demonstrate
     73     Meet
     74     Phone-Write
     75     Start-Position
     76     End-Position
     77     Nominate
     78     Elect
     79     Arrest-Jail
     80     Release-Parole
     81     Charge-Indict
     82     Trial-Hearing
     83     Sue
     84     Convict
     85     Sentence
     86     Fine
     87     Execute
     88     Extradite
     89     Acquit
     90     Pardon
     91     Appeal""".split('\n')]
     92 
     93     test_tokenizer = AutoTokenizer.from_pretrained('t5-small')
     94 
     95     suptype_tree = get_label_name_tree(
     96         event_suptype_name, test_tokenizer)
     97     # role_tree = get_label_name_tree(ACEEventMetaData.event_role_name)
     98     print_tree(suptype_tree)
     99     # print_tree(role_tree)