<filename>src/python/serif/embeddings/get_pretrained_embeddings.py
import argparse
import codecs
import json
import os
import sys
import logging

import numpy as np

import serifxml3
from serif.theory.mention import Mention
from serif.theory.syn_node import SynNode
from serif.theory.value_mention import ValueMention


def get_docid_to_embeddings_from_dir(embeddings_dir):
    embeddings_map = {}
    if embeddings_dir is not None:
        for npz in os.listdir(embeddings_dir):
            if npz.endswith('.npz'):
                docid = npz[:-len('.npz')]
                embeddings_map[docid] = os.path.join(embeddings_dir, npz)
    return embeddings_map


def get_docid_to_embeddings_from_list(embeddings_filelist):
    embeddings_map = {}
    with open(embeddings_filelist, 'r') as fh:
        for line in fh:
            line = line.strip()
            docid = line[:-len('.npz')].rsplit('/')[-1]
            embeddings_map[docid] = line
    return embeddings_map


def load_embeddings(docid, npz_files):
    """
    Loads an npz file containing embeddings generated by NlpLingo.
    The output of this function has the following format:
    {
        "triggers": {
            <trigger_extractor_name>: {

                # key for triggers:
                (int(<start_offset>), int(<end_offset>)): (

                    int(<start_offset>),
                    int(<end_offset>),
                    ndarray(<trigger_embedding>)), [...]

            }, [...]
        },

        "arguments": {
            <argument_extractor_name>: {

                # key for arguments:
                (int(<trigger_start_offset>),
                 int(<trigger_end_offset>),
                 int(<argument_start_offset>),
                 int(<argument_end_offset>)): (

                    int(<trigger_start_offset>),
                    int(<trigger_end_offset>),
                    int(<argument_start_offset>),
                    int(<argument_end_offset>),
                    ndarray(<argument_embedding>)), [...]

            }, [...]
        },

        "bert": {
            # for triggers, the trigger key above:
            (int(<start_offset>), int(<end_offset>)): (

                int(<sentence_id>),
                int(<token_id>)),

            # for arguments, the argument key above:
            (int(<trigger_start_offset>),
             int(<trigger_end_offset>),
             int(<argument_start_offset>),
             int(<argument_end_offset>)): (

                int(<argument_start_offset>),
                int(<argument_end_offset>)), [...]

        }
    }

    :param docid: a docid
    :param npz_files: a docid2filepath dict for NlpLingo npz output files
    :return: dict as described above
    """
    npz_file = npz_files.get(docid)
    if npz_file is None:
        return None
    else:
        ret = {}
        np_object = np.load(npz_file, allow_pickle=True)
        ret["bert"] = np_object["bert"].item()
        try:
            ret["triggers"] = np_object["triggers"].item()
        except KeyError:
            ret["trigger"] = None
        try:
            ret["arguments"] = np_object["arguments"].item()
        except KeyError:
            ret["arguments"] = None
        return ret


def get_serif_objects(serifxml_path):
    serif_doc = serifxml3.Document(serifxml_path)
    event_mentions = []

    for sentence_id, sentence in enumerate(serif_doc.sentences):
        st = sentence.sentence_theories[0]
        """:type: serifxml3.SentenceTheory"""
        em_set = []
        if sentence.event_mention_set is not None:
            em_set = sentence.event_mention_set
        elif st.event_mention_set is not None:
            em_set = st.event_mention_set
        event_mentions.extend(list(em_set))

    return serif_doc, event_mentions


def get_embedding_keys(event_mentions):
    """
    Find vector keys for event mentions and arguments.
    :param event_mentions:
    :return:
    """
    event_mention_map = {}
    argument_map = {}
    for event_mention in event_mentions:

        # get the offset pairs for all anchors in this EM
        if len(event_mention.anchors) > 0:
            trig_offset_list = [
                # end offset in NLPLingo is Serif's edt offset -1
                (a.anchor_node.start_char, a.anchor_node.end_char + 1)
                for a in event_mention.anchors]
        else:
            trig_offset_list = [(event_mention.anchor_node.start_char,
                                 event_mention.anchor_node.end_char + 1)]
        event_mention_map[event_mention] = trig_offset_list

        for argument in event_mention.arguments:
            arg_offset_set = set()
            if isinstance(argument.value, Mention):
                arg_node = argument.value.head
            elif isinstance(argument.value, ValueMention):
                arg_node = argument.value
            else:
                raise ValueError(
                    "argument.value should be a Mention or a ValueMention")

            # end offset in NLPLingo is Serif's edt offset -1
            arg_offset_set.add((arg_node.start_char, arg_node.end_char + 1))
            # NlpLingo and SERIF disagree on how many tokens a "head" can be.
            if isinstance(arg_node, SynNode):
                for node in arg_node:
                    arg_offset_set.add((node.start_char, node.end_char + 1))

            # get all possible keys of the mention
            arg_offset_list = [anchor_offsets + arg_offsets
                               for anchor_offsets in trig_offset_list
                               for arg_offsets in arg_offset_set]
            argument_map[argument] = arg_offset_list

    return event_mention_map, argument_map


