# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/ml.kg.cons.ipynb.

# %% auto 0
__all__ = ['Entity', 'Relation', 'Triplet', 'evaluate_joint_er_extraction', 'evaluate_joint_er_extractions',
           'parse_triplet_strings', 'parse_triplets']

# %% ../../../nbs/ml.kg.cons.ipynb 3
from typing import TypeAlias, Iterable, List, Set, Tuple, Callable, Any
import numpy as np

# %% ../../../nbs/ml.kg.cons.ipynb 4
Entity: TypeAlias = tuple[str, str]
Relation: TypeAlias = str
Triplet: TypeAlias = tuple[Entity, Relation, Entity]

# %% ../../../nbs/ml.kg.cons.ipynb 5
def evaluate_joint_er_extraction(*, reference: Iterable[Triplet], prediction: Iterable[Triplet]):
    """
    Example: [(('John', 'PERSON'), 'works_at', ('Google', 'ORG'))]
    """

    reference_set = set(reference)
    prediction_set = set(prediction)
    assert len(reference) == len(reference_set), "Duplicates found in references"

    TP = len(reference_set & prediction_set)
    FP = len(prediction_set - reference_set)
    FN = len(reference_set - prediction_set)
    
    # Calculate metrics
    precision = TP / (TP + FP) if TP + FP > 0 else 0
    recall = TP / (TP + FN) if TP + FN > 0 else 0
    f1_score = 2 * (precision * recall) / (precision + recall) if precision + recall > 0 else 0
    
    return {
        'precision': precision,
        'recall': recall,
        'f1': f1_score
    }

def evaluate_joint_er_extractions(*, references: Iterable[Iterable[Triplet]], predictions: Iterable[Iterable[Triplet]]):
    score_dicts = [
        evaluate_joint_er_extraction(reference=reference, prediction=prediction) 
        for reference, prediction in zip(references, predictions)
    ]
    return {('mean_' + key): np.mean([scores[key] for scores in score_dicts]) for key in score_dicts[0].keys()}

# %% ../../../nbs/ml.kg.cons.ipynb 8
def parse_triplet_strings(text: str, delimiter: str=" | ") -> List[str]:
    return [line for line in text.splitlines() if line and line.count(delimiter) == 2]

def parse_triplets(text: str, delimiter: str=" | ") -> List[Tuple]:
    return [tuple(triplet_string.split(delimiter)) for triplet_string in parse_triplet_strings(text, delimiter=delimiter)]
