base_generator.py

from abc import ABC, abstractmethod
from typing import Any

from graphgen.bases.base_llm_wrapper import BaseLLMWrapper


class BaseGenerator(ABC):
    """
    Generate QAs based on given prompts.
    """

    def __init__(self, llm_client: BaseLLMWrapper):
        self.llm_client = llm_client

    @staticmethod
    @abstractmethod
    def build_prompt(
        batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
    ) -> str:
        """Build prompt for LLM based on the given batch"""

    @staticmethod
    @abstractmethod
    def parse_response(response: str) -> list[dict]:
        """Parse the LLM response and return the generated QAs"""

    async def generate(
        self,
        batch: tuple[
            list[tuple[str, dict]], list[tuple[Any, Any, dict] | tuple[Any, Any, Any]]
        ],
    ) -> list[dict]:
        """
        Generate QAs based on a given batch.
        :param batch
        :return: QA pairs
        """
        prompt = self.build_prompt(batch)
        response = await self.llm_client.generate_answer(prompt)
        qa_pairs = self.parse_response(response)  # generate one or more QA pairs
        return qa_pairs

    @staticmethod
    def format_generation_results(
        result: dict, output_data_format: str
    ) -> dict[str, Any]:
        question = result.get("question", "")
        answer = result.get("answer", "")
        if "options" in result and result["options"]:
            options = result["options"]
            options_str = "\n".join(
                [f"{key}. {options[key]}" for key in sorted(options.keys())]
            )
            question += f"\nOptions:\n{options_str}"

        if output_data_format == "Alpaca":
            return {
                "instruction": question,
                "input": "",
                "output": answer,
            }

        if output_data_format == "Sharegpt":
            return {
                "conversations": [
                    {"from": "human", "value": question},
                    {"from": "gpt", "value": answer},
                ]
            }
        if output_data_format == "ChatML":
            return {
                "messages": [
                    {"role": "user", "content": question},
                    {"role": "assistant", "content": answer},
                ]
            }
        raise ValueError(f"Unknown output data format: {output_data_format}")