evaluation.py

import json
from langchain_huggingface import HuggingFaceEmbeddings
from sklearn.metrics.pairwise import cosine_similarity
from preprocess_FrameNet import parse_frames
from pathlib import Path
from typing import Any
from collections import defaultdict
import re
import warnings

def span_overlap(span1:str, span2:str)->float:

    """Compute token-level overlap ratio (Jaccard) between two spans.

    :param span1: First text span
    :param span2: Second text span

    :return: Jaccard similarity between span1 and span2.
    """

    tokens1 = set(span1.lower().split())
    tokens2 = set(span2.lower().split())

    intersection = tokens1 & tokens2
    union = tokens1 | tokens2

    return len(intersection) / len(union)

def count_spans(text:str, frames:list[dict[str,Any]])->dict[str, int]:

    """
    Count how many times each span string appears in the frame annotations.

    :param text: Input text
    :param frames: Frame annotations containing frame targets (text span that evokes frame) and frame elements.

    :return: Mapping of span strings to their occurrence counts.
    """

    # collect span strings from both frame targets and frame elements for a text
    spans = set()
    for f in frames:
        spans.add(f["target"])
        for fe in f.get("frame_elements", []):
            spans.add(fe["span"])

    # then count how many times these spans appear in the text (then we know which are ambiguous and need index matching)
    counts = {}
    for span in spans:
        # escape regex special characters
        pattern = re.escape(span)
        # make spaces flexible (handles multiple/irregular spacing)
        pattern = pattern.replace(r"\ ", r"\s+")
        # add custom boundaries (span not inside words)
        pattern = r"(?<!\w)" + pattern + r"(?!\w)"
        # pattern = r"\b{}\b".format(re.escape(span))
        matches = re.findall(pattern, text)
        counts[span] = len(matches)

    return counts

def compute_similarity_frames(pred_name: str, gold_name: str, frames: list[dict[str,Any]],
                              embedding_model:HuggingFaceEmbeddings) -> float:

    """
    Compute semantic similarity between predicted and gold frames using embeddings.

    :param pred_name: predicted frame name
    :param gold_name: gold frame name
    :param frames: frame definitions
    :param embedding_model: embedding model to compute semantic similarity

    :return: Cosine similarity between predicted and gold frame representations
    """

    # -1 if the opposite, 0 if dissimilar, 1 if identical

    pred_frame_def, gold_frame_def = "", ""

    for frame in frames:
        if frame["name"] == pred_name:
            pred_frame_def = frame["definition"]
        if frame["name"] == gold_name:
            gold_frame_def = frame["definition"]
        if pred_frame_def and gold_frame_def:
            break

    pred_frame, gold_frame = pred_name, gold_name
    if pred_frame_def:
        pred_frame = f'{pred_name}: {pred_frame_def}'
    if gold_frame_def:
        gold_frame = f'{gold_name}: {gold_frame_def}'

    embeds = embedding_model.embed_documents([pred_frame, gold_frame])

    return cosine_similarity([embeds[0]], [embeds[1]])[0][0]

def compute_graph_relatedness(pred_name: str, gold_name: str, frames):

    pred_frame_relations, gold_frame_relations = dict(), dict()

    for frame in frames:
        if frame["name"] == pred_name:
            pred_frame_relations = frame["frame_relations"]
        if frame["name"] == gold_name:
            gold_frame_relations = frame["frame_relations"]

    if not gold_frame_relations:
        warnings.warn(f'Frame relations of gold frame {gold_name} not found.')
        return False, None
    if not pred_frame_relations:
        warnings.warn(f'Frame relations of predicted frame {pred_name} not found.')
        return False, None

    for relation_type in pred_frame_relations:
        for related_frame in pred_frame_relations[relation_type]:
            if related_frame == gold_name:
                # print(f'Predicted frame: {pred_name}, gold frame: {gold_name}, relation type: {relation_type}')
                return True, relation_type

    for relation_type in gold_frame_relations:
        for related_frame in gold_frame_relations[relation_type]:
            if related_frame == pred_name:
                # print(f'Gold frame: {gold_name}, pred frame: {pred_name}, relation type: {relation_type}')
                return True, relation_type

    return False, None

