# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/jerx.eval.ipynb.

# %% auto 0
__all__ = ['log', 'evaluate_joint_er_extraction', 'evaluate_joint_er_extractions', 'evaluate_pipe_jer', 'evaluate_model_jer']

# %% ../../nbs/jerx.eval.ipynb 3
from typing import Iterable, Callable, Any, Generator
import numpy as np
from .utils import Triplet, parse_triplet_strings
from ..logging import get_logger

log = get_logger(__name__)

# %% ../../nbs/jerx.eval.ipynb 4
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/jerx.eval.ipynb 7
def evaluate_pipe_jer(dataset, pipe):
    import evaluate

    log.info(f"Evaluating model for JER on dataset with {len(dataset)} samples.")

    tokenizer = pipe.tokenizer

    def _clean(text):
        return text.replace(tokenizer.special_tokens_map["eos_token"], "").strip()

    results = pipe(dataset["input"])
    generations = [result[0]["generated_text"] for result in results]
    predictions = [parse_triplet_strings(_clean(text)) for text in generations]
    references = [parse_triplet_strings(_clean(text)) for text in dataset["output"]]

    dataf = dataset.to_pandas()
    dataf["generation"] = generations
    dataf["prediction"] = predictions
    dataf["reference"] = references

    metric = evaluate.load("bdsaglam/jer")
    scores = metric.compute(predictions=predictions, references=references)

    return scores, dataf


def evaluate_model_jer(
    dataset,
    *,
    response_template: str,
    tokenizer,
    model,
    max_new_tokens=256,
    batch_size=4,
    **kwargs,
):
    assert len(dataset) > 0, "Dataset is empty!"

    def extract_input_output(example):
        input, output = example["text"].rsplit(response_template, 1)
        input += response_template
        return {"input": input, "output": output}

    dataset = dataset.map(extract_input_output)

    # setup generation pipeline
    from transformers import pipeline

    pipe = pipeline(
        task="text-generation",
        model=model,
        tokenizer=tokenizer,
        max_new_tokens=max_new_tokens,
        batch_size=batch_size,
        return_full_text=False,
        **kwargs,
    )

    return evaluate_pipe_jer(dataset, pipe)
