"""Synthetic QA dataset generator. Given a corpus of documents, this generates question-answer pairs for evaluation. Why you need this: - Real eval datasets are expensive to collect manually - Companies need to eval on THEIR data, not public benchmarks - Synthetic data lets you bootstrap an eval suite in minutes How it works (3-step pipeline): Step 1: CHUNK — split documents into digestible passages Step 2: GENERATE — ask LLM to generate Q&A pairs from each chunk Step 3: VALIDATE — ask a second LLM call to verify the Q&A is good Teaching note: This "generate then filter" pattern is used everywhere in ML pipelines. Generate a lot cheaply, then filter expensively. You'll see it in data augmentation, synthetic training data, and here in eval data generation. """ import asyncio import logging import uuid import json import re from typing import Optional from ..types import QAPair, DatasetExample, EvaluationMetric from ..utils.llm_client import LLMClient logger = logging.getLogger(__name__) # ─── Prompts ───────────────────────────────────────────────────────────────── # Prompt engineering for dataset generation is an art. # Key insight: ask for diverse question types to avoid a dataset that only # has easy factoid questions. Real systems fail on reasoning and multi-hop queries. QA_GENERATION_SYSTEM_PROMPT = """You are an expert at creating evaluation datasets for AI systems. Given a passage of text, generate high-quality question-answer pairs for testing a RAG (Retrieval Augmented Generation) system. Generate diverse question types: - FACTOID: Direct fact lookup ("What is X?", "When did Y happen?") - REASONING: Requires inference ("Why did X cause Y?", "What would happen if...?") - MULTI-HOP: Requires combining multiple facts from the passage - COMPARISON: Comparing entities or concepts ("How does X differ from Y?") Rules: 1. Questions must be answerable ONLY from the passage (not general knowledge) 2. Answers must be grounded in specific text from the passage 3. Include the relevant quote from the passage as supporting evidence 4. Make questions challenging enough to test the system meaningfully 5. Vary difficulty: some easy, some medium, some hard Respond with a JSON array of QA pairs.""" QA_GENERATION_USER_TEMPLATE = """Generate {n_questions} diverse question-answer pairs from this passage. PASSAGE: {passage} Return a JSON array with this structure: [ {{ "question": "...", "answer": "...", "question_type": "factoid|reasoning|multi_hop|comparison", "difficulty": "easy|medium|hard", "supporting_quote": "exact quote from passage", "tags": ["topic1", "topic2"] }} ] Generate exactly {n_questions} pairs. Make them varied and challenging.""" QA_VALIDATION_SYSTEM_PROMPT = """You are a quality reviewer for AI evaluation datasets. Review each question-answer pair and check: 1. Is the question clear and unambiguous? 2. Is the answer correct and grounded in the passage? 3. Can the question be answered from the passage alone? 4. Is the answer complete (not missing key information)? 5. Is it a meaningful test (not trivially easy)? Be strict — only approve pairs that would genuinely test a RAG system.""" QA_VALIDATION_USER_TEMPLATE = """Review this question-answer pair for quality. PASSAGE: {passage} QUESTION: {question} ANSWER: {answer} Is this a good evaluation pair? Respond with JSON: {{ "is_valid": true/false, "issues": ["list of any problems"], "quality_score": 0.0-1.0, "reasoning": "brief explanation" }}""" # ─── Core classes ───────────────────────────────────────────────────────────── class DatasetGenerator: """Generates synthetic evaluation datasets from document corpora. Usage: generator = DatasetGenerator(llm_client, questions_per_chunk=3) # From a list of documents dataset = await generator.generate_from_documents( documents=["doc text 1", "doc text 2", ...], n_questions_per_chunk=3, ) print(f"Generated {len(dataset)} QA pairs") for ex in dataset[:3]: print(f"Q: {ex.qa_pair.question}") print(f"A: {ex.qa_pair.answer}") """ def __init__( self, llm_client: LLMClient, chunk_size: int = 800, chunk_overlap: int = 100, validate: bool = True, min_quality_score: float = 0.7, concurrency: int = 3, ): """ Args: llm_client: LLM to use for generation and validation. chunk_size: Characters per chunk (800 ≈ 200 tokens, sweet spot for Q&A). chunk_overlap: Overlap between chunks to avoid cutting mid-sentence. validate: Whether to run a validation pass (recommended, costs more). min_quality_score: Threshold for keeping a QA pair after validation. concurrency: Max concurrent LLM calls (keep low, generation is expensive). """ self.llm_client = llm_client self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap self.validate = validate self.min_quality_score = min_quality_score self._semaphore = asyncio.Semaphore(concurrency) def chunk_document(self, text: str) -> list[str]: """Split a document into overlapping chunks. We use character-level chunking with overlap so context isn't lost at chunk boundaries. In production, you'd use semantic chunking (split on paragraph/section boundaries), but this is good enough. """ chunks = [] start = 0 while start < len(text): end = start + self.chunk_size chunk = text[start:end] # Try to end at a sentence boundary for cleaner chunks if end < len(text): last_period = chunk.rfind(". ") if last_period > self.chunk_size // 2: chunk = chunk[: last_period + 1] chunks.append(chunk.strip()) start += len(chunk) - self.chunk_overlap return [c for c in chunks if len(c) > 100] # Skip tiny chunks async def _generate_qa_for_chunk( self, chunk: str, n_questions: int = 3, ) -> list[dict]: """Generate raw QA pairs for a single chunk.""" async with self._semaphore: user_prompt = QA_GENERATION_USER_TEMPLATE.format( passage=chunk, n_questions=n_questions, ) try: response = await self.llm_client.generate( system_prompt=QA_GENERATION_SYSTEM_PROMPT, user_message=user_prompt, temperature=0.7, # Higher temp = more diverse questions max_tokens=2000, ) # Extract JSON array from response json_match = re.search(r'\[.*\]', response, re.DOTALL) if json_match: return json.loads(json_match.group()) return json.loads(response) except (json.JSONDecodeError, Exception) as e: logger.warning(f"QA generation failed for chunk: {e}") return [] async def _validate_qa_pair( self, chunk: str, question: str, answer: str, ) -> tuple[bool, float]: """Validate a QA pair — returns (is_valid, quality_score).""" async with self._semaphore: user_prompt = QA_VALIDATION_USER_TEMPLATE.format( passage=chunk, question=question, answer=answer, ) try: response = await self.llm_client.generate( system_prompt=QA_VALIDATION_SYSTEM_PROMPT, user_message=user_prompt, temperature=0.1, # Low temp for consistent validation max_tokens=300, ) json_match = re.search(r'\{.*\}', response, re.DOTALL) if json_match: data = json.loads(json_match.group()) is_valid = data.get("is_valid", False) quality = float(data.get("quality_score", 0.5)) return is_valid and quality >= self.min_quality_score, quality return False, 0.0 except Exception as e: logger.warning(f"Validation failed: {e}") return False, 0.0 async def generate_from_documents( self, documents: list[str], n_questions_per_chunk: int = 3, ) -> list[DatasetExample]: """Main entry point: generate a dataset from a list of document strings. Args: documents: Raw document texts (e.g., loaded from PDFs, web pages). n_questions_per_chunk: How many Q&A pairs to generate per chunk. Returns: List of DatasetExample objects ready for evaluation. """ # Step 1: Chunk all documents all_chunks: list[str] = [] for doc in documents: chunks = self.chunk_document(doc) all_chunks.extend(chunks) logger.info(f"Chunked document into {len(chunks)} chunks") logger.info( f"Total: {len(all_chunks)} chunks across {len(documents)} documents. " f"Generating {n_questions_per_chunk} Q&A per chunk = " f"~{len(all_chunks) * n_questions_per_chunk} pairs before filtering." ) # Step 2: Generate QA pairs for all chunks concurrently generation_tasks = [ self._generate_qa_for_chunk(chunk, n_questions_per_chunk) for chunk in all_chunks ] chunk_results = await asyncio.gather(*generation_tasks, return_exceptions=True) # Step 3: Validate and convert to DatasetExample dataset_examples: list[DatasetExample] = [] for chunk, raw_pairs in zip(all_chunks, chunk_results): if isinstance(raw_pairs, Exception): logger.warning(f"Chunk generation failed: {raw_pairs}") continue for raw_qa in raw_pairs: question = raw_qa.get("question", "") answer = raw_qa.get("answer", "") if not question or not answer: continue # Validate if requested if self.validate: is_valid, quality_score = await self._validate_qa_pair( chunk, question, answer ) if not is_valid: logger.debug( f"Filtered low-quality pair (score={quality_score:.2f}): {question[:50]}..." ) continue else: quality_score = 1.0 example = DatasetExample( id=str(uuid.uuid4()), qa_pair=QAPair( question=question, answer=answer, context=chunk, ), source_document=chunk[:200] + "..." if len(chunk) > 200 else chunk, difficulty=raw_qa.get("difficulty", "medium"), tags=raw_qa.get("tags", []) + [raw_qa.get("question_type", "factoid")], ) dataset_examples.append(example) logger.info( f"Generated {len(dataset_examples)} validated QA pairs " f"(filtered {len(all_chunks) * n_questions_per_chunk - len(dataset_examples)} low-quality pairs)" ) return dataset_examples async def generate_from_text( self, text: str, n_questions: int = 10, ) -> list[DatasetExample]: """Convenience method: generate from a single text string.""" return await self.generate_from_documents( documents=[text], n_questions_per_chunk=max(1, n_questions // max(1, len(self.chunk_document(text)))), )