Matrix-Corp/Zenith-32b-p300-V1
0
1"""Advanced Data Augmentation for Training"""
2
3import json
4import logging
5import random
6from abc import ABC, abstractmethod
7from dataclasses import dataclass
8from typing import Any, Dict, List, Optional, Tuple
9
10import numpy as np
11
12logger = logging.getLogger(__name__)
13
14
15@dataclass
16class AugmentationConfig:
17 """Configuration for data augmentation."""
18 enabled_methods: List[str] = field(default_factory=lambda: [
19 "synonym_replacement",
20 "back_translation",
21 "code_perturbation",
22 "paraphrasing",
23 "noise_injection",
24 ])
25 probabilities: Dict[str, float] = field(default_factory=lambda: {
26 "synonym_replacement": 0.3,
27 "back_translation": 0.2,
28 "code_perturbation": 0.4,
29 "paraphrasing": 0.3,
30 "noise_injection": 0.1,
31 })
32 max_augmentations_per_sample: int = 2
33
34
35class AugmentationMethod(ABC):
36 """Base class for augmentation methods."""
37
38 @abstractmethod
39 def augment(self, sample: Dict[str, Any]) -> Optional[Dict[str, Any]]:
40 """Apply augmentation to sample. Return augmented sample or None if failed."""
41 pass
42
43 @abstractmethod
44 def can_augment(self, sample: Dict[str, Any]) -> bool:
45 """Check if sample can be augmented by this method."""
46 pass
47
48
49class SynonymReplacement(AugmentationMethod):
50 """Replace words with synonyms."""
51
52 def __init__(self, replacement_prob: float = 0.1):
53 self.replacement_prob = replacement_prob
54 # Simple synonym dictionary (in practice, use WordNet or embeddings)
55 self.synonyms = {
56 "good": ["excellent", "great", "fine", "quality", "superb"],
57 "bad": ["poor", "terrible", "awful", "inferior", "subpar"],
58 "big": ["large", "huge", "enormous", "massive", "giant"],
59 "small": ["tiny", "little", "miniature", "compact", "petite"],
60 "fast": ["quick", "rapid", "speedy", "swift", "expedited"],
61 "slow": ["sluggish", "leisurely", "unhurried", "gradual", "delayed"],
62 "create": ["build", "generate", "produce", "develop", "construct"],
63 "use": ["utilize", "employ", "apply", "leverage", "harness"],
64 "find": ["discover", "locate", "detect", "identify", "uncover"],
65 "improve": ["enhance", "upgrade", "optimize", "refine", "better"],
66 }
67
68 def can_augment(self, sample: Dict[str, Any]) -> bool:
69 """Check if sample has text to augment."""
70 text = self._extract_text(sample)
71 return len(text.split()) > 10
72
73 def augment(self, sample: Dict[str, Any]) -> Optional[Dict[str, Any]]:
74 """Replace random words with synonyms."""
75 text = self._extract_text(sample)
76 words = text.split()
77
78 # Replace random words
79 new_words = []
80 for word in words:
81 if random.random() < self.replacement_prob and word.lower() in self.synonyms:
82 synonym = random.choice(self.synonyms[word.lower()])
83 # Preserve capitalization
84 if word[0].isupper():
85 synonym = synonym.capitalize()
86 new_words.append(synonym)
87 else:
88 new_words.append(word)
89
90 new_text = " ".join(new_words)
91 if new_text == text:
92 return None
93
94 augmented = sample.copy()
95 self._replace_text(augmented, new_text)
96 augmented["augmentation"] = "synonym_replacement"
97 return augmented
98
99 def _extract_text(self, sample: Dict[str, Any]) -> str:
100 """Extract text from sample."""
101 if "conversations" in sample:
102 conv = sample["conversations"]
103 if isinstance(conv, list):
104 return " ".join(msg.get("content", "") for msg in conv if isinstance(msg, dict))
105 return sample.get("text", sample.get("content", ""))
106
107 def _replace_text(self, sample: Dict[str, Any], new_text: str):
108 """Replace text in sample."""
109 if "conversations" in sample:
110 conv = sample["conversations"]
111 if isinstance(conv, list):
112 # Replace content of first message
113 for msg in conv:
114 if isinstance(msg, dict) and "content" in msg:
115 msg["content"] = new_text[:len(msg["content"])]
116 break
117 else:
118 sample["text"] = new_text
119 sample["content"] = new_text
120
121
122class CodePerturbation(AugmentationMethod):
123 """Perturb code while preserving functionality."""
124
125 def __init__(self):
126 self.perturbations = [
127 self._rename_variables,
128 self._reorder_statements,
129 self._add_redundant_parentheses,
130 self._change_loop_style,
131 self._add_comments,
132 ]
133
134 def can_augment(self, sample: Dict[str, Any]) -> bool:
135 """Check if sample has code."""
136 return "code" in sample or any(
137 "```" in str(conv.get("content", ""))
138 for conv in sample.get("conversations", [])
139 if isinstance(conv, dict)
140 )
141
142 def augment(self, sample: Dict[str, Any]) -> Optional[Dict[str, Any]]:
143 """Apply random code perturbation."""
144 code = self._extract_code(sample)
145 if not code:
146 return None
147
148 # Apply random perturbation
149 perturbation = random.choice(self.perturbations)
150 new_code = perturbation(code)
151
152 if new_code == code:
153 return None
154
155 augmented = sample.copy()
156 self._replace_code(augmented, new_code)
157 augmented["augmentation"] = f"code_perturbation:{perturbation.__name__}"
158 return augmented
159
160 def _extract_code(self, sample: Dict[str, Any]) -> str:
161 """Extract code from sample."""
162 if "code" in sample:
163 return sample["code"]
164 # Look for code blocks in conversations
165 for conv in sample.get("conversations", []):
166 if isinstance(conv, dict):
167 content = conv.get("content", "")
168 if "```" in content:
169 # Extract code block
170 parts = content.split("```")
171 if len(parts) >= 2:
172 return parts[1].strip()
173 return ""
174
175 def _replace_code(self, sample: Dict[str, Any], new_code: str):
176 """Replace code in sample."""
177 if "code" in sample:
178 sample["code"] = new_code
179 else:
180 for conv in sample.get("conversations", []):
181 if isinstance(conv, dict) and "```" in conv.get("content", ""):
182 parts = conv["content"].split("```")
183 conv["content"] = f"```{new_code}```"
184
185 def _rename_variables(self, code: str) -> str:
186 """Rename variables to random names (simple version)."""
187 # This is a simplified version - in practice use AST parsing
188 import re
189 # Find variable names (simplistic)
190 variables = re.findall(r'\b([a-zA-Z_][a-zA-Z0-9_]*)\b', code)
191 unique_vars = set(variables)
192
193 # Generate random replacements
194 replacements = {}
195 for var in unique_vars:
196 if len(var) > 1 and var not in ["if", "for", "while", "def", "class", "return", "import", "from"]:
197 new_name = f"var_{random.randint(1000, 9999)}"
198 replacements[var] = new_name
199
200 # Replace
201 for old, new in replacements.items():
202 code = code.replace(old, new)
203
204 return code
205
206 def _reorder_statements(self, code: str) -> str:
207 """Reorder independent statements."""
208 lines = code.split('\n')
209 # Simple: shuffle non-indented lines (top-level statements)
210 # This is risky - only apply to simple code
211 return code # TODO: Implement safely
212
213 def _add_redundant_parentheses(self, code: str) -> str:
214 """Add redundant parentheses."""
215 # Simplistic: add parentheses around binary operations
216 import re
217 # This is placeholder - would need proper parsing
218 return code
219
220 def _change_loop_style(self, code: str) -> str:
221 """Change between for loops and while loops where possible."""
222 # This requires AST parsing - placeholder
223 return code
224
225 def _add_comments(self, code: str) -> str:
226 """Add explanatory comments."""
227 lines = code.split('\n')
228 new_lines = []
229 for i, line in enumerate(lines):
230 new_lines.append(line)
231 if line.strip() and not line.strip().startswith('#'):
232 if random.random() < 0.2:
233 new_lines.append(f"# TODO: Explain this line")
234 return '\n'.join(new_lines)
235
236
237class BackTranslation(AugmentationMethod):
238 """Simulate back-translation by paraphrasing."""
239
240 def __init__(self):
241 self.paraphrase_templates = [
242 "In other words, {text}",
243 "To put it differently, {text}",
244 "That is to say, {text}",
245 "Alternatively, {text}",
246 ]
247
248 def can_augment(self, sample: Dict[str, Any]) -> bool:
249 """Check if sample has text suitable for back translation."""
250 text = self._extract_text(sample)
251 return len(text) > 50
252
253 def augment(self, sample: Dict[str, Any]) -> Optional[Dict[str, Any]]:
254 """Apply back translation simulation."""
255 text = self._extract_text(sample)
256 template = random.choice(self.paraphrase_templates)
257 new_text = template.format(text=text[:200]) + text[200:] # Add prefix
258
259 if new_text == text:
260 return None
261
262 augmented = sample.copy()
263 self._replace_text(augmented, new_text)
264 augmented["augmentation"] = "back_translation"
265 return augmented
266
267 def _extract_text(self, sample: Dict[str, Any]) -> str:
268 """Extract text from sample."""
269 if "conversations" in sample:
270 conv = sample["conversations"]
271 if isinstance(conv, list):
272 return " ".join(msg.get("content", "") for msg in conv if isinstance(msg, dict))
273 return sample.get("text", sample.get("content", ""))
274
275 def _replace_text(self, sample: Dict[str, Any], new_text: str):
276 """Replace text in sample."""
277 if "conversations" in sample:
278 conv = sample["conversations"]
279 if isinstance(conv, list):
280 for msg in conv:
281 if isinstance(msg, dict) and "content" in msg:
282 msg["content"] = new_text[:len(msg["content"])]
283 break
284 else:
285 sample["text"] = new_text
286 sample["content"] = new_text
287
288
289class Paraphrasing(AugmentationMethod):
290 """Paraphrase text using templates."""
291
292 def __init__(self):
293 self.paraphrase_patterns = [
294 (r"\b(is)\b", ["represents", "constitutes", "means"]),
295 (r"\b(has)\b", ["contains", "possesses", "includes"]),
296 (r"\b(use)\b", ["utilize", "employ", "leverage"]),
297 (r"\b(make)\b", ["create", "build", "produce"]),
298 (r"\b(find)\b", ["discover", "locate", "identify"]),
299 ]
300
301 def can_augment(self, sample: Dict[str, Any]) -> bool:
302 """Check if sample has text."""
303 text = self._extract_text(sample)
304 return len(text) > 30
305
306 def augment(self, sample: Dict[str, Any]) -> Optional[Dict[str, Any]]:
307 """Apply paraphrasing."""
308 text = self._extract_text(sample)
309 new_text = text
310
311 # Apply random pattern
312 pattern, replacements = random.choice(self.paraphrase_patterns)
313 import re
314 matches = re.findall(pattern, text, re.IGNORECASE)
315 if matches:
316 # Replace first occurrence
317 old_word = matches[0]
318 new_word = random.choice(replacements)
319 new_text = re.sub(pattern, new_word, text, count=1, flags=re.IGNORECASE)
320
321 if new_text == text:
322 return None
323
324 augmented = sample.copy()
325 self._replace_text(augmented, new_text)
326 augmented["augmentation"] = "paraphrasing"
327 return augmented
328
329 def _extract_text(self, sample: Dict[str, Any]) -> str:
330 """Extract text from sample."""
331 if "conversations" in sample:
332 conv = sample["conversations"]
333 if isinstance(conv, list):
334 return " ".join(msg.get("content", "") for msg in conv if isinstance(msg, dict))
335 return sample.get("text", sample.get("content", ""))
336
337 def _replace_text(self, sample: Dict[str, Any], new_text: str):
338 """Replace text in sample."""
339 if "conversations" in sample:
340 conv = sample["conversations"]
341 if isinstance(conv, list):
342 for msg in conv:
343 if isinstance(msg, dict) and "content" in msg:
344 msg["content"] = new_text[:len(msg["content"])]
345 break
346 else:
347 sample["text"] = new_text
348 sample["content"] = new_text
349
350
351class NoiseInjection(AugmentationMethod):
352 """Inject noise into text."""
353
354 def __init__(self, noise_prob: float = 0.01):
355 self.noise_prob = noise_prob
356 self.noise_tokens = ["[MASK]", "<noise>", "...", "[UNK]"]
357
358 def can_augment(self, sample: Dict[str, Any]) -> bool:
359 """Check if sample has text."""
360 text = self._extract_text(sample)
361 return len(text) > 20
362
363 def augment(self, sample: Dict[str, Any]) -> Optional[Dict[str, Any]]:
364 """Inject noise tokens."""
365 text = self._extract_text(sample)
366 words = text.split()
367
368 # Randomly replace words with noise
369 new_words = []
370 for word in words:
371 if random.random() < self.noise_prob and len(word) > 3:
372 new_words.append(random.choice(self.noise_tokens))
373 else:
374 new_words.append(word)
375
376 new_text = " ".join(new_words)
377 if new_text == text:
378 return None
379
380 augmented = sample.copy()
381 self._replace_text(augmented, new_text)
382 augmented["augmentation"] = "noise_injection"
383 return augmented
384
385 def _extract_text(self, sample: Dict[str, Any]) -> str:
386 """Extract text from sample."""
387 if "conversations" in sample:
388 conv = sample["conversations"]
389 if isinstance(conv, list):
390 return " ".join(msg.get("content", "") for msg in conv if isinstance(msg, dict))
391 return sample.get("text", sample.get("content", ""))
392
393 def _replace_text(self, sample: Dict[str, Any], new_text: str):
394 """Replace text in sample."""
395 if "conversations" in sample:
396 conv = sample["conversations"]
397 if isinstance(conv, list):
398 for msg in conv:
399 if isinstance(msg, dict) and "content" in msg:
400 msg["content"] = new_text[:len(msg["content"])]
401 break
402 else:
403 sample["text"] = new_text
404 sample["content"] = new_text
405
406
407class DataAugmenter:
408 """Manages multiple augmentation methods."""
409
410 def __init__(self, config: AugmentationConfig):
411 self.config = config
412 self.methods: Dict[str, AugmentationMethod] = {
413 "synonym_replacement": SynonymReplacement(),
414 "back_translation": BackTranslation(),
415 "code_perturbation": CodePerturbation(),
416 "paraphrasing": Paraphrasing(),
417 "noise_injection": NoiseInjection(),
418 }
419
420 def augment(self, sample: Dict[str, Any]) -> Optional[Dict[str, Any]]:
421 """Apply random augmentation to sample."""
422 # Choose random enabled method
423 enabled_methods = [
424 m for m in self.config.enabled_methods
425 if m in self.methods and self.methods[m].can_augment(sample)
426 ]
427
428 if not enabled_methods:
429 return None
430
431 method_name = random.choice(enabled_methods)
432 method = self.methods[method_name]
433
434 # Apply augmentation
435 augmented = method.augment(sample)
436
437 if augmented:
438 augmented["augmentation_applied"] = method_name
439
440 return augmented
441
442 def augment_batch(
443 self,
444 batch: List[Dict[str, Any]],
445 augmentation_ratio: float = 0.1,
446 ) -> List[Dict[str, Any]]:
447 """Augment a batch of samples."""
448 augmented_batch = []
449
450 for sample in batch:
451 augmented_batch.append(sample)
452
453 if random.random() < augmentation_ratio:
454 augmented = self.augment(sample)
455 if augmented:
456 augmented_batch.append(augmented)
457
458 return augmented_batch
459
460
461def augment_sample(
462 sample: Dict[str, Any],
463 methods: List[str],
464 max_augmentations: int = 2,
465) -> List[Dict[str, Any]]:
466 """Augment a single sample with multiple methods."""
467 augmenter = DataAugmenter(AugmentationConfig(enabled_methods=methods))
468 results = [sample]
469
470 for _ in range(max_augmentations):
471 augmented = augmenter.augment(sample)
472 if augmented:
473 results.append(augmented)
474
475 return results
476