File size: 5,381 Bytes
c8b77b5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | """
Synthetic Data Generator for MangoMAS Local
This module provides a framework for generating synthetic training data
for different specialized capabilities, adaptable across all modules.
"""
import json
import logging
import os
import random
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple
import yaml
from tqdm import tqdm
logger = logging.getLogger(__name__)
class SyntheticDataGenerator(ABC):
"""
Abstract base class for synthetic data generators.
Each specialized module can implement this interface for data generation.
"""
def __init__(self, config: Dict[str, Any], output_dir: str = "data/processed"):
"""
Initialize the synthetic data generator.
Args:
config: Configuration for the data generator
output_dir: Directory to save generated data
"""
self.config = config
self.output_dir = Path(output_dir)
self.output_dir.mkdir(parents=True, exist_ok=True)
# Set number of examples to generate
self.num_examples = config.get("synthetic_examples", 1000)
# Template pool for generation
self.templates = self._load_templates()
logger.info(f"Initialized {self.__class__.__name__} with {self.num_examples} examples")
@abstractmethod
def _load_templates(self) -> List[Dict[str, Any]]:
"""
Load templates for data generation.
Each implementation should define its own templates.
Returns:
List of template dictionaries
"""
pass
@abstractmethod
def generate_example(self) -> Dict[str, Any]:
"""
Generate a single synthetic training example.
Returns:
Dictionary with the generated example
"""
pass
def generate_dataset(self, filename: str, num_examples: Optional[int] = None) -> str:
"""
Generate a synthetic dataset and save to a JSONL file.
Args:
filename: Name of the output file
num_examples: Number of examples to generate (overrides config)
Returns:
Path to the generated dataset file
"""
n = num_examples if num_examples is not None else self.num_examples
output_file = self.output_dir / filename
logger.info(f"Generating {n} synthetic examples for {self.__class__.__name__}")
with open(output_file, 'w', encoding='utf-8') as f:
for _ in tqdm(range(n), desc=f"Generating {filename}"):
example = self.generate_example()
f.write(json.dumps(example) + '\n')
logger.info(f"Generated dataset saved to {output_file}")
return str(output_file)
def augment_existing_dataset(self, input_file: str, output_file: Optional[str] = None,
ratio: float = 0.5) -> str:
"""
Augment an existing dataset with synthetic examples.
Args:
input_file: Path to the existing dataset
output_file: Path to save the augmented dataset (or None to overwrite)
ratio: Ratio of synthetic to original examples
Returns:
Path to the augmented dataset
"""
if output_file is None:
output_file = input_file
# Load existing data
existing_data = []
try:
with open(input_file, 'r', encoding='utf-8') as f:
for line in f:
existing_data.append(json.loads(line.strip()))
except (FileNotFoundError, json.JSONDecodeError) as e:
logger.warning(f"Error loading existing data: {e}")
existing_data = []
# Calculate number of synthetic examples to generate
n_existing = len(existing_data)
n_synthetic = int(n_existing * ratio)
# Generate synthetic examples
synthetic_data = [self.generate_example() for _ in tqdm(range(n_synthetic),
desc=f"Generating augmentation data")]
# Combine datasets
combined_data = existing_data + synthetic_data
random.shuffle(combined_data)
# Save augmented dataset
with open(output_file, 'w', encoding='utf-8') as f:
for item in combined_data:
f.write(json.dumps(item) + '\n')
logger.info(f"Augmented dataset with {n_synthetic} synthetic examples, saved to {output_file}")
return output_file
class SyntheticDataGeneratorRegistry:
"""Registry for all synthetic data generators in the system."""
_generators = {}
@classmethod
def register(cls, module_type: str, generator_class):
"""Register a generator class for a module type."""
cls._generators[module_type] = generator_class
@classmethod
def get_generator(cls, module_type: str, config: Dict[str, Any], output_dir: str) -> SyntheticDataGenerator:
"""Get a generator instance for a module type."""
if module_type not in cls._generators:
raise ValueError(f"No generator registered for module type: {module_type}")
return cls._generators[module_type](config, output_dir)
@classmethod
def list_generators(cls) -> List[str]:
"""List all registered generator types."""
return list(cls._generators.keys())
"""
|