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

# %% auto 0
__all__ = ['log', 'parse_triplet_response', 'make_kg_triplet_extract_fn']

# %% ../../../nbs/jerx.offline.llm.ipynb 3
from pathlib import Path
import pandas as pd
from ..utils import parse_triplets
from ...logging import get_logger

log = get_logger(__name__)

# %% ../../../nbs/jerx.offline.llm.ipynb 5
def parse_triplet_response(response: str, *args, **kwargs) -> list[tuple[str, str, str]]:
    triplets = parse_triplets(response.strip())
    return [(e1, rel, e2) if e1 != e2 else (e1, rel, e2 + "(obj)") for e1, rel, e2 in triplets]


def make_kg_triplet_extract_fn(inference_cache_filepath: Path):
    df = pd.read_json(inference_cache_filepath, lines=True)
    mapping = {text: generation for text, generation in zip(df["text"], df["generation"])}
    del df

    def extract_kg_triplets(text: str) -> list[tuple[str, str, str]]:
        return parse_triplet_response(mapping[text])

    return extract_kg_triplets