def get_argument_vectors_from_keys(embeddings, keys):
    return _get_vectors_from_keys(embeddings, keys, "arguments")


def get_trigger_vectors_from_keys(embeddings, keys):
    return _get_vectors_from_keys(embeddings, keys, "triggers")


def _get_vectors_from_keys(embeddings, keys, vector_type):
    ret = []
    for key in keys:
        if key in embeddings['bert']:
            model_vectors = {
                extractor: vectors[key]
                for extractor, vectors in embeddings[vector_type].items()
                if key in vectors}
            model_vectors['bert'] = embeddings['bert'][key]
            ret.append(model_vectors)
    # # debug
    # if not ret:
    #     print('$$$', vector_type, keys, embeddings['bert'].keys())
    return ret


def get_pretrained_embeddings(serifxml_list, npz_path, event_mention_trigger_callback, event_mention_argument_callback):
    # collect npzs
    npz_path = os.path.realpath(os.path.abspath(npz_path))
    if os.path.isdir(npz_path):
        npz_map = get_docid_to_embeddings_from_dir(npz_path)
    else:
        npz_map = get_docid_to_embeddings_from_list(npz_path)

    # collect serifxmls
    serifxml_list = os.path.realpath(os.path.abspath(serifxml_list))
    with codecs.open(serifxml_list, 'r', encoding='utf8') as fh:
        serifxmls = [path.strip() for path in fh.readlines()]

    # process serifxmls
    for path in serifxmls:
        serif_doc, event_mentions = get_serif_objects(path)
        docid = serif_doc.docid
        npz = load_embeddings(docid, npz_map)

        # get maps from serif objects to vector addresses
        event_mention_to_vector_keys, argument_to_vector_keys = (
            get_embedding_keys(event_mentions))

        # get embeddings themselves
        for event_mention, trigger_keys in event_mention_to_vector_keys.items():
            trigger_vecs = get_trigger_vectors_from_keys(npz, trigger_keys)

            event_mention_trigger_callback(serif_doc, event_mention, trigger_vecs)

            for argument in event_mention.arguments:
                argument_keys = argument_to_vector_keys[argument]
                argument_vecs = get_argument_vectors_from_keys(
                    npz, argument_keys)
                event_mention_argument_callback(serif_doc, event_mention, trigger_vecs, argument, argument_vecs)


def event_mention_trigger_print_callback(serif_doc, event_mention, trigger_vecs):
    print("### {} \n{}".format(serif_doc.docid, event_mention.anchor_node, event_mention.event_type))
    print(trigger_vecs)


def event_mention_argument_print_callback(serif_doc, event_mention, trigger_vecs, argument, argument_vecs):
    print("--- {} \n{}".format(argument.role, argument.value))
    print(argument_vecs)


if __name__ == '__main__':
    serifxmls = sys.argv[1]
    npzs = sys.argv[2]
    # serifxmls = "/nfs/raid88/u10/users/hqiu/runjob/expts/Hume/fuzzy_test.012020/dump_serifxml_into_features/batch_files/00000"
    # npzs = "/home/hqiu/ld100/Hume_pipeline_2/Hume/expts/nlplingo_decoder.012020.modified/nn_events/embeddings"
    get_pretrained_embeddings(serifxmls, npzs, event_mention_trigger_print_callback,
                              event_mention_argument_print_callback)

<filename>src/python/serif/embeddings/get_pretrained_embeddings.py
import argparse
import codecs
import json
import os
import sys
import logging

import numpy as np

import serifxml3
from serif.theory.mention import Mention
from serif.theory.syn_node import SynNode
from serif.theory.value_mention import ValueMention


def get_docid_to_embeddings_from_dir(embeddings_dir):
    embeddings_map = {}
    if embeddings_dir is not None:
        for npz in os.listdir(embeddings_dir):
            if npz.endswith('.npz'):
                docid = npz[:-len('.npz')]
                embeddings_map[docid] = os.path.join(embeddings_dir, npz)
    return embeddings_map


def get_docid_to_embeddings_from_list(embeddings_filelist):
    embeddings_map = {}
    with open(embeddings_filelist, 'r') as fh:
        for line in fh:
            line = line.strip()
            docid = line[:-len('.npz')].rsplit('/')[-1]
            embeddings_map[docid] = line
    return embeddings_map


