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

# %% auto 0
__all__ = ['log', 'Entity', 'Relation', 'Triplet', 'DEFAULT_SYSTEM_PROMPT_TEMPLATE', 'DEFAULT_RELATION_SET_PROMPT_TEMPLATE',
           'DEFAULT_FEW_SHOT_EXAMPLES_PROMPT_TEMPLATE', 'DEFAULT_SYSTEM_PROMPT_TEMPLATE2',
           'DEFAULT_RELATION_SET_PROMPT_TEMPLATE2', 'evaluate_joint_er_extraction', 'evaluate_joint_er_extractions',
           'parse_triplet_strings', 'parse_triplets', 'format_triplets', 'format_few_shot_example',
           'format_few_shot_examples', 'ERX2AlpacaFormatter', 'ERX2ChatFormatter', 'evaluate_pipe_jer',
           'evaluate_model_jer']

# %% ../../../nbs/ml.kg.cons.ipynb 3
import random
from dataclasses import dataclass
from typing import TypeAlias, Iterable, Callable, Any, Generator
import numpy as np

from ...logging import get_logger

log = get_logger(__name__)

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

# %% ../../../nbs/ml.kg.cons.ipynb 6
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 9
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[Triplet]:
    return [tuple(triplet_string.split(delimiter)) for triplet_string in parse_triplet_strings(text, delimiter=delimiter)]

# %% ../../../nbs/ml.kg.cons.ipynb 12
def format_triplets(triplets: Iterable[str]) -> str:
    return '\n'.join(triplets)

def format_few_shot_example(example, text_prefix="# Text\n", triplets_prefix="# Triplets\n"):
    text = example['text']
    triplets = format_triplets(example['triplets'])
    return f"{text_prefix}{text}\n{triplets_prefix}{triplets}"

def format_few_shot_examples(examples):
    return "\n\n".join([format_few_shot_example(example) for example in examples])

# %% ../../../nbs/ml.kg.cons.ipynb 13
DEFAULT_SYSTEM_PROMPT_TEMPLATE = """You are a helpful assistant that extracts entity-relation-entity triplets from given text.
{relation_set_prompt}
{few_shot_prompt}
""".strip()

DEFAULT_RELATION_SET_PROMPT_TEMPLATE = """Here are the list of relations that you can use:
{relation_set}
"""

DEFAULT_FEW_SHOT_EXAMPLES_PROMPT_TEMPLATE = """Use the same format for triplets as in the examples provided below.
{few_shot_examples}
"""

@dataclass
class ERX2AlpacaFormatter:
    system_prompt_template: str = DEFAULT_SYSTEM_PROMPT_TEMPLATE
    relation_set_prompt_template: str = DEFAULT_RELATION_SET_PROMPT_TEMPLATE
    relation_set: set|None = None
    few_shot_examples_prompt_template: str = DEFAULT_FEW_SHOT_EXAMPLES_PROMPT_TEMPLATE
    few_shot_examples: list[dict]|None = None
    n_few_shot_examples: int = 3

    def __post_init__(self):
        if self.relation_set:
            self.relation_set = sorted(self.relation_set)
        if self.few_shot_examples is not None:
            self.few_shot_examples = list(self.few_shot_examples)

    def format(self, example: dict):
        instruction = self.make_system_prompt()
        input = example['text']
        output = '\n'.join(example['triplets'])
        return {
            'instruction': instruction.strip(),
            'input': input.strip(),
            'output': output.strip(),
        }

    def make_system_prompt(self) -> str:
        rsp = self.relation_set_prompt_template.format(relation_set=','.join(self.relation_set)) if self.relation_set else ""
        fsp = self.few_shot_examples_prompt_template.format(few_shot_examples=format_few_shot_examples(self._choose_few_shot_examples())) if self.few_shot_examples else ""
        return self.system_prompt_template.format(relation_set_prompt=rsp, few_shot_prompt=fsp)

    def _choose_few_shot_examples(self) -> list[dict]:
        if len(self.few_shot_examples) <= self.n_few_shot_examples:
            return self.few_shot_examples
        else:
            return random.sample(self.few_shot_examples, k=self.n_few_shot_examples)

# %% ../../../nbs/ml.kg.cons.ipynb 16
DEFAULT_SYSTEM_PROMPT_TEMPLATE2 = """You are a helpful assistant that extracts entity-relation-entity triplets from given text. Use '|' as delimiter and provide one triplet per line.
{relation_set_prompt}
""".strip()

DEFAULT_RELATION_SET_PROMPT_TEMPLATE2 = """Here are the list of relations that you can use:
{relation_set}
""".strip()

@dataclass
class ERX2ChatFormatter:
    system_prompt_template: str = DEFAULT_SYSTEM_PROMPT_TEMPLATE2
    relation_set_prompt_template: str = DEFAULT_RELATION_SET_PROMPT_TEMPLATE2
    relation_set: set|None = None
    few_shot_examples: list[dict]|None = None
    n_few_shot_examples: int = 3

    def __post_init__(self):
        if self.relation_set:
            self.relation_set = sorted(self.relation_set)
        if self.few_shot_examples is not None:
            self.few_shot_examples = list(self.few_shot_examples)

    def format(self, example: dict):
        messages = [
            {"role": "system", "content": self.make_system_message()},
            *list(self.make_messages(*self._choose_few_shot_examples(), example)),
        ]
        return {'conversations': messages}

    def make_system_message(self) -> str:
        rsp = self.relation_set_prompt_template.format(relation_set=','.join(self.relation_set)) if self.relation_set else ""
        return self.system_prompt_template.format(relation_set_prompt=rsp)

    def make_messages(self, *examples) -> Generator[dict, None, None]:
        for example in examples:
            yield {"role": "user", "content": example["text"]}
            yield {"role": "assistant", "content": format_triplets(example["triplets"])}

    def _choose_few_shot_examples(self) -> list[dict]:
        if len(self.few_shot_examples) <= self.n_few_shot_examples:
            return self.few_shot_examples
        else:
            return random.sample(self.few_shot_examples, k=self.n_few_shot_examples)

# %% ../../../nbs/ml.kg.cons.ipynb 19
def evaluate_pipe_jer(dataset, pipe):
    import evaluate

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

    results = pipe(dataset["input"])
    generations = [result[0]["generated_text"] for result in results]
    predictions = [parse_triplet_strings(text.strip()) for text in generations]
    references = [parse_triplet_strings(text.strip()) 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)
