preprocess_FrameNet.py

import xml.etree.ElementTree as ET
import json
from pathlib import Path
import re
import warnings
from typing import Any

# Namespace used throughout the file
NS = {"fn": "http://framenet.icsi.berkeley.edu"}

class Frame:

    def __init__(self, xml:str):
        """
        Representation of a FrameNet frame parsed from XML.

        :param xml: Path to the XML file containing the frame data.

        Attributes:
            name (str): Frame name.
            id (str): Frame ID.
            definition (str): Cleaned frame definition.
            examples (List[str]): Example sentences.
            frame_elements (List[Dict[str, Any]]): Frame elements data.
            lexical_units (List[Dict[str, str]]): Lexical units evoking the frame.
        """

        tree = ET.parse(xml)
        root = tree.getroot()

        self.name:str = root.attrib.get("name")
        self.id:str = root.attrib.get("ID")
        self.definition:str = self.parse_frame_definition(root)
        self.examples:list[str] = self.parse_frame_examples(root)
        self.frame_elements:list[dict[str,Any]] = self.extract_frame_elements(root)
        self.lexical_units:list[dict[str,str]] = self.extract_lexical_units(root)
        self.frame_relations:dict[str,list[str]] = self.extract_frame_relations(root)

    @staticmethod
    def parse_frame_definition(root:ET.Element)->str:

        """
        Extract and clean the frame definition text.

        :param root: XML root element.
        :return: cleaned frame definition text.
        """

        text = root.find("fn:definition", NS).text

        if re.search(r'<def-root>.*\n', text):
            text = re.search(r'<def-root>(.*?)\n', text)[0]
            text = text.replace('<def-root>', '') # remove def-root tag from definition text
        else:
            text = re.sub(r'<def-root>(.*?)</def-root>', # remove def-root tag from definition text
                          r'\1', text)

        text = re.sub(r'<fex\s+name="([^"]+)">(.*?)</fex>', # replace fex tags with the element name
                      r'\2', text)

        text = re.sub(r'<fen>(.*?)</fen>', # remove fen tags
                      r'\1', text)

        return text.strip()

    @staticmethod
    def parse_frame_examples(root:ET.Element)->list[str]:

        """
        Extract example sentences from frame definition.

        :param root: XML root element.
        :return: list of cleaned example strings.
        """

        text = root.find("fn:definition", NS).text

        examples = []

        for line in text.split("\n"):

            if line.startswith('<ex>'):
                # remove the ex and def-root wrappers
                line = re.sub(r'</?ex>', '', line)
                line = line.replace('</def-root>', '')
                # replace fex wrappers with frame element name
                line = re.sub(
                        r'<fex\s+name="([^"]+)">(.*?)</fex>',
                        r'<\1>\2</\1>',
                        line)
                # convert \t wrapper into \Target for clarity
                line = re.sub(
                    r'<t>(.*?)</t>',
                    r'<Target>\1</Target>',
                    line)
                if line:
                    examples.append(line)

        return examples

    @staticmethod
    def parse_frame_element_definition(text:str)->tuple[str, list[str]]:

        """
        Parse a frame element definition and its examples.

        :param text: raw frame element definition text in XML file
        :return: definition and list of examples
        """
        # split definition and example.txt
        if "<ex>" in text:
            def_parts = text.split("<ex>")
            definition = def_parts[0]
            examples = def_parts[1:]
        else:
            definition, examples = text, []

        # remove def-root tag from definition text
        definition = re.sub(r'</?def-root>', '', definition)
        # remove fen tags
        definition = re.sub(r'<fen>(.*?)</fen>', r'\1', definition)
        # remove fex tags
        definition = re.sub(r'<fex\s+name="([^"]+)">(.*?)</fex>', r'\2', definition)
        definition = definition.strip()

        if examples:
            for i, example in enumerate(examples):
                # strip out possible text outside example.txt
                example = re.sub(r'\n.*', '', example)
                # replace fex tags with frame element names
                example = re.sub(r'<fex\s+name="([^"]+)">(.*?)</fex>', r'<\1>\2</\1>', example)
                # replace t tag with Target for clarity
                example = re.sub(r'<t>(.*?)</t>',r'<Target>\1</Target>', example)
                # remove </ex> tag
                example = re.sub(r'</ex>', '', example)
                # remove </def-root> tag
                example = re.sub(r'</def-root>', '', example)
                example = example.strip()
                examples[i] = example

        return definition, examples

    def extract_frame_elements(self,root:ET.Element)->list[dict[str, Any]]:

        """
        Extract frame elements from XML.

        :param root: XML root element.
        :return: List of frame element dictionaries with parsed frame element info.
        """

        frame_elements = []

        fe_root = root.findall("fn:FE", NS)

        for fe in fe_root:
            frame_element = {'name': fe.attrib.get("name"),
                             'ID': fe.attrib.get("ID"),
                             'type': fe.attrib.get("coreType"),
                             'definition': '',
                             'examples': []}

            definition_root = fe.find("fn:definition", NS).text
            if definition_root:
                frame_element['definition'], frame_element['examples'] = self.parse_frame_element_definition(definition_root)

            frame_elements.append(frame_element)

        return frame_elements

    @staticmethod
    def extract_lexical_units(root:ET.Element)->list[dict[str, str]]:

        """
        Extract lexical units for a frame from XML.

        :param root: XML root element.
        :return: List of lexical unit dictionaries with parsed lexical unit info.
        """

        lexical_units = []

        lu_root = root.findall("fn:lexUnit", NS)

        for lu in lu_root:
            lexical_unit = {'name': lu.find("fn:lexeme", NS).attrib.get("name"),
                            'ID': lu.attrib.get("ID"),
                            'pos_tag': lu.attrib.get("POS"),
                            'definition': lu.find("fn:definition", NS).text.replace('FN: ', '')
                            }
            lexical_units.append(lexical_unit)

        return lexical_units

    @staticmethod
    def extract_frame_relations(root:ET.Element)->(list[str], list[str]):

        """
        Extract frame relations from XML.
        :param root: XML root element.
        :return: List of super frames and sub frames relative to the current frame.
        """
        frame_relations = {'inherits_from': [], "uses": [], "perspective_on": [], "subframe_of": [],
                           "precedes": [], "inchoative_of": [], "causative_of": [], "see_also": []}

        for relation in root.findall("fn:frameRelation", NS):
            relation_type = relation.attrib.get("type")
            frame_names = relation.findall("fn:relatedFrame", NS)
            if frame_names:
                for frame_name in frame_names:
                    if relation_type == "Inherits from":
                        frame_relations['inherits_from'].append(frame_name.text)
                    elif relation_type == "Uses":
                        frame_relations['uses'].append(frame_name.text)
                    elif relation_type == "Perspective on":
                        frame_relations['perspective_on'].append(frame_name.text)
                    elif relation_type == "Subframe of":
                        frame_relations['subframe_of'].append(frame_name.text)
                    elif relation_type == "Precedes":
                        frame_relations['precedes'].append(frame_name.text)
                    elif relation_type == "Is Inchoative of":
                        frame_relations['inchoative_of'].append(frame_name.text)
                    elif relation_type == "Is Causative of":
                        frame_relations['causative_of'].append(frame_name.text)
                    elif relation_type == "See also":
                        frame_relations['see_also'].append(frame_name.text)

        return frame_relations

    def to_dict(self):
        return {
            'name': self.name,
            'ID': self.id,
            'definition': self.definition,
            'examples': self.examples,
            'frame_elements': self.frame_elements,
            'lexical_units': self.lexical_units,
            'frame_relations': self.frame_relations
        }