def match_item(pred:dict[str,Any],
               gold:dict[str,Any],
               span_counts:dict[str,int],
               match_cfg:dict[str,Any],
               is_frame:bool=False,
               span_match:bool=True)->bool:

    """
    Check whether a predicted item (frame or frame element) matches a gold item.

    Matching includes:
    - label equality or semantic similarity
    - span string similarity
    - optional positional tolerance for ambiguous spans

    Explanation on span matching: if span string match exactly, check if span appears only once in text. If so, exact string matches suffices.
    If not, span is ambiguous. In this case, check span start and end positions with a tolerance of <tol> positions to the left or to the right
    (since trailing spaces of orginal texts were removed for model generation, and thus the annotated span indices may not correspond exactly.
    Another reason for a wrong span index mismatch is that the model sometimes misses the final word letter or includes the space right after the word, for instance)

    :param pred: Predicted frame.
    :param gold: Gold frame.
    :param span_counts: Span occurrence counts in text.
    :param match_cfg: Matching configuration parameters.
    :param is_frame: Whether matching frames (True) or frame elements (False).
    :param span_match: Whether matching spans as well (True) or only labels (False).

    :return: True if the predicted item matches the gold item, False otherwise.
    """

    assert pred["name"] and gold["name"]
    assert match_cfg["match_type"].lower() in ["exact", "semantic", "graph"], "Matching type not supported. Choose from 'exact', 'semantic', or 'graph' using command line argument --match_type."

    # 1. Label matching
    if pred["name"] != gold["name"]:
        # labels must strictly match (exact matching)
        if match_cfg["match_type"].lower() == "exact":
            return False
        elif match_cfg["match_type"].lower() == "semantic":
            # check similarity between frame embeddings
            similarity_score = compute_similarity_frames(pred_name=pred["name"], gold_name=gold["name"],
                                                         embedding_model=match_cfg["embedding_model"],
                                                         frames=match_cfg["frames"])
            if similarity_score < match_cfg["semantic_match_threshold"]:
                return False
        elif match_cfg["match_type"].lower() == "graph":
            if is_frame:
                related, _ = compute_graph_relatedness(pred_name=pred["name"], gold_name=gold["name"], frames=match_cfg["frames"])
                if not related:
                    return False
            else:
                return False

    if span_match:
        # 2. Span matching
        pred_span = pred["target"] if is_frame else pred["span"]
        gold_span = gold["target"] if is_frame else gold["span"]
        assert gold_span, f'Span missing from gold item:\n{pred}\nvs\n{gold}'

        if not pred_span:
            warnings.warn(f'Span missing from predicted item (thus considered incorrect):\n{pred}')
            return False

        # relaxed span string match
        if span_overlap(pred_span, gold_span) < match_cfg["span_str_sim_threshold"]:
            return False

        # if unique span (based on gold span string), only string matters
        if span_counts.get(gold_span, 0) <= 1:
            return True

        # if multiple occurrences, check tolerant boundaries
        assert isinstance(pred["start"], int) and isinstance(gold["start"], int)
        return abs(pred["start"] - gold["start"]) <= match_cfg["span_idx_tolerance"] and abs(pred["end"] - gold["end"]) <= match_cfg["span_idx_tolerance"]

    else: # no span matching (only label matching, e.g. candidate frame eval) and labels do match
        return True

def match_list(pred_list:list[dict[str,Any]],
               gold_list:list[dict[str,Any]],
               span_counts:dict[str,int],
               match_cfg:dict[str,Any],
               is_frame:bool=False,
               span_match:bool=True)-> set[tuple[int,int]]:

    """
    Greedy matching to count correct generated frames.

    :param pred_list: List of predicted items.
    :param gold_list: List of gold items.
    :param span_counts: Span occurrence counts in text.
    :param match_cfg: Matching configuration parameters.
    :param is_frame: Whether matching frames (True) or frame elements (False).
    :param span_match: Whether matching spans as well (True) or only labels (False).

    :return: Set of matched index pairs (pred_idx, gold_idx).
    """

    matched = set()
    gold_matched = set()

    for p_idx, pred in enumerate(pred_list):
        for g_idx, gold in enumerate(gold_list):
            if g_idx in gold_matched:
                continue
            if match_item(pred=pred,
                          gold=gold,
                          span_counts=span_counts,
                          match_cfg=match_cfg,
                          is_frame=is_frame,
                          span_match=span_match):
                # print(f'Generated:{pred}')
                # print(f'Annotated:{gold}')
                # print()
                matched.add((p_idx, g_idx))
                gold_matched.add(g_idx)
                break

    return matched