def load_embeddings(docid, npz_files):
    """
    Loads an npz file containing embeddings generated by NlpLingo.
    The output of this function has the following format:
    {
        "triggers": {
            <trigger_extractor_name>: {

                # key for triggers:
                (int(<start_offset>), int(<end_offset>)): (

                    int(<start_offset>),
                    int(<end_offset>),
                    ndarray(<trigger_embedding>)), [...]

            }, [...]
        },

        "arguments": {
            <argument_extractor_name>: {

                # key for arguments:
                (int(<trigger_start_offset>),
                 int(<trigger_end_offset>),
                 int(<argument_start_offset>),
                 int(<argument_end_offset>)): (

                    int(<trigger_start_offset>),
                    int(<trigger_end_offset>),
                    int(<argument_start_offset>),
                    int(<argument_end_offset>),
                    ndarray(<argument_embedding>)), [...]

            }, [...]
        },

        "bert": {
            # for triggers, the trigger key above:
            (int(<start_offset>), int(<end_offset>)): (

                int(<sentence_id>),
                int(<token_id>)),

            # for arguments, the argument key above:
            (int(<trigger_start_offset>),
             int(<trigger_end_offset>),
             int(<argument_start_offset>),
             int(<argument_end_offset>)): (

                int(<argument_start_offset>),
                int(<argument_end_offset>)), [...]

        }
    }

    :param docid: a docid
    :param npz_files: a docid2filepath dict for NlpLingo npz output files
    :return: dict as described above
    """
    npz_file = npz_files.get(docid)
    if npz_file is None:
        return None
    else:
        ret = {}
        np_object = np.load(npz_file, allow_pickle=True)
        ret["bert"] = np_object["bert"].item()
        try:
            ret["triggers"] = np_object["triggers"].item()
        except KeyError:
            ret["trigger"] = None
        try:
            ret["arguments"] = np_object["arguments"].item()
        except KeyError:
            ret["arguments"] = None
        return ret


def get_serif_objects(serifxml_path):
    serif_doc = serifxml3.Document(serifxml_path)
    event_mentions = []

    for sentence_id, sentence in enumerate(serif_doc.sentences):
        st = sentence.sentence_theories[0]
        """:type: serifxml3.SentenceTheory"""
        em_set = []
        if sentence.event_mention_set is not None:
            em_set = sentence.event_mention_set
        elif st.event_mention_set is not None:
            em_set = st.event_mention_set
        event_mentions.extend(list(em_set))

    return serif_doc, event_mentions


def get_embedding_keys(event_mentions):
    """
    Find vector keys for event mentions and arguments.
    :param event_mentions:
    :return:
    """
    event_mention_map = {}
    argument_map = {}
    for event_mention in event_mentions:

        # get the offset pairs for all anchors in this EM
        if len(event_mention.anchors) > 0:
            trig_offset_list = [
                # end offset in NLPLingo is Serif's edt offset -1
                (a.anchor_node.start_char, a.anchor_node.end_char + 1)
                for a in event_mention.anchors]
        else:
            trig_offset_list = [(event_mention.anchor_node.start_char,
                                 event_mention.anchor_node.end_char + 1)]
        event_mention_map[event_mention] = trig_offset_list

        for argument in event_mention.arguments:
            arg_offset_set = set()
            if isinstance(argument.value, Mention):
                arg_node = argument.value.head
            elif isinstance(argument.value, ValueMention):
                arg_node = argument.value
            else:
                raise ValueError(
                    "argument.value should be a Mention or a ValueMention")

            # end offset in NLPLingo is Serif's edt offset -1
            arg_offset_set.add((arg_node.start_char, arg_node.end_char + 1))
            # NlpLingo and SERIF disagree on how many tokens a "head" can be.
            if isinstance(arg_node, SynNode):
                for node in arg_node:
                    arg_offset_set.add((node.start_char, node.end_char + 1))

            # get all possible keys of the mention
            arg_offset_list = [anchor_offsets + arg_offsets
                               for anchor_offsets in trig_offset_list
                               for arg_offsets in arg_offset_set]
            argument_map[argument] = arg_offset_list

    return event_mention_map, argument_map


def get_argument_vectors_from_keys(embeddings, keys):
    return _get_vectors_from_keys(embeddings, keys, "arguments")


def get_trigger_vectors_from_keys(embeddings, keys):
    return _get_vectors_from_keys(embeddings, keys, "triggers")


def _get_vectors_from_keys(embeddings, keys, vector_type):
    ret = []
    for key in keys:
        if key in embeddings['bert']:
            model_vectors = {
                extractor: vectors[key]
                for extractor, vectors in embeddings[vector_type].items()
                if key in vectors}
            model_vectors['bert'] = embeddings['bert'][key]
            ret.append(model_vectors)
    # # debug
    # if not ret:
    #     print('$$$', vector_type, keys, embeddings['bert'].keys())
    return ret