class AnnotatedFrame:

    def __init__(self, annotation:ET.Element, sentence_text:str):

        """
        Representation of a frame annotation in a sentence.

        :param annotation: XML annotation element.
        :param sentence_text: Original sentence text.

        Attributes:
        - name: name of the annotated frame.
        - lexical_unit: the lexical unit that triggered the frame annotation.
        - target: the text span that triggered the frame annotation.
        - start: the starting character index of the target span in the sentence.
        - end: the ending character index of the target span in the sentence.
        - frame_elements: list of frame element dictionaries belonging to the annotated frame.
        """

        self.name = annotation.attrib.get("frameName")
        # self.id = annotation.attrib.get("frameID"),
        self.lexical_unit = annotation.attrib.get("luName")
        self.target, self.start, self.end = self.parse_target(annotation, sentence_text)
        self.frame_elements = self.parse_frame_elements(annotation, sentence_text)

    @staticmethod
    def parse_target(annotation:ET.Element, sentence_text:str)->tuple[str, None|int, None|int]:

        """
        Extract frame target (text span that evokes the frame) from annotation.

        :param annotation: XML annotation element.
        :param sentence_text: Sentence text.

        :return: Span string and start and end indices from frame target
        """

        target, target_start, target_end = "", None, None
        layer_root = annotation.findall("fn:layer", NS)
        for layer in layer_root:
            if layer.attrib.get("name") == "Target":
                target_info = layer.find("fn:label", NS)
                if target_info is not None:
                    target_span_start = target_info.attrib.get('start')
                    target_span_end = target_info.attrib.get('end')
                    if target_span_start and target_span_end:
                        target = sentence_text[int(target_span_start):int(target_span_end) + 1]
                        target_start = int(target_span_start)
                        target_end = int(target_span_end)

        return target, target_start, target_end

    @staticmethod
    def parse_frame_elements(annotation:ET.Element, sentence_text:str)->list[dict[str,Any]]:

        """
        Extract frame elements from frame annotation.

        :param annotation: Annotation element.
        :param sentence_text: Sentence text.

        :return: name, span, start index, and end index from frame elements belonging to the annotated frame.
        """

        layer_root = annotation.findall("fn:layer", NS)
        frame_elements = []
        for layer in layer_root:
            if layer.attrib.get("name") == "FE":
                fe_info_list = layer.findall("fn:label", NS)
                if fe_info_list:
                    for fe_info in fe_info_list:
                        fe_start = fe_info.attrib.get("start")
                        fe_end = fe_info.attrib.get("end")
                        if fe_start and fe_end:
                            frame_element = {'name': fe_info.attrib.get("name"),
                                             'span': sentence_text[int(fe_start):int(fe_end) + 1],
                                             'start': int(fe_start),
                                             'end': int(fe_end)}
                            frame_elements.append(frame_element)

        return frame_elements

    def to_dict(self):
        return {"name": self.name,
                # "ID": self.id,
                # "lexical_unit": self.lexical_unit,
                "target": self.target,
                "start": self.start,
                "end": self.end,
                "frame_elements": self.frame_elements}

class AnnotatedSentence:

    def __init__(self, sentence_root:ET.Element):

        """
        Representation of an annotated sentence.

        :param sentence_root: XML sentence element.
        Attributes:
        - text: sentence text
        - frames: list of frames annotated in the sentence.
        """

        # self.sent_id = counter # sentence_root.attrib.get("sentNo") AL: sentNo does not make sense in raw xml (e.g. first 3 sentences had ID 1)
        self.text = sentence_root.find("fn:text", NS).text.strip()
        self.frames = self.parse_frames(sentence_root)

    def parse_frames(self, sentence_root)->list[AnnotatedFrame]:

        """
        Extract manually annotated frames.

        :param sentence_root: XML sentence element.
        :return: List of frame annotations.
        """

        annotated_frames = []
        annotation_root = sentence_root.findall("fn:annotationSet", NS)

        for annotation in annotation_root:
            if annotation.attrib.get("status") == "MANUAL":
                frame = AnnotatedFrame(annotation, self.text)
                annotated_frames.append(frame)

        return annotated_frames

    def to_dict(self):
        return {"text": self.text,
                "frames": [frame.to_dict() for frame in self.frames]}

class AnnotatedText:

    def __init__(self, xml):

        """
        Representation of a FrameNet annotated text.

        :param xml: path to XML file with annotated text.

        Attributes:
        - corpus: name of corpus text belongs to.
        - description: description of text.
        - ID: ID of the text.
        - sentences: list of sentences annotated in the text.
        """

        tree = ET.parse(xml)
        root = tree.getroot()
        self.corpus = root.find("fn:header", NS).find("fn:corpus", NS).attrib.get("description")
        self.description = root.find("fn:header", NS).find("fn:corpus", NS).find("fn:document", NS).attrib.get("description")
        self.ID = root.find("fn:header", NS).find("fn:corpus", NS).find("fn:document", NS).attrib.get("ID")
        self.sentences = self.parse_sentences(root)

    @staticmethod
    def parse_sentences(root:ET.Element)->list[AnnotatedSentence]:

        """
        Extract sentences from XML.
        :param root: XML root element.
        :return: Parsed sentences.
        """

        sentences = []
        sentence_root = root.findall("fn:sentence", NS)
        # counter = 0
        for sentence in sentence_root:
            sentences.append(AnnotatedSentence(sentence))
            # counter += 1
        return sentences

    def to_dict(self):
        return {"corpus": self.corpus,
                "description": self.description,
                "ID": self.ID,
                "sentences": [sentence.to_dict() for sentence in self.sentences]}

