feat: add multi answer qa generation

Author: ChenZiHong-Gavin

examples/generate/generate_multi_answer_qa/README.md

@@ -0,0 +1,3 @@
+# Generate Multi-Answer QAs
+
+Multi-answer question answering (QA) involves generating questions that can have multiple valid answers. This is particularly useful in educational settings, surveys, and research where diverse perspectives are valuable. 

examples/generate/generate_multi_answer_qa/generate_multi_answer.sh

@@ -0,0 +1,2 @@
+python3 -m graphgen.run \
+--config_file examples/generate/generate_multi_answer_qa/multi_answer_config.yaml

examples/generate/generate_multi_answer_qa/multi_answer_config.yaml

@@ -0,0 +1,80 @@
+global_params:
+  working_dir: cache
+  graph_backend: kuzu # graph database backend, support: kuzu, networkx
+  kv_backend: rocksdb # key-value store backend, support: rocksdb, json_kv
+
+nodes:
+  - id: read_files # id is unique in the pipeline, and can be referenced by other steps
+    op_name: read
+    type: source
+    dependencies: []
+    params:
+        input_path:
+          - examples/input_examples/jsonl_demo.jsonl # input file path, support json, jsonl, txt, pdf. See examples/input_examples for examples
+
+  - id: chunk_documents
+    op_name: chunk
+    type: map_batch
+    dependencies:
+      - read_files
+    execution_params:
+      replicas: 4
+    params:
+        chunk_size: 1024 # chunk size for text splitting
+        chunk_overlap: 100 # chunk overlap for text splitting
+
+  - id: build_kg
+    op_name: build_kg
+    type: map_batch
+    dependencies:
+      - chunk_documents
+    execution_params:
+      replicas: 1
+      batch_size: 128
+
+  - id: quiz
+    op_name: quiz
+    type: map_batch
+    dependencies:
+      - build_kg
+    execution_params:
+      replicas: 1
+      batch_size: 128
+    params:
+      quiz_samples: 2 # number of quiz samples to generate
+
+  - id: judge
+    op_name: judge
+    type: map_batch
+    dependencies:
+      - quiz
+    execution_params:
+      replicas: 1
+      batch_size: 128
+
+  - id: partition
+    op_name: partition
+    type: aggregate
+    dependencies:
+      - judge
+    params:
+      method: ece # ece is a custom partition method based on comprehension loss
+      method_params:
+        max_units_per_community: 20 # max nodes and edges per community
+        min_units_per_community: 5 # min nodes and edges per community
+        max_tokens_per_community: 10240 # max tokens per community
+        unit_sampling: max_loss # unit sampling strategy, support: random, max_loss, min_loss
+
+  - id: generate
+    op_name: generate
+    type: map_batch
+    dependencies:
+      - partition
+    execution_params:
+      replicas: 1
+      batch_size: 128
+    save_output: true # save output
+    params:
+      method: multi_answer
+      num_of_questions: 5
+      data_format: Alpaca # Alpaca, Sharegpt, ChatML

graphgen/models/__init__.py

@@ -11,6 +11,7 @@
     AggregatedGenerator,
     AtomicGenerator,
     CoTGenerator,
+    MultiAnswerGenerator,
     MultiChoiceGenerator,
     MultiHopGenerator,
     QuizGenerator,

graphgen/models/generator/__init__.py

@@ -1,6 +1,7 @@
 from .aggregated_generator import AggregatedGenerator
 from .atomic_generator import AtomicGenerator
 from .cot_generator import CoTGenerator
+from .multi_answer_generator import MultiAnswerGenerator
 from .multi_choice_generator import MultiChoiceGenerator
 from .multi_hop_generator import MultiHopGenerator
 from .quiz_generator import QuizGenerator

graphgen/models/generator/multi_answer_generator.py

@@ -0,0 +1,121 @@
+import re
+from typing import Any
+
+from graphgen.bases import BaseGenerator
+from graphgen.templates import MAQ_GENERATION_PROMPT
+from graphgen.utils import compute_content_hash, detect_main_language, logger
+
+
+class MultiAnswerGenerator(BaseGenerator):
+    def __init__(self, llm_client, num_of_questions) -> None:
+        super().__init__(llm_client)
+        self.num_of_questions = num_of_questions
+
+    @staticmethod
+    def parse_response(response: str) -> Any:
+        """
+        Parse multiple-answer QA pairs from the LLM response.
+        Each QA pair contains question text, four options, and the correct answers (one or more).
+
+        :param response: The LLM response containing XML-formatted QA pairs
+        :return: Dictionary mapping question hash to question data, where each
+                 value is a dict with "question", "options", "answer", and
+                 "correct_answer_text" keys
+        """
+        qa_pairs = {}
+
+        # Extract all QA pair blocks
+        qa_blocks = re.findall(r"<qa_pair>(.*?)</qa_pair>", response, re.DOTALL)
+
+        if not qa_blocks:
+            logger.warning("No QA pairs found in response: %s", response)
+            return {}
+
+        for block in qa_blocks:
+            # Extract and clean question text
+            q_match = re.search(r"<question>(.*?)</question>", block, re.DOTALL)
+            if not q_match:
+                logger.warning("Failed to parse question from block: %s", block)
+                continue
+            question = q_match.group(1).strip().strip('"').strip("'")
+
+            # Extract and parse options (A, B, C, D)
+            opt_match = re.search(r"<options>(.*?)</options>", block, re.DOTALL)
+            if not opt_match:
+                logger.warning("Failed to parse options from block: %s", block)
+                continue
+
+            options = {}
+            options_text = opt_match.group(1).strip()
+            for line in options_text.split("\n"):
+                line = line.strip()
+                if not line:
+                    continue
+                # Match patterns like "A. text" or "B. text"
+                if m := re.match(r"^([A-Z])[.\s]\s*(.*)$", line):
+                    letter, text = m.groups()
+                    options[letter] = text.strip()
+
+            # Extract and validate answer
+            ans_match = re.search(r"<answer>(.*?)</answer>", block, re.DOTALL)
+            if not ans_match:
+                logger.warning("Failed to parse answer from block: %s", block)
+                continue
+            answer_text = ans_match.group(1).strip().strip('"').strip("'")
+            answers = [ans.strip().upper() for ans in answer_text.split(",")]
+
+            answers = [ans for ans in answers if ans]
+            invalid_answers = [ans for ans in answers if ans not in options]
+            if invalid_answers:
+                logger.warning(
+                    "Answers %s not found in options: %s",
+                    invalid_answers,
+                    list(options.keys()),
+                )
+                continue
+
+            # Ensure at least one valid answer
+            if len(answers) == 0:
+                logger.warning("No valid answers found in: %s", answer_text)
+                continue
+
+            # Build result entry with question hash as key
+            question_hash = compute_content_hash(question)
+            qa_pairs[question_hash] = {
+                "question": question,
+                "options": options,  # Dict like {"A": "text", "B": "text", ...}
+                "answer": ", ".join(answers),
+            }
+
+            logger.debug("Successfully parsed MCQ: %s", question[:50])
+
+        if not qa_pairs:
+            logger.error("Failed to parse any valid MAQ pairs from response")
+
+        return qa_pairs
+
+    # pylint: disable=W0221
+    def build_prompt(
+        self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
+    ) -> str:
+        nodes, edges = batch
+        entities_str = "\n".join(
+            [
+                f"{index + 1}. {node[0]}: {node[1]['description']}"
+                for index, node in enumerate(nodes)
+            ]
+        )
+
+        relationships_str = "\n".join(
+            [
+                f"{index + 1}. {edge[0]} -- {edge[1]}: {edge[2]['description']}"
+                for index, edge in enumerate(edges)
+            ]
+        )
+        context = entities_str + "\n" + relationships_str
+        language = detect_main_language(entities_str + relationships_str)
+        prompt = MAQ_GENERATION_PROMPT[language].format(
+            context=context,
+            num_of_questions=self.num_of_questions,
+        )
+        return prompt

graphgen/models/generator/multi_choice_generator.py

@@ -83,9 +83,6 @@ def parse_response(response: str) -> Any:
                 "question": question,
                 "options": options,  # Dict like {"A": "text", "B": "text", ...}
                 "answer": answer,  # Single letter: "A", "B", "C", or "D"
-                "correct_answer_text": options[
-                    answer
-                ],  # The actual text of correct answer
             }
 
             logger.debug("Successfully parsed MCQ: %s", question[:50])

graphgen/operators/generate/generate_service.py

@@ -50,6 +50,13 @@ def __init__(
                 self.llm_client,
                 num_of_questions=generate_kwargs.get("num_of_questions", 5),
             )
+        elif self.method == "multi_answer":
+            from graphgen.models import MultiAnswerGenerator
+
+            self.generator = MultiAnswerGenerator(
+                self.llm_client,
+                num_of_questions=generate_kwargs.get("num_of_questions", 3),
+            )
         else:
             raise ValueError(f"Unsupported generation mode: {method}")
 

graphgen/templates/__init__.py

@@ -6,6 +6,7 @@
     AGGREGATED_GENERATION_PROMPT,
     ATOMIC_GENERATION_PROMPT,
     COT_GENERATION_PROMPT,
+    MAQ_GENERATION_PROMPT,
     MCQ_GENERATION_PROMPT,
     MULTI_HOP_GENERATION_PROMPT,
     VQA_GENERATION_PROMPT,

graphgen/templates/generation/__init__.py

@@ -1,6 +1,7 @@
 from .aggregated_generation import AGGREGATED_GENERATION_PROMPT
 from .atomic_generation import ATOMIC_GENERATION_PROMPT
 from .cot_generation import COT_GENERATION_PROMPT
+from .multi_answer_generation import MAQ_GENERATION_PROMPT
 from .multi_choice_generation import MCQ_GENERATION_PROMPT
 from .multi_hop_generation import MULTI_HOP_GENERATION_PROMPT
 from .vqa_generation import VQA_GENERATION_PROMPT