|
| 1 | +from typing import Any |
| 2 | + |
| 3 | +from graphgen.bases import BaseGenerator |
| 4 | +from graphgen.templates import DESCRIPTION_REPHRASING_PROMPT |
| 5 | +from graphgen.utils import detect_main_language, logger |
| 6 | + |
| 7 | + |
| 8 | +class QuizGenerator(BaseGenerator): |
| 9 | + """ |
| 10 | + Quiz Generator rephrases given descriptions to create quiz questions. |
| 11 | + """ |
| 12 | + |
| 13 | + @staticmethod |
| 14 | + def build_prompt( |
| 15 | + batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]] |
| 16 | + ) -> str: |
| 17 | + """ |
| 18 | + Build prompt for rephrasing the description. |
| 19 | + :param batch: A tuple containing (nodes, edges) where nodes/edges |
| 20 | + contain description information |
| 21 | + :return: Prompt string |
| 22 | + """ |
| 23 | + # Extract description from batch |
| 24 | + # For quiz generator, we expect a special format where |
| 25 | + # the description is passed as the first node's description |
| 26 | + nodes, edges = batch |
| 27 | + if nodes: |
| 28 | + description = nodes[0][1].get("description", "") |
| 29 | + template_type = nodes[0][1].get("template_type", "TEMPLATE") |
| 30 | + elif edges: |
| 31 | + description = edges[0][2].get("description", "") |
| 32 | + template_type = edges[0][2].get("template_type", "TEMPLATE") |
| 33 | + else: |
| 34 | + raise ValueError("Batch must contain at least one node or edge with description") |
| 35 | + |
| 36 | + return QuizGenerator.build_prompt_for_description(description, template_type) |
| 37 | + |
| 38 | + @staticmethod |
| 39 | + def build_prompt_for_description(description: str, template_type: str = "TEMPLATE") -> str: |
| 40 | + """ |
| 41 | + Build prompt for rephrasing a single description. |
| 42 | + :param description: The description to rephrase |
| 43 | + :param template_type: Either "TEMPLATE" (same meaning) or "ANTI_TEMPLATE" (opposite meaning) |
| 44 | + :return: Prompt string |
| 45 | + """ |
| 46 | + language = detect_main_language(description) |
| 47 | + prompt = DESCRIPTION_REPHRASING_PROMPT[language][template_type].format( |
| 48 | + input_sentence=description |
| 49 | + ) |
| 50 | + return prompt |
| 51 | + |
| 52 | + @staticmethod |
| 53 | + def parse_rephrased_text(response: str) -> str: |
| 54 | + """ |
| 55 | + Parse the rephrased text from the response. |
| 56 | + :param response: |
| 57 | + :return: |
| 58 | + """ |
| 59 | + rephrased_text = response.strip().strip('"') |
| 60 | + logger.debug("Rephrased Text: %s", rephrased_text) |
| 61 | + return rephrased_text |
| 62 | + |
| 63 | + @staticmethod |
| 64 | + def parse_response(response: str) -> Any: |
| 65 | + """ |
| 66 | + Parse the LLM response. For quiz generator, this returns the rephrased text. |
| 67 | + :param response: LLM response |
| 68 | + :return: Rephrased text |
| 69 | + """ |
| 70 | + return QuizGenerator.parse_rephrased_text(response) |
0 commit comments