graphgen/operators/resolute_coreference.py

from typing import List
from graphgen.models import Chunk
from graphgen.models import OpenAIModel
from graphgen.templates import COREFERENCE_RESOLUTION_TEMPLATE
from graphgen.utils import detect_main_language

async def resolute_coreference(
        llm_client: OpenAIModel,
        chunks: List[Chunk]) -> List[Chunk]:
    """
    Resolute conference

    :param llm_client: LLM model
    :param chunks: List of chunks
    :return: List of chunks
    """

    if len(chunks) == 0:
        return chunks

    results = [chunks[0]]

    for _, chunk in enumerate(chunks[1:]):
        language = detect_main_language(chunk.content)
        result = await llm_client.generate_answer(
            COREFERENCE_RESOLUTION_TEMPLATE[language].format(
                reference = results[0].content,
                input_sentence = chunk.content
            )
        )
        results.append(Chunk(id=chunk.id, content=result))

    return results