def compute_metrics(correct:int, gold:int, pred:int)-> tuple[float,float,float]:

    """
    Compute precision, recall, and F1 score.

    Precision: from the extracted frames, which ones are annotated? (correct extracted frames (or frame elements) / total extracted frames (or frame elements))
    Recall: from the annotated frames, which ones are extracted? (correct extracted frames (or frame elements) / total annotated frames (or frame elements))
    F1-score: 2*P*R / P+R

    :param correct: Number of correct predictions.
    :param gold: Total gold items.
    :param pred: Total predicted items.

    :return: Precision, recall, and F1 score.
    """

    precision = correct / pred if pred else 0
    recall = correct / gold if gold else 0
    f1 = (2 * precision * recall / (precision + recall) if precision + recall > 0 else 0)

    return precision, recall, f1

def compute_scores(gold:list[dict[str,Any]],
                   pred:list[dict[str,Any]],
                   match_cfg:dict[str,Any],
                   fes_within_matched_frame:bool = True)-> tuple[dict[str, dict[str,float]], list[dict[Any,Any]], defaultdict(dict[str,dict[str,int]])]:

    """
    Compute evaluation scores for frames and frame elements.

    :param gold: Gold annotated frames.
    :param pred: Predicted frames.
    :param match_cfg: Matching configuration.
    :param fes_within_matched_frame: Whether to restrict FE scoring to matched frames.

    :return: Tuple:
    - scores (Dict): Precision, recall, and F1 scores for frames, frame elements, and both.
    - matches (List): Matching details for each text.
    - repeated_spans (Dict): Repeated spans frequencies for each text.
    """

    # global counters
    total_gold_frames, total_pred_frames, correct_frames = 0, 0, 0
    total_gold_fes, total_pred_fes, correct_fes = 0, 0, 0
    correct_full_frames = 0
    matches = []
    repeated_spans = defaultdict(dict)

    for gold_doc, pred_doc in zip(gold, pred):

        # if pred item (text) not in gold data for whatever reason, skip it
        if not gold_doc:
            continue

        text_idx_matches = dict()

        gold_frames = gold_doc.get("frames", [])
        pred_frames = pred_doc.get("frames", [])
        text = gold_doc.get("text", pred_doc.get("input_text", ""))

        span_counts = count_spans(text, gold_frames)
        for span, count in span_counts.items():
            repeated_spans[text][span] = count

        ###### Frame level
        total_gold_frames += len(gold_frames)
        total_pred_frames += len(pred_frames)
        frame_index_matches = match_list(pred_list=pred_frames,
                                         gold_list=gold_frames,
                                         span_counts=span_counts,
                                         match_cfg=match_cfg,
                                         is_frame=True)
        correct_frames += len(frame_index_matches)

        ###### Frame element level
        if not fes_within_matched_frame: # total is the total number of fes in and out of matched frames
            for g_frame in gold_frames:
                total_gold_fes += len(g_frame.get("frame_elements", []))
            for p_frame in pred_frames:
                total_pred_fes += len(p_frame.get("frame_elements", []))
        # match FEs only within matched frames
        for p_idx, p_frame in enumerate(pred_frames):
            for g_idx, g_frame in enumerate(gold_frames):
                if (p_idx, g_idx) not in frame_index_matches:
                    continue
                if fes_within_matched_frame: # total is the total number of fes in matched frames only
                    total_gold_fes += len(g_frame.get("frame_elements", []))
                    total_pred_fes += len(p_frame.get("frame_elements", []))
                pred_fes = p_frame.get("frame_elements", [])
                gold_fes = g_frame.get("frame_elements", [])
                fe_index_matches =  match_list(pred_list=pred_fes,
                                               gold_list=gold_fes,
                                               span_counts=span_counts,
                                               match_cfg=match_cfg,
                                               is_frame=False)
                correct_fes += len(fe_index_matches)
                text_idx_matches[(p_idx, g_idx)] = fe_index_matches

                ####### Full match (next to the frames matching, all frame elements match)
                if len(pred_fes) == len(gold_fes) and len(fe_index_matches) == len(gold_fes):
                    correct_full_frames += 1

        matches.append(text_idx_matches)

    # Compute precision, recall, and f1 scores
    frame_p, frame_r, frame_f1 = compute_metrics(correct=correct_frames, pred=total_pred_frames, gold=total_gold_frames)
    fe_p, fe_r, fe_f1 = compute_metrics(correct=correct_fes, pred=total_pred_fes, gold=total_gold_fes)
    full_p, full_r, full_f1 = compute_metrics(correct=correct_full_frames, pred=total_pred_frames, gold=total_gold_frames)

    scores = {"frame": {"precision": frame_p, "recall": frame_r, "f1": frame_f1},
              "frame_element": {"precision": fe_p, "recall": fe_r, "f1": fe_f1},
              "both": {"precision": full_p, "recall": full_r, "f1": full_f1}}

    return scores, matches, repeated_spans

def compute_scores_retrieval(gold:list[dict[str,Any]],
                           pred:list[dict[str,Any]],
                           match_cfg:dict[str,Any]):

    correct_candidate_frames = 0
    total_gold_frames = 0
    total_retrieved_frames = 0

    for gold_doc, pred_doc in zip(gold, pred):

        # if pred item (text) not in gold data for whatever reason, skip it
        if not gold_doc:
            continue

        candidate_frames = pred_doc.get("candidate_frames", [])
        gold_frames = gold_doc.get("frames", [])
        text = gold_doc.get("text", pred_doc.get("input_text", ""))
        span_counts = count_spans(text, gold_frames)

        total_gold_frames += len(gold_frames)
        total_retrieved_frames += len(candidate_frames)

        ###### Candidate frames (recall of RAG)
        frame_index_matches = match_list(pred_list=candidate_frames,
                                         gold_list=gold_frames,
                                         span_counts=span_counts,
                                         match_cfg=match_cfg,
                                         is_frame=True,
                                         span_match=False)
        correct_candidate_frames += len(frame_index_matches)

    c_frame_p, c_frame_r, c_frame_f1 = compute_metrics(correct=correct_candidate_frames, pred=total_retrieved_frames, gold=total_gold_frames)

    scores = {"candidate_frames": {"precision": c_frame_p, "recall": c_frame_r, "f1": c_frame_f1}}

    return scores


def find_graph_relations(gold, pred, frames):

    relations = []

    for doc_idx, (gold_doc, pred_doc) in enumerate(zip(gold, pred)):

        # if pred item (text) not in gold data for whatever reason, skip it
        if not gold_doc:
            continue

        text_eval = {"text": pred_doc["input_text"],
                     "related_frames": []}

        human_frames = gold_doc.get("frames", [])
        pred_frames = pred_doc.get("frames", [])

        if pred_frames and human_frames:
            for pred_frame in pred_frames:
                for gold_frame in human_frames:
                    related, relation = compute_graph_relatedness(pred_name=pred_frame.get("name"),
                                                            gold_name=gold_frame.get("name"),
                                                            frames=frames)
                    if related:
                        text_eval["related_frames"].append({'pred': pred_frame.get("name"),
                                                            'gold': gold_frame.get("name"),
                                                            'relation': relation})
        relations.append(text_eval)

    return relations

def find_hallucinations(pred, gold):

    pass

def get_eval_info(gold: list[dict[str,Any]], pred:list[dict[str,Any]], matches:list[dict[Any,Any]])-> list[dict[str,Any]]:

    """
    Generate detailed evaluation breakdown of correct, incorrect, and missed predictions.

    :param gold: Annotated frames.
    :param pred: Predicted frames.
    :param matches: Matching indices.

    :return: Detailed error analysis per text.
    """

    error_eval = []

    for doc_idx, (gold_doc, pred_doc) in enumerate(zip(gold, pred)):

        # if pred item (text) not in gold data for whatever reason, skip it
        if not gold_doc:
            continue

        text_eval = {"text": pred_doc["input_text"],
                     "correct_frames": [],
                     "incorrect_frames": [],
                     "missed_frames": []}

        human_frames = gold_doc.get("frames", [])
        pred_frames = pred_doc.get("frames", [])

        frame_matches = matches[doc_idx]
        matched_pred_frames, matched_gold_frames = set(), set() # 4 computing incorrect and missed frames

        ############ CORRECT FRAMES + FEs
        for (p_idx, g_idx), fe_matches in frame_matches.items():

            matched_pred_frames.add(p_idx)
            matched_gold_frames.add(g_idx)

            model_frame = pred_frames[p_idx]
            human_frame = human_frames[g_idx]

            frame_info = {
                "model": {k: v for k, v in model_frame.items() if k != 'frame_elements'},
                "human": {k: v for k, v in human_frame.items() if k != 'frame_elements'},
                "fes": {"correct_fes": [], "incorrect_fes": [], "missed_fes": []}
            }

            matched_pred_fes, matched_gold_fes = set(), set() # 4 computing incorrect and missed fes

            for p_fe_idx, g_fe_idx in fe_matches:

                matched_pred_fes.add(p_fe_idx)
                matched_gold_fes.add(g_fe_idx)

                fe_info = {"model": model_frame["frame_elements"][p_fe_idx],
                           "human": human_frame["frame_elements"][g_fe_idx]}
                frame_info["fes"]["correct_fes"].append(fe_info)

            # INCORRECT FEs (extra predictions)
            for p_fe_idx, p_fe in enumerate(model_frame.get("frame_elements", [])):
                if p_fe_idx not in matched_pred_fes:
                    fe_info = model_frame["frame_elements"][p_fe_idx]
                    frame_info["fes"]["incorrect_fes"].append(fe_info)

            # MISSED FEs (not predicted)
            for g_fe_idx, g_fe in enumerate(human_frame.get("frame_elements", [])):
                if g_fe_idx not in matched_gold_fes:
                    fe_info = human_frame["frame_elements"][g_fe_idx]
                    frame_info["fes"]["missed_fes"].append(fe_info)

            text_eval["correct_frames"].append(frame_info)

        # INCORRECT FRAMES (predicted but unmatched)
        for p_idx, p_frame in enumerate(pred_frames):
            if p_idx not in matched_pred_frames:
                model_frame = pred_frames[p_idx]
                frame_info = {k: v for k, v in model_frame.items() if k != 'frame_elements'}
                text_eval["incorrect_frames"].append(frame_info)

        # MISSED FRAMES (gold not predicted)
        for g_idx, g_frame in enumerate(human_frames):
            if g_idx not in matched_gold_frames:
                gold_frame = human_frames[g_idx]
                frame_info = {k: v for k, v in gold_frame.items() if k != 'frame_elements'}
                text_eval["missed_frames"].append(frame_info)

        error_eval.append(text_eval)

    return error_eval

def evaluate(args:Any, annotated_data:dict[str,Any], generated_data:dict[str,Any])->None:

    """
    Run full evaluation pipeline: scoring + error analysis.

    :param args: Argument object containing evaluation configs.
    :param annotated_data: Annotated dataset.
    :param generated_data: Model output dataset.

    :raise ValueError: if required data fields are missing.
    """

    match_cfg = {"match_type": args.match_type,
                 "semantic_match_threshold": .5,
                 "graph_match_threshold": .5,
                 "span_str_sim_threshold": .3,
                 "span_idx_tolerance": 3,
                 "embedding_model": None,
                 "frames": None}

    # if matching semantically similar frames
    if args.match_type.lower() == "semantic":
        match_cfg["embedding_model"] = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2",
                                                model_kwargs={"device": "cpu"},
                                                encode_kwargs={"batch_size": 16,
                                                               "normalize_embeddings": True})
    if args.match_type.lower() in ["semantic", "graph"]:
        match_cfg["frames"] = parse_frames(frame_dir=args.frame_dir, output_filepath=args.frame_output_filepath)

    # get generated and annotated frames
    gold_docs = annotated_data.get("texts")
    pred_docs = generated_data.get("responses")

    if gold_docs is None:
        raise ValueError(f'No annotated texts under "texts" key. Make sure annotated texts have the correct format.')

    if pred_docs is None:
        raise ValueError(f'No generated texts under "responses" key. Make sure generated texts have the correct format.')

    if len(gold_docs) == 0:
        raise ValueError(f'No annotated texts found. Make sure annotated texts contain at least one text.')

    if len(pred_docs) == 0:
        raise ValueError(f'No generated texts found. Make sure generated texts contain at least one text.')

    print(f'Running evaluation...')

    # select texts in gold data that are also in pred data
    selected_gold_docs = []
    for item in pred_docs:
        gold_index = [i for i, gold_item in enumerate(gold_docs) if item["input_text"] == gold_item["text"]]
        if gold_index:
            selected_gold_docs.append(gold_docs[gold_index[0]])
        else:
            selected_gold_docs.append(dict())
            warnings.warn(f'Text "{item["input_text"]}" not found in gold data. Make sure model input texts are in the gold data for evaluation.')
    assert len(selected_gold_docs) == len(pred_docs)

    # compute scores
    scores = compute_scores_retrieval(gold=selected_gold_docs, pred=pred_docs, match_cfg=match_cfg)

    # handle output directory
    if not args.output_eval_dir:
        args.output_eval_dir = Path.cwd()
    Path(args.output_eval_dir).mkdir(exist_ok=True)

    if not args.retrieval_only:

        _scores, matches, repeated_spans = compute_scores(gold=selected_gold_docs, pred=pred_docs, match_cfg=match_cfg)
        scores.update(_scores)

        # get frame relations between predicted and gold frames
        relations = find_graph_relations(gold=selected_gold_docs, pred=pred_docs, frames=match_cfg["frames"])

        # get detailed info on correct and incorrect frames and fes
        error_eval = get_eval_info(gold=selected_gold_docs, pred=pred_docs, matches=matches)

        # save
        counts_filename = f"count_spans_{args.model_name}.json"
        error_filename = f"error_eval_{args.model_name}.json"
        relations_filename = f"relations_{args.model_name}.json"

        if args.match_type.lower() == "exact":
            counts_filename = counts_filename.replace(".json", "_exact_match.json")
            error_filename = error_filename.replace(".json", "_exact_match.json")
        elif args.match_type.lower() == "semantic":
            counts_filename = counts_filename.replace(".json", "_semantic_match.json")
            error_filename = error_filename.replace(".json", "_semantic_match.json")
        elif args.match_type.lower() == "graph":
            counts_filename = counts_filename.replace(".json", "_graph_match.json")
            error_filename = error_filename.replace(".json", "_graph_match.json")

        with open(f'{args.output_eval_dir}/{counts_filename}', "w") as f:
            json.dump(repeated_spans, f, indent=2, ensure_ascii=False)
        with open(f'{args.output_eval_dir}/{error_filename}', 'w') as f:
            json.dump(error_eval, f, indent=4)
        with open(f'{args.output_eval_dir}/{relations_filename}', 'w') as f:
            json.dump(relations, f, indent=2)

    # save scores
    filename = f"eval_scores_{args.model_name}.json"
    if args.retrieval_only:
        filename = filename.replace(f'{args.model_name}', 'retrieval-only')
    if args.match_type.lower() == "exact":
        filename = filename.replace(".json","_exact_match.json")
    elif args.match_type.lower() == "semantic":
        filename = filename.replace(".json", "_semantic_match.json")
    elif args.match_type.lower() == "graph":
        filename = filename.replace(".json", "_graph_match.json")
    with open(f'{args.output_eval_dir}/{filename}', 'w') as f:
        json.dump(scores, f, indent=2)

    print('Evaluation done.')