def get_pretrained_embeddings(serifxml_list, npz_path, event_mention_trigger_callback, event_mention_argument_callback):
    # collect npzs
    npz_path = os.path.realpath(os.path.abspath(npz_path))
    if os.path.isdir(npz_path):
        npz_map = get_docid_to_embeddings_from_dir(npz_path)
    else:
        npz_map = get_docid_to_embeddings_from_list(npz_path)

    # collect serifxmls
    serifxml_list = os.path.realpath(os.path.abspath(serifxml_list))
    with codecs.open(serifxml_list, 'r', encoding='utf8') as fh:
        serifxmls = [path.strip() for path in fh.readlines()]

    # process serifxmls
    for path in serifxmls:
        serif_doc, event_mentions = get_serif_objects(path)
        docid = serif_doc.docid
        npz = load_embeddings(docid, npz_map)

        # get maps from serif objects to vector addresses
        event_mention_to_vector_keys, argument_to_vector_keys = (
            get_embedding_keys(event_mentions))

        # get embeddings themselves
        for event_mention, trigger_keys in event_mention_to_vector_keys.items():
            trigger_vecs = get_trigger_vectors_from_keys(npz, trigger_keys)

            event_mention_trigger_callback(serif_doc, event_mention, trigger_vecs)

            for argument in event_mention.arguments:
                argument_keys = argument_to_vector_keys[argument]
                argument_vecs = get_argument_vectors_from_keys(
                    npz, argument_keys)
                event_mention_argument_callback(serif_doc, event_mention, trigger_vecs, argument, argument_vecs)


def event_mention_trigger_print_callback(serif_doc, event_mention, trigger_vecs):
    print("### {} \n{}".format(serif_doc.docid, event_mention.anchor_node, event_mention.event_type))
    print(trigger_vecs)


def event_mention_argument_print_callback(serif_doc, event_mention, trigger_vecs, argument, argument_vecs):
    print("--- {} \n{}".format(argument.role, argument.value))
    print(argument_vecs)


if __name__ == '__main__':
    serifxmls = sys.argv[1]
    npzs = sys.argv[2]
    # serifxmls = "/nfs/raid88/u10/users/hqiu/runjob/expts/Hume/fuzzy_test.012020/dump_serifxml_into_features/batch_files/00000"
    # npzs = "/home/hqiu/ld100/Hume_pipeline_2/Hume/expts/nlplingo_decoder.012020.modified/nn_events/embeddings"
    get_pretrained_embeddings(serifxmls, npzs, event_mention_trigger_print_callback,
                              event_mention_argument_print_callback)

Loads an npz file containing embeddings generated by NlpLingo.     The output of this function has the following format:     {         "triggers": {             <trigger_extractor_name>: {                  # key for triggers:                 (int(<start_offset>), int(<end_offset>)): (                      int(<start_offset>),                     int(<end_offset>),                     ndarray(<trigger_embedding>)), [...]              }, [...]         },          "arguments": {             <argument_extractor_name>: {                  # key for arguments:                 (int(<trigger_start_offset>),                  int(<trigger_end_offset>),                  int(<argument_start_offset>),                  int(<argument_end_offset>)): (                      int(<trigger_start_offset>),                     int(<trigger_end_offset>),                     int(<argument_start_offset>),                     int(<argument_end_offset>),                     ndarray(<argument_embedding>)), [...]              }, [...]         },          "bert": {             # for triggers, the trigger key above:             (int(<start_offset>), int(<end_offset>)): (                  int(<sentence_id>),                 int(<token_id>)),              # for arguments, the argument key above:             (int(<trigger_start_offset>),              int(<trigger_end_offset>),              int(<argument_start_offset>),              int(<argument_end_offset>)): (                  int(<argument_start_offset>),                 int(<argument_end_offset>)), [...]          }     }      :param docid: a docid     :param npz_files: a docid2filepath dict for NlpLingo npz output files     :return: dict as described above :type: serifxml3.SentenceTheory Find vector keys for event mentions and arguments.     :param event_mentions:     :return: # get the offset pairs for all anchors in this EM # end offset in NLPLingo is Serif's edt offset -1 # end offset in NLPLingo is Serif's edt offset -1 # NlpLingo and SERIF disagree on how many tokens a "head" can be. # get all possible keys of the mention # # debug # if not ret: #     print('$$$', vector_type, keys, embeddings['bert'].keys()) # collect npzs # collect serifxmls # process serifxmls # get maps from serif objects to vector addresses # get embeddings themselves ## {} \n{}".format(serif_doc.docid, event_mention.anchor_node, event_mention.event_type)) # serifxmls = "/nfs/raid88/u10/users/hqiu/runjob/expts/Hume/fuzzy_test.012020/dump_serifxml_into_features/batch_files/00000" # npzs = "/home/hqiu/ld100/Hume_pipeline_2/Hume/expts/nlplingo_decoder.012020.modified/nn_events/embeddings"