def parse_frames(frame_dir:str|None, output_filepath:str|None)->list[dict[str,Any]]:

    """
    Parse FrameNet frame XML files into JSON representations.

    :param frame_dir: Directory containing FrameNet frame XML files.
    :param output_filepath: Path to save combined frames JSON.

    :return: List of parsed frame dictionaries.

    Raises:
        ValueError: If frame_dir is not provided.
        NotADirectoryError: If frame_dir is invalid.
    """

    if not frame_dir:
        raise ValueError("Provide a path to a folder with FrameNet frames (in xml or json) using --frame_dir.")

    frame_dir = Path(frame_dir)

    if not frame_dir.is_dir():
        raise NotADirectoryError(f'Directory {frame_dir} (folder with frames) does not exist. '
                                 f'Provide path to a valid directory with FrameNet frames using --frame_dir.')

    if not any(frame_dir.iterdir()):
        raise ValueError(
            f"Directory {frame_dir} is empty. "
            f"Provide path to directory with FrameNet frames using --frame_dir.")

    all_frames = []

    # only xml files in input directory
    for xml_path in frame_dir.glob("*.xml"):

        json_path = frame_dir / (xml_path.stem + ".json")

        # if json_path.is_file():
        #     # load it if already exists
        #     with open(json_path, encoding="utf-8") as f:
        #         frame_data = json.load(f)
        # else:
        frame_data = Frame(xml_path).to_dict()
        # save out as json file
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(frame_data, f, indent=2, ensure_ascii=False)

        all_frames.append(frame_data)

    # save out all frames as json file if not already there
    if not output_filepath:
        output_filepath = frame_dir / "all_frames.json"
    if not Path(output_filepath).is_file():
        with open(output_filepath, "w", encoding="utf-8") as f:
            json.dump(all_frames, f, indent=4, ensure_ascii=False)

    return all_frames

def parse_annotated_texts(texts_dir:str|None)->None:

    """
    Parse annotated FrameNet XML texts into JSON files.

    :param texts_dir: Directory containing FrameNet annotated text XML files.
    
    Raises:
        NotADirectoryError: If texts_dir is invalid.
    """

    if not texts_dir:
        warnings.warn('Path to annotated texts from FrameNet not given. Annotated texts will not be parsed.')

    else:
        texts_dir = Path(texts_dir)

        if not texts_dir.is_dir():
            raise NotADirectoryError(f'Directory "{texts_dir}" (folder with annotated texts from FrameNet) does not exist. '
                                     'Provide a valid directory path to FrameNet annotated texts using --annotated_texts_dir.')
        if not any(texts_dir.iterdir()):
            raise ValueError(
                f"Directory {texts_dir} is empty. Provide a directory with FrameNet annotated texts using --annotated_texts_dir.")

        # only xml files in input directory
        for xml_path in texts_dir.glob("*.xml"):

            json_path = texts_dir / (xml_path.stem + ".json")

            # if parsed text does not already exist, parse it
            if not json_path.is_file():

                # convert frames in xml to dictionary with relevant info
                text_data = AnnotatedText(xml_path).to_dict()
                # save out as json file
                with open(json_path, "w", encoding="utf-8") as f:
                    json.dump(text_data, f, indent=2, ensure_ascii=False)