src/data/prepare_sciarg_crosssection_annotations.py

import argparse
import logging
import os
import re
import shutil
from collections import defaultdict
from typing import Dict, List, Optional, Tuple

import pandas as pd
from pie_datasets import Dataset, IterableDataset, load_dataset
from pie_datasets.builders.brat import BratDocumentWithMergedSpans

logger = logging.getLogger(__name__)


def find_span_idx(raw_text: str, span_string: str) -> Optional[List]:
    """
    Match span string to raw text (document).
    Return either
    1) Tuple, 2) List of Tuples (more than one span match), or 3) empty List (no span match).
    """
    # remove possibly accidentally added white spaces
    span_string.strip()
    # use raw text input as regex-safe pattern
    safe = re.escape(span_string)
    pattern = rf"{safe}"
    # find match(es)
    out = [(s.start(), s.end()) for s in re.finditer(pattern, raw_text)]
    return out


def append_spans_start_and_end(
    raw_text: str,
    pd_table: pd.DataFrame,
    input_cols: List[str],
    input_idx_cols: List[str],
    output_cols: List[str],
    doc_id_col: str = "doc ID",
) -> pd.DataFrame:
    """
    Create new column(s) for span indexes (i.e. start and end as Tuple) in pd.DataFrame from span strings.
    Warn if
    1) span string does not match anything in document -> None,
    2) span string is not unique in the document -> List[Tuple].
    """
    pd_table = pd_table.join(pd.DataFrame(columns=output_cols))
    for idx, pd_row in pd_table.iterrows():
        for in_col, idx_col, out_col in zip(input_cols, input_idx_cols, output_cols):
            span_indices = find_span_idx(raw_text, pd_row[in_col])
            str_idx = pd_row[idx_col]
            span_idx = None
            if span_indices is None or len(span_indices) == 0:
                logger.warning(
                    f'The span "{pd_row[in_col]}" in Column "{in_col}" does not exist in {pd_row[doc_id_col]}.'
                )
            elif len(span_indices) == 1:
                # warn if column is not empty, but span is unique
                if str_idx == str_idx:
                    logger.warning(f'Column "{idx_col}" is not empty. It has value: {str_idx}.')
                span_idx = span_indices.pop()
            else:
                # warn if span not unique, but column is empty
                if str_idx != str_idx:
                    logger.warning(
                        f'The span "{pd_row[in_col]}" in Column "{in_col}" is not unique,'
                        f'but, column "{idx_col}" is empty. '
                        f"Need a string index to specify the non-unique span."
                    )
                else:
                    span_idx = span_indices.pop(int(str_idx))

            if span_idx is not None:
                pd_table.at[idx, out_col] = span_idx

                # sanity check (NOTE: this should live in a test)
                search_string = pd_row[in_col]
                reconstructed_string = raw_text[span_idx[0] : span_idx[1]]
                if search_string != reconstructed_string:
                    raise ValueError(
                        f"Reconstructed string does not match the original string. "
                        f"Original: {search_string}, Reconstructed: {reconstructed_string}"
                    )
    return pd_table


def get_texts_from_pie_dataset(
    doc_ids: List[str], **dataset_kwargs
) -> Dict[str, BratDocumentWithMergedSpans]:
    """Get texts from a PIE dataset for a list of document IDs.

    :param doc_ids: list of document IDs
    :param dataset_kwargs: keyword arguments to pass to load_dataset

    :return: a dictionary with document IDs as keys and texts as values
    """

    text_based_dataset = load_dataset(**dataset_kwargs)
    if not isinstance(text_based_dataset, (Dataset, IterableDataset)):
        raise ValueError(
            f"Expected a PIE Dataset or PIE IterableDataset, but got a {type(text_based_dataset)} instead."
        )
    if not issubclass(text_based_dataset.document_type, BratDocumentWithMergedSpans):
        raise ValueError(
            f"Expected a PIE Dataset with BratDocumentWithMergedSpans as document type, "
            f"but got {text_based_dataset.document_type} instead."
        )
    doc_id2text = {doc.id: doc for doc in text_based_dataset}
    return {doc_id: doc_id2text[doc_id] for doc_id in doc_ids}


def set_span_annotation_ids(
    table: pd.DataFrame,
    doc_id2doc: Dict[str, BratDocumentWithMergedSpans],
    doc_id_col: str,
    span_idx_cols: List[str],
    span_id_cols: List[str],
) -> pd.DataFrame:
    """
    Create new column(s) for span annotation IDs in pd.DataFrame from span indexes. The annotation IDs are
    retrieved from the TextBasedDocument object using the span indexes.

    :param table: pd.DataFrame with span indexes, document IDs, and other columns
    :param doc_id2doc: dictionary with document IDs as keys and BratDocumentWithMergedSpans objects as values
    :param doc_id_col: column name that contains document IDs
    :param span_idx_cols: column names that contain span indexes
    :param span_id_cols: column names for new span ID columns

    :return: pd.DataFrame with new columns for span annotation IDs
    """
    table = table.join(pd.DataFrame(columns=span_id_cols))
    span2id: Dict[str, Dict[Tuple[int, int], str]] = defaultdict(dict)
    for doc_id, doc in doc_id2doc.items():
        for span_id, span in zip(doc.metadata["span_ids"], doc.spans):
            span2id[doc_id][(span.start, span.end)] = span_id

    for span_idx_col, span_id_col in zip(span_idx_cols, span_id_cols):
        table[span_id_col] = table.apply(
            lambda row: span2id[row[doc_id_col]][tuple(row[span_idx_col])], axis=1
        )

    return table


def set_relation_annotation_ids(
    table: pd.DataFrame,
    doc_id2doc: Dict[str, BratDocumentWithMergedSpans],
    doc_id_col: str,
    relation_id_col: str,
) -> pd.DataFrame:
    """create new column for relation annotation IDs in pd.DataFrame. They are simply new ids starting from the last
    relation annotation id in the document.

    :param table: pd.DataFrame with document IDs and other columns
    :param doc_id2doc: dictionary with document IDs as keys and BratDocumentWithMergedSpans objects as values
    :param doc_id_col: column name that contains document IDs
    :param relation_id_col: column name for new relation ID column

    :return: pd.DataFrame with new column for relation annotation IDs
    """

    table = table.join(pd.DataFrame(columns=[relation_id_col]))
    doc_id2highest_relation_id = defaultdict(int)

    for doc_id, doc in doc_id2doc.items():
        # relation ids are prefixed with "R" in the dataset
        doc_id2highest_relation_id[doc_id] = max(
            [int(relation_id[1:]) for relation_id in doc.metadata["relation_ids"]]
        )

    for idx, row in table.iterrows():
        doc_id = row[doc_id_col]
        doc_id2highest_relation_id[doc_id] += 1
        table.at[idx, relation_id_col] = f"R{doc_id2highest_relation_id[doc_id]}"

    return table


def main(
    input_path: str,
    output_path: str,
    brat_data_dir: str,
    input_encoding: str,
    include_unsure: bool = False,
    doc_id_col: str = "doc ID",
    unsure_col: str = "unsure",
    span_str_cols: List[str] = ["head argument string", "tail argument string"],
    str_idx_cols: List[str] = ["head string index", "tail string index"],
    span_idx_cols: List[str] = ["head_span_idx", "tail_span_idx"],
    span_id_cols: List[str] = ["head_span_id", "tail_span_id"],
    relation_id_col: str = "relation_id",
    set_annotation_ids: bool = False,
    relation_type: str = "relation",
) -> None:
    """
    Convert long dependency annotations from a CSV file to a JSON format. The input table should have
    columns for document IDs, argument span strings, and string indexes (required in the case that the
    span string occurs multiple times in the base text). The argument span strings are matched to the
    base text to get the start and end indexes of the span. The output JSON file will have the same
    columns as the input file, plus two additional columns for the start and end indexes of the spans.

    :param input_path: path to a CSV/Excel file that contains annotations
    :param output_path: path to save JSON output
    :param brat_data_dir: directory where the BRAT data (base texts and annotations) is located
    :param input_encoding: encoding of the input file. Only used for CSV files. Default: "cp1252"
    :param include_unsure: include annotations marked as unsure. Default: False
    :param doc_id_col: column name that contains document IDs. Default: "doc ID"
    :param unsure_col: column name that contains unsure annotations. Default: "unsure"
    :param span_str_cols: column names that contain span strings. Default: ["head argument string", "tail argument string"]
    :param str_idx_cols: column names that contain string indexes. Default: ["head string index", "tail string index"]
    :param span_idx_cols: column names for new span-index columns. Default: ["head_span_idx", "tail_span_idx"]
    :param span_id_cols: column names for new span-ID columns. Default: ["head_span_id", "tail_span_id"]
    :param relation_id_col: column name for new relation-ID column. Default: "relation_id"
    :param set_annotation_ids: set annotation IDs for the spans and relations. Default: False
    :param relation_type: specify the relation type for the BRAT output. Default: "relation"

    :return: None
    """
    # get annotations from a csv file
    if input_path.lower().endswith(".csv"):
        input_df = pd.read_csv(input_path, encoding=input_encoding)
    elif input_path.lower().endswith(".xlsx"):
        logger.warning(
            f"encoding parameter (--input-encoding={input_encoding}) is ignored for Excel files."
        )
        input_df = pd.read_excel(input_path)
    else:
        raise ValueError("Input file has unexpected format. Please provide a CSV or Excel file.")

    # remove unsure
    if not include_unsure:
        input_df = input_df[input_df[unsure_col].isna()]
    # remove all empty columns
    input_df = input_df.dropna(axis=1, how="all")

    # define output DataFrame
    result_df = pd.DataFrame(columns=[*input_df.columns, *span_idx_cols])

    # get unique document IDs
    doc_ids = list(input_df[doc_id_col].unique())

    # get base texts from a PIE SciArg dataset
    doc_id2doc = get_texts_from_pie_dataset(
        doc_ids=doc_ids,
        path="pie/brat",
        name="merge_fragmented_spans",
        split="train",
        revision="769a15e44e7d691148dd05e54ae2b058ceaed1f0",
        base_dataset_kwargs=dict(data_dir=brat_data_dir),
    )

    for doc_id in doc_ids:

        # iterate over each sub-df that contains annotations for a single document
        doc_df = input_df[input_df[doc_id_col] == doc_id]
        input_df = input_df.drop(doc_df.index)
        # get spans' start and end indexes as new columns
        doc_with_span_indices_df = append_spans_start_and_end(
            raw_text=doc_id2doc[doc_id].text,
            pd_table=doc_df,
            input_cols=span_str_cols,
            input_idx_cols=str_idx_cols,
            output_cols=span_idx_cols,
        )
        # append this sub-df (with spans' indexes columns) to result_df
        result_df = pd.concat(
            [result_df if not result_df.empty else None, doc_with_span_indices_df]
        )

    out_ext = os.path.splitext(output_path)[1]
    save_as_brat = out_ext == ""

    if set_annotation_ids or save_as_brat:
        result_df = set_span_annotation_ids(
            table=result_df,
            doc_id2doc=doc_id2doc,
            doc_id_col=doc_id_col,
            span_idx_cols=span_idx_cols,
            span_id_cols=span_id_cols,
        )
        result_df = set_relation_annotation_ids(
            table=result_df,
            doc_id2doc=doc_id2doc,
            doc_id_col=doc_id_col,
            relation_id_col=relation_id_col,
        )

    base_dir = os.path.dirname(output_path)
    os.makedirs(base_dir, exist_ok=True)

    if out_ext.lower() == ".json":
        logger.info(f"Saving output in JSON format to {output_path} ...")
        result_df.to_json(
            path_or_buf=output_path,
            orient="records",
            lines=True,
        )  # possible orient values: 'split','index', 'table','records', 'columns', 'values'
    elif save_as_brat:
        logger.info(f"Saving output in BRAT format to {output_path} ...")
        os.makedirs(output_path, exist_ok=True)
        for doc_id in doc_ids:
            # handle the base text file (just copy from the BRAT data directory)
            shutil.copy(
                src=os.path.join(brat_data_dir, f"{doc_id}.txt"),
                dst=os.path.join(output_path, f"{doc_id}.txt"),
            )

            # handle the annotation file
            # first, read the original annotation file
            input_ann_path = os.path.join(brat_data_dir, f"{doc_id}.ann")
            with open(input_ann_path, "r") as f:
                ann_lines = f.readlines()
            # then, append new relation annotations
            # The format for each line is (see https://brat.nlplab.org/standoff.html):
            # R{relation_id}\t{relation_type} Arg1:{span_id1} Arg2:{span_id2}
            doc_df = result_df[result_df[doc_id_col] == doc_id]
            logger.info(f"Adding {len(doc_df)} relation annotations to {doc_id}.ann ...")
            for idx, row in doc_df.iterrows():
                head_span_id = row[span_id_cols[0]]
                tail_span_id = row[span_id_cols[1]]
                relation_id = row[relation_id_col]
                ann_line = (
                    f"{relation_id}\t{relation_type} Arg1:{head_span_id} Arg2:{tail_span_id}\n"
                )
                ann_lines.append(ann_line)
            # finally, write the new annotation file
            output_ann_path = os.path.join(output_path, f"{doc_id}.ann")
            with open(output_ann_path, "w") as f:
                f.writelines(ann_lines)
    else:
        raise ValueError(
            "Output file has unexpected format. Please provide a JSON file or a directory."
        )

    logger.info("Done!")


if __name__ == "__main__":

    """
    example call:
    python src/data/prepare_sciarg_crosssection_annotations.py
    // or //
    python src/data/prepare_sciarg_crosssection_annotations.py \
    --input-path data/annotations/sciarg-cross-section/aligned_input.csv \
    --output-path data/annotations/sciarg-with-abstracts-and-cross-section-rels \
    --brat-data-dir data/annotations/sciarg-abstracts/v0.9.3/alisa
    """

    logging.basicConfig(level=logging.INFO)

    parser = argparse.ArgumentParser(
        description="Read text files in a directory and a CSV file that contains cross-section annotations. "
        "Transform the CSV file to a JSON format and save at a specified output directory."
    )
    parser.add_argument(
        "--input-path",
        type=str,
        default="data/annotations/sciarg-cross-section/aligned_input.csv",
        help="Locate a CSV/Excel file.",
    )
    parser.add_argument(
        "--output-path",
        type=str,
        default="data/annotations/sciarg-with-abstracts-and-cross-section-rels",
        help="Specify a path where output will be saved. Should be a JSON file or a directory for BRAT output.",
    )
    parser.add_argument(
        "--brat-data-dir",
        type=str,
        default="data/annotations/sciarg-abstracts/v0.9.3/alisa",
        help="Specify the directory where the BRAT data (base texts and annotations) is located.",
    )
    parser.add_argument(
        "--relation-type",
        type=str,
        default="semantically_same",
        help="Specify the relation type for the BRAT output.",
    )
    parser.add_argument(
        "--input-encoding",
        type=str,
        default="cp1252",
        help="Specify encoding for reading an input file.",
    )
    parser.add_argument(
        "--include-unsure",
        action="store_true",
        help="Include annotations marked as unsure.",
    )
    parser.add_argument(
        "--set-annotation-ids",
        action="store_true",
        help="Set BRAT annotation IDs for the spans and relations.",
    )
    args = parser.parse_args()
    kwargs = vars(args)

    main(**kwargs)