utils.py

import pandas as pd
import random
from preprocess_FrameNet import parse_frames, parse_annotated_texts
from pathlib import Path
import json
from typing import Any, Dict, List
import re

def remove_opt_fes(annotated_frame:Dict[str, Any], all_frames:List[Dict[str,Any]]) -> Dict[str,Any]:

    """
    Remove peripheral (optional) frame elements from an annotated frame.

    :param annotated_frame: frame annotated in text.
    :param all_frames: List of frame definitions, each including frame elements and their types.

    :return: The updated frame annotation without peripheral frame elements.

    Raises:
        ValueError: if an annotated frame is not found in frame definitions.
    """

    frame_name = annotated_frame["name"]

    # find the annotated frame in frame definitions
    frame_def = next((frame for frame in all_frames if frame.get("name") == frame_name),None)

    if frame_def is None:
        raise ValueError(
            f"Could not find annotated frame {frame_name} in frame resources when trying to remove peripheral frame elements")

    # collect peripheral FE names
    peripheral_fes = {fe["name"] for fe in frame_def.get("frame_elements", []) if  "peripheral" in fe.get("type").lower()}

    # filter annotated frame elements that are peripheral
    annotated_frame["frame_elements"] = [fe for fe in annotated_frame["frame_elements"] if fe["name"] not in peripheral_fes]

    return annotated_frame

def select_annotated_sentences(text: Dict[str, Any],
                               all_frames: List[Dict[str,Any]],
                               proportion:float=.1,
                               min_frames:int=1,
                               max_frames:int=29) -> List[Dict[str,Any]]:

    """
    Randomly select annotated sentences for evaluation.

    :param text: A text object containing:
    - "ID": Identifier of the text.
    - "sentences": List of sentence dicts with annotations.
    :param all_frames: Frame definitions used to filter frame elements.
    :param proportion: Fraction of sentences to select. Defaults to 0.1.
    :param min_frames: Minimum number of frames required in a sentence. Defaults to 2.

    :return: List of selected sentence dictionaries, each containing: "corpus_ID", "text", "frames"
    """

    # randomly select <proportion> of the annotated sentences in the text to be evaluated
    # the sentences should have at least <min_frames> frames, and we only keep core frame elements
    n_selected_sentences = int(len(text["sentences"]) * proportion)
    sentence_indices = [i for i in range(len(text["sentences"]))]
    random.seed(42)
    random.shuffle(sentence_indices)
    selected_sentences: list[dict] = []
    counter = 0

    for i in sentence_indices:

        if counter >= n_selected_sentences:
            break

        sentence_text = text["sentences"][i]["text"]
        # remove trailing spaces around punctuations
        sentence_text = re.sub(r"\s+([.,!?;:])", r"\1", sentence_text)
        sentence_text = sentence_text.strip()

        sentence = {"corpus_ID": text["ID"],
                    "text": sentence_text,
                    "frames": text["sentences"][i]["frames"]}

        frames = sentence.get("frames")

        # check if sentence has at least 2 annotated frames
        if isinstance(frames, list) and min_frames <= len(frames) <= max_frames:

            # remove trailing spaces around punctuations in span strings
            for i, f in enumerate(frames):
                if f.get('target'):
                    sentence["frames"][i]["target"] = re.sub(r"\s+([.,!?;:])", r"\1", f['target'])
                if f.get('frame_elements'):
                    for i_fe, fe in enumerate(f["frame_elements"]):
                        if fe.get('span'):
                            sentence["frames"][i]["frame_elements"][i_fe]["span"] = re.sub(r"\s+([.,!?;:])", r"\1", fe['span'])

            # make sure only core frame elements are stored
            if isinstance(all_frames, list):
                frames = sentence.get("frames")
                only_core_fes = []
                for f in frames:
                    only_core_fes.append(remove_opt_fes(f, all_frames))
                sentence["frames"] = only_core_fes

            selected_sentences.append(sentence)
            counter += 1

    # assert len(selected_sentences) == n_selected_sentences, f'{len(selected_sentences)}, {n_selected_sentences}, {text}'
    # some files/texts do not have any annotated frames

    return selected_sentences

def select_texts_4_eval(args:Any)-> List[Dict[str,Any]]:

    """
    Select annotated sentences from a directory of annotated texts for evaluation.
    This function either loads a precomputed selection or parses and selects sentences from raw annotated text files.

    :param args: Argument object expected to have:
    - annotated_texts_dir (str): Directory containing annotated text files.
    - frame_dir (str): Directory containing frame definitions.
    - frame_output_filepath (str): Path to store/load parsed frame definitions.

    :return: Dictionary with key "texts" containing selected sentences.
    Raises:
        ValueError: If directory is missing or empty.
        NotADirectoryError: If provided path does not exist.
    """

    # handle case of directory with annotated texts not give
    if not args.annotated_texts_dir:
        raise ValueError('Directory containing annotated texts is not specified. Please specify --annotated_texts_dir.')

    texts_dir = Path(args.annotated_texts_dir)

    # handle case of directory given non-existent or it's empty
    if not texts_dir.exists():
        raise NotADirectoryError(f"Directory {texts_dir} given as --annotated_texts_dir does not exist. Provide a valid directory.")
    if not any(texts_dir.iterdir()):
        raise ValueError(f"Directory {texts_dir} given as --annotated_texts_dir is empty. Provide a directory with annotated texts.")

    annotated_texts = {"texts": []}
    output_filepath = texts_dir / "texts_4_eval.json"

    # try loading texts if already selected, otherwise select texts for evaluation
    if output_filepath in texts_dir.iterdir():
        with open(output_filepath) as f:
            annotated_texts = json.load(f)
        print(f"Loaded annotated texts for evaluation from {texts_dir / 'texts_4_eval.json'}")

    if not annotated_texts.get("texts"):

        # parse raw texts in directory given
        print(f"Parsing annotated texts from {texts_dir}")
        parse_annotated_texts(texts_dir=texts_dir)

        # parse or load frames
        all_frames = parse_frames(frame_dir=args.frame_dir, output_filepath=args.frame_output_filepath)

        print(f'Selecting annotated texts for evaluation from {texts_dir}')
        # select sentences for evaluation
        for text_path in texts_dir.glob("*.json"):

            try:
                with open(text_path) as f:
                    text_data = json.load(f)
                selected_sentences = select_annotated_sentences(text_data, all_frames)
                annotated_texts["texts"].extend(selected_sentences)

            except json.JSONDecodeError:
                print(f"Invalid JSON in: {text_path}")

        # if --proportion_eval_texts is .1, 605 annotated sentences selected (1700 frames)
        # (10% of sentences with at least 2 annotated frames randomly selected from each text)
        with open(output_filepath, "w") as f:
            json.dump(annotated_texts, f, indent=2, ensure_ascii=False)
        print(f'Saved selected texts for evaluation to {output_filepath}')

    return annotated_texts

def define_text_input(args:Any)->List[str]:

    """
    Load input texts from a CSV file or a plain text file.

    :param args: ArgumentParser object expected to have:
    - input_filepath (str): path to input text file (.csv or .txt)
    - sep (str): Separator used in CSV file
    - text_column_name (str): Name of column in input text file (required for CSV).

    :return: A list of text strings extracted from the input file.
    """

    # handle different kinds of input

    texts = []

    if not args.input_filepath:
        # if not input filepath is given, try selecting texts from FrameNet annotated texts for generation
        annotated_texts = select_texts_4_eval(args)
        texts = [text["text"] for text in annotated_texts["texts"]]

    elif args.input_filepath.endswith(".csv"):

        texts_df = pd.read_csv(args.input_filepath, sep=args.sep)

        if not args.text_column_name:
            raise ValueError(
                f"No column name specified for {args.input_filepath}. Use --text_column_name to specify a column name.")

        if not args.text_column_name in texts_df.columns:
            raise ValueError(f'Column {args.text_column_name} not found in {args.input_filepath}.')

        texts = texts_df[args.text_column_name].tolist()

        if not texts:
            raise ValueError(f'No texts found in column {args.text_column_name}.')

    else: # e.g. text files

        try:
            with open(args.input_filepath, encoding="utf-8") as f:
                texts = f.read().splitlines()
        except FileNotFoundError as e:
            raise FileNotFoundError(f"File {args.input_filepath} given as --input_filepath not found. Please check the path and try again.") from e

    if not texts:
        raise ValueError(f"No texts found. Please provide a file with texts (either .csv or .txt) with --input_filepath, or provide a text string directly with --input_text.")

    return texts