wip: trace generated data

Author: ChenZiHong-Gavin

graphgen/bases/base_generator.py

@@ -21,72 +21,67 @@ def build_prompt(
 
     @staticmethod
     @abstractmethod
-    def parse_response(response: str) -> Any:
+    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]]
         ],
-    ) -> dict[str, Any]:
+    ) -> list[dict]:
         """
         Generate QAs based on a given batch.
         :param batch
         :return: QA pairs
         """
-        result = {}
         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
-        result.update(qa_pairs)
-        return result
+        return qa_pairs
 
     @staticmethod
     def format_generation_results(
         results: list[dict], output_data_format: str
     ) -> list[dict[str, Any]]:
 
         flat_results = []
-        for item in results:
-            for _, qa_data in item.items():
-                question = qa_data.get("question", "")
-                answer = qa_data.get("answer", "")
-                if "options" in qa_data and qa_data["options"]:
-                    options = qa_data["options"]
-                    options_str = "\n".join(
-                        [f"{key}. {options[key]}" for key in sorted(options.keys())]
-                    )
-                    question += f"\nOptions:\n{options_str}"
+        for qa_data in results:
+            question = qa_data.get("question", "")
+            answer = qa_data.get("answer", "")
+            if "options" in qa_data and qa_data["options"]:
+                options = qa_data["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":
-                    flat_results.append(
-                        {
-                            "instruction": question,
-                            "input": "",
-                            "output": answer,
-                        }
-                    )
-                elif output_data_format == "Sharegpt":
-                    flat_results.append(
-                        {
-                            "conversations": [
-                                {"from": "human", "value": question},
-                                {"from": "gpt", "value": answer},
-                            ]
-                        }
-                    )
-                elif output_data_format == "ChatML":
-                    flat_results.append(
-                        {
-                            "messages": [
-                                {"role": "user", "content": question},
-                                {"role": "assistant", "content": answer},
-                            ]
-                        }
-                    )
-                else:
-                    raise ValueError(
-                        f"Unknown output data format: {output_data_format}"
-                    )
+            if output_data_format == "Alpaca":
+                flat_results.append(
+                    {
+                        "instruction": question,
+                        "input": "",
+                        "output": answer,
+                    }
+                )
+            elif output_data_format == "Sharegpt":
+                flat_results.append(
+                    {
+                        "conversations": [
+                            {"from": "human", "value": question},
+                            {"from": "gpt", "value": answer},
+                        ]
+                    }
+                )
+            elif output_data_format == "ChatML":
+                flat_results.append(
+                    {
+                        "messages": [
+                            {"role": "user", "content": question},
+                            {"role": "assistant", "content": answer},
+                        ]
+                    }
+                )
+            else:
+                raise ValueError(f"Unknown output data format: {output_data_format}")
         return flat_results

graphgen/models/generator/aggregated_generator.py

@@ -3,7 +3,7 @@
 
 from graphgen.bases import BaseGenerator
 from graphgen.templates import AGGREGATED_GENERATION_PROMPT
-from graphgen.utils import compute_content_hash, detect_main_language, logger
+from graphgen.utils import detect_main_language, logger
 
 
 class AggregatedGenerator(BaseGenerator):
@@ -101,30 +101,26 @@ async def generate(
         batch: tuple[
             list[tuple[str, dict]], list[tuple[Any, Any, dict] | tuple[Any, Any, Any]]
         ],
-    ) -> dict[str, Any]:
+    ) -> list[dict]:
         """
         Generate QAs based on a given batch.
         :param batch
         :return: QA pairs
         """
-        result = {}
         rephrasing_prompt = self.build_prompt(batch)
         response = await self.llm_client.generate_answer(rephrasing_prompt)
         context = self.parse_rephrased_text(response)
         if not context:
-            return result
+            return []
         question_generation_prompt = self._build_prompt_for_question_generation(context)
         response = await self.llm_client.generate_answer(question_generation_prompt)
         question = self.parse_response(response)["question"]
         if not question:
-            return result
+            return []
         logger.debug("Question: %s", question)
         logger.debug("Answer: %s", context)
         qa_pairs = {
-            compute_content_hash(question): {
-                "question": question,
-                "answer": context,
-            }
+            "question": question,
+            "answer": context,
         }
-        result.update(qa_pairs)
-        return result
+        return [qa_pairs]

graphgen/models/generator/atomic_generator.py

@@ -3,7 +3,7 @@
 
 from graphgen.bases import BaseGenerator
 from graphgen.templates import ATOMIC_GENERATION_PROMPT
-from graphgen.utils import compute_content_hash, detect_main_language, logger
+from graphgen.utils import detect_main_language, logger
 
 
 class AtomicGenerator(BaseGenerator):
@@ -23,7 +23,7 @@ def build_prompt(
         return prompt
 
     @staticmethod
-    def parse_response(response: str) -> dict:
+    def parse_response(response: str) -> list[dict]:
         """
         AtomicGenerator normally generates one QA pair per response.
         So we just need to parse one QA pair from the response.
@@ -38,15 +38,10 @@ def parse_response(response: str) -> dict:
             answer = answer_match.group(1).strip()
         else:
             logger.warning("Failed to parse response: %s", response)
-            return {}
+            return []
 
         question = question.strip('"').strip("'")
         answer = answer.strip('"').strip("'")
         logger.debug("Question: %s", question)
         logger.debug("Answer: %s", answer)
-        return {
-            compute_content_hash(question): {
-                "question": question,
-                "answer": answer,
-            }
-        }
+        return [{"question": question, "answer": answer}]

graphgen/models/generator/cot_generator.py

@@ -100,28 +100,25 @@ async def generate(
         batch: tuple[
             list[tuple[str, dict]], list[tuple[Any, Any, dict] | tuple[Any, Any, Any]]
         ],
-    ) -> dict[str, Any]:
+    ) -> list[dict]:
         """
         Generate QAs based on a given batch.
         :param batch
         :return: QA pairs
         """
-        result = {}
         prompt = self.build_prompt(batch)
         response = await self.llm_client.generate_answer(prompt)
         response = self.parse_response(response)
         if not response:
-            return result
+            return []
         question, reasoning_path = response["question"], response["reasoning_path"]
         prompt = self.build_prompt_for_cot_generation(batch, question, reasoning_path)
         cot_answer = await self.llm_client.generate_answer(prompt)
         logger.debug("CoT Answer: %s", cot_answer)
-        qa_pairs = {
-            compute_content_hash(question): {
+        return [
+            {
                 "question": question,
                 "answer": cot_answer,
                 "reasoning_path": reasoning_path,
             }
-        }
-        result.update(qa_pairs)
-        return result
+        ]

graphgen/models/generator/fill_in_blank_generator.py

@@ -12,7 +12,7 @@ def __init__(self, llm_client, num_of_questions) -> None:
         self.num_of_questions = num_of_questions
 
     @staticmethod
-    def parse_response(response: str) -> Any:
+    def parse_response(response: str) -> list[dict]:
         """
         Parse fill-in-the-blank QA pairs from the LLM response.
         Each QA pair contains question text with placeholders and the correct answer(s).
@@ -21,14 +21,14 @@ def parse_response(response: str) -> Any:
         :return: Dictionary mapping question hash to question data, where each
                  value is a dict with "question", "answer", and "answers" keys
         """
-        qa_pairs = {}
+        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 {}
+            return qa_pairs
 
         for block in qa_blocks:
             # Extract and clean question text
@@ -55,13 +55,13 @@ def parse_response(response: str) -> Any:
                 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,
-                "answer": answer_text,  # Original answer text with commas
-                "answers": answers,  # List of individual answers: ["A8X"] or ["A8X", "八百万"]
-            }
+            qa_pairs.append(
+                {
+                    "question": question,
+                    "answer": answer_text,  # Original answer text with commas
+                    "answers": answers,  # List of individual answers: ["A8X"] or ["A8X", "八百万"]
+                }
+            )
 
             logger.debug(
                 "Successfully parsed fill-in-the-blank question: %s", question[:50]

graphgen/models/generator/multi_answer_generator.py

@@ -12,7 +12,7 @@ def __init__(self, llm_client, num_of_questions) -> None:
         self.num_of_questions = num_of_questions
 
     @staticmethod
-    def parse_response(response: str) -> Any:
+    def parse_response(response: str) -> list[dict]:
         """
         Parse multiple-answer QA pairs from the LLM response.
         Each QA pair contains question text, four options, and the correct answers (one or more).
@@ -21,14 +21,14 @@ def parse_response(response: str) -> Any:
         :return: Dictionary mapping question hash to question data, where each
                  value is a dict with "question", "options", and "answer" keys
         """
-        qa_pairs = {}
+        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 {}
+            return qa_pairs
 
         for block in qa_blocks:
             # Extract and clean question text
@@ -61,7 +61,9 @@ def parse_response(response: str) -> Any:
                 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(",") if ans.strip()]
+            answers = [
+                ans.strip().upper() for ans in answer_text.split(",") if ans.strip()
+            ]
             invalid_answers = [ans for ans in answers if ans not in options]
             if invalid_answers:
                 logger.warning(
@@ -76,13 +78,13 @@ def parse_response(response: str) -> Any:
                 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),
-            }
+            qa_pairs.append(
+                {
+                    "question": question,
+                    "options": options,  # Dict like {"A": "text", "B": "text", ...}
+                    "answers": answers,  # List of correct answers: ["A", "C"]
+                }
+            )
 
             logger.debug("Successfully parsed MAQ: %s", question[:50])
 

graphgen/models/generator/multi_choice_generator.py

@@ -12,7 +12,7 @@ def __init__(self, llm_client, num_of_questions) -> None:
         self.num_of_questions = num_of_questions
 
     @staticmethod
-    def parse_response(response: str) -> Any:
+    def parse_response(response: str) -> list[dict]:
         """
         Parse multiple choice QA pairs from the LLM response.
         Each QA pair contains question text, four options, and the correct answer.
@@ -21,14 +21,14 @@ def parse_response(response: str) -> Any:
         :return: Dictionary mapping question hash to question data, where each
                  value is a dict with "question", "options", and "answer" keys
         """
-        qa_pairs = {}
+        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 {}
+            return qa_pairs
 
         for block in qa_blocks:
             # Extract and clean question text
@@ -76,13 +76,13 @@ def parse_response(response: str) -> Any:
                 )
                 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": answer,  # Single letter: "A", "B", "C", or "D"
-            }
+            qa_pairs.append(
+                {
+                    "question": question,
+                    "options": options,  # Dict like {"A": "text", "B": "text", ...}
+                    "answer": answer,  # Single letter: "A", "B", "C", or "D"
+                }
+            )
 
             logger.debug("Successfully parsed MCQ: %s", question[:50])