#!/usr/bin/env python3 """ Dataset Generator v2 - Refactored with SOLID principles Q&A -> TTS -> Features -> Save """ import os import sys import re import time import gc import logging import multiprocessing as mp from abc import ABC, abstractmethod from dataclasses import dataclass, field from pathlib import Path from typing import Dict, List, Optional, Any, Callable from concurrent.futures import ThreadPoolExecutor, as_completed from enum import Enum import numpy as np import requests import torch # ============================================================================= # CONFIGURATION (Single source of truth) # ============================================================================= class TimeoutConfig: """Timeout constants in seconds.""" HEARTBEAT_INTERVAL = 30 STUCK_WORKER_THRESHOLD = 120 STALL_WARNING = 30 STALL_CHECKPOINT = 60 STALL_EXIT = 180 NO_PROGRESS_EXIT = 600 QUEUE_GET = 5 QUEUE_PUT = 30 DRAIN_LOOP = 2.0 WORKER_JOIN = 5 class MemoryConfig: """Memory management constants.""" CLEANUP_INTERVAL_ITEMS = 500 CLEANUP_INTERVAL_BATCHES = 100 MAX_PUT_ATTEMPTS = 10 @dataclass class BatchSizeConfig: """Batch size configuration.""" tts: int = 100 whisper_workers: int = 4 snac: int = 20 @classmethod def from_vram(cls, vram_gb: float, shared_gpu: bool = True) -> 'BatchSizeConfig': """Calculate optimal batch sizes based on VRAM.""" if vram_gb >= 80: config = cls(tts=200, whisper_workers=8, snac=30) elif vram_gb >= 40: config = cls(tts=150, whisper_workers=6, snac=25) elif vram_gb >= 24: config = cls(tts=100, whisper_workers=4, snac=16) elif vram_gb >= 16: config = cls(tts=66, whisper_workers=2, snac=12) else: config = cls(tts=40, whisper_workers=1, snac=8) if not shared_gpu: config.snac = min(50, int(config.snac * 1.5)) config.whisper_workers = min(8, config.whisper_workers + 2) return config @dataclass class PipelineConfig: """Main pipeline configuration.""" output_path: str = "./data/dataset.pt" target_count: int = 1000 num_gpus: int = 1 log_file: Optional[str] = None batch_sizes: BatchSizeConfig = field(default_factory=BatchSizeConfig) # API Configuration groq_api_key: str = "" groq_model: str = "openai/gpt-oss-20b" groq_parallel_requests: int = 20 qa_per_request: int = 100 # ============================================================================= # LOGGING (Extracted responsibility) # ============================================================================= class PipelineLogger: """Centralized logging with file and console output.""" def __init__(self, log_file: Optional[str] = None): self._logger = logging.getLogger("dataset_generator") self._logger.setLevel(logging.INFO) self._logger.handlers.clear() # Console handler console = logging.StreamHandler(sys.stdout) console.setFormatter(logging.Formatter('%(message)s')) self._logger.addHandler(console) # File handler if log_file: file_handler = logging.FileHandler(log_file, mode='a') file_handler.setFormatter( logging.Formatter('%(asctime)s | %(message)s', datefmt='%H:%M:%S') ) self._logger.addHandler(file_handler) def log(self, msg: str) -> None: self._logger.info(msg) for handler in self._logger.handlers: handler.flush() def error(self, msg: str) -> None: self._logger.error(msg) def warning(self, msg: str) -> None: self._logger.warning(msg) # Global logger instance (initialized in main) _logger: Optional[PipelineLogger] = None def log(msg: str) -> None: """Global log function.""" if _logger: _logger.log(msg) else: print(msg) sys.stdout.flush() # ============================================================================= # INCREMENTAL SAVER (Append-only - instant saves) # ============================================================================= class IncrementalSaver: """ Append-only saves using batch files for instant I/O. Each batch is saved to a separate small file - no rewriting. Files: - output.pt: Base file (existing items) - output.pt.batches/: Directory with batch files - batch_0000.pt, batch_0001.pt, ... """ def __init__(self, config: PipelineConfig): self.config = config self._base_count = 0 self._batch_count = 0 # Number of batch files self._items_in_batches = 0 # Total items across all batch files self._pending_items: List[Any] = [] # Items not yet saved @property def _batches_dir(self) -> Path: return Path(f"{self.config.output_path}.batches") def load_existing(self) -> int: """ Count existing items from base file + batch files. Returns total count. """ # Count base file items if Path(self.config.output_path).exists(): try: data = torch.load( self.config.output_path, map_location="cpu", weights_only=False, mmap=True ) self._base_count = len(data) del data log(f"[Resume] Base file: {self._base_count} items") except Exception as e: log(f"[Resume] Failed to read base: {e}") # Count items in batch files # Use metadata file for true count, or estimate from max batch number metadata_path = self._batches_dir / "metadata.txt" if self._batches_dir.exists(): batch_files = sorted(self._batches_dir.glob("batch_*.pt")) num_local_files = len(batch_files) if num_local_files > 0: # Use max batch number + 1 for naming (prevents overwrites if files deleted) max_batch_num = max( int(f.stem.replace("batch_", "")) for f in batch_files ) self._batch_count = max_batch_num + 1 # Check metadata for true total count (survives deletions) if metadata_path.exists(): try: meta = metadata_path.read_text().strip().split("\n") for line in meta: if line.startswith("total_items="): self._items_in_batches = int(line.split("=")[1]) log(f"[Resume] Metadata: {self._items_in_batches} items, next batch: {self._batch_count}") except Exception: pass # If no metadata, estimate from local files if self._items_in_batches == 0: self._items_in_batches = num_local_files * self.config.batch_sizes.tts log(f"[Resume] Estimated from {num_local_files} local files: {self._items_in_batches} items, next batch: {self._batch_count}") else: log(f"[Resume] No batch files found") else: self._batches_dir.mkdir(parents=True, exist_ok=True) total = self._base_count + self._items_in_batches if total > 0: log(f"[Resume] Total: {total} items") return total def _save_metadata(self): """Write metadata file tracking true total count.""" metadata_path = self._batches_dir / "metadata.txt" try: metadata_path.write_text( f"total_items={self._items_in_batches}\n" f"next_batch={self._batch_count}\n" ) except Exception: pass def add_batch(self, items: List[Any]) -> int: """ Save batch instantly to a new file (append-only). No rewriting - just create a new small file. """ # Save batch to new file batch_path = self._batches_dir / f"batch_{self._batch_count:06d}.pt" try: torch.save(items, batch_path) self._batch_count += 1 self._items_in_batches += len(items) total = self.get_count() log(f"[Save] {total} items (+{len(items)})") # Update metadata every 10 batches if self._batch_count % 10 == 0: self._save_metadata() except Exception as e: log(f"[Save] ERROR: {e}") # Keep items for retry self._pending_items.extend(items) return self.get_count() def get_count(self) -> int: """Get total item count.""" return self._base_count + self._items_in_batches + len(self._pending_items) def finalize(self) -> int: """ Finalize dataset generation. Keeps batch files as-is (no merge) to avoid OOM on large datasets. Returns final count. """ # Save any pending items first if self._pending_items: batch_path = self._batches_dir / f"batch_{self._batch_count:06d}.pt" torch.save(self._pending_items, batch_path) self._batch_count += 1 self._items_in_batches += len(self._pending_items) self._pending_items = [] total = self.get_count() if self._items_in_batches == 0: log(f"[Final] No new items, keeping {self._base_count} base items") return self._base_count # Save metadata and log self._save_metadata() batch_files = sorted(self._batches_dir.glob("batch_*.pt")) log(f"[Final] Dataset complete: {total} items in {len(batch_files)} batch files") log(f"[Final] Batch dir: {self._batches_dir}") log(f"[Final] Skipping merge (too large for RAM). Use batch files directly for training.") return total # ============================================================================= # GPU UTILITIES # ============================================================================= class GPUManager: """GPU detection and memory management utilities.""" @staticmethod def get_num_gpus() -> int: if not torch.cuda.is_available(): return 0 return torch.cuda.device_count() @staticmethod def get_vram_gb(device_id: int = 0) -> float: if not torch.cuda.is_available(): return 0 return torch.cuda.get_device_properties(device_id).total_memory / 1024**3 @staticmethod def get_device_name(device_id: int = 0) -> str: if not torch.cuda.is_available(): return "CPU" return torch.cuda.get_device_properties(device_id).name @staticmethod def clear_memory() -> None: """Aggressively clear GPU memory.""" torch.cuda.empty_cache() torch.cuda.synchronize() gc.collect() @staticmethod def supports_lmdeploy(device_name: str) -> bool: """Check if GPU supports lmdeploy backend.""" unsupported = ["5090", "5080", "B100", "B200"] return not any(x in device_name for x in unsupported) # ============================================================================= # MESSAGE TYPES (Type Safety) # ============================================================================= class MessageType(Enum): """Status message types for inter-process communication.""" TTS_READY = "tts_ready" TTS_PROGRESS = "tts" TTS_DONE = "tts_done" TTS_HEARTBEAT = "tts_heartbeat" TTS_ERROR = "tts_error" FEAT_READY = "feat_ready" FEAT_PROGRESS = "feat" FEAT_DONE = "feat_done" FEAT_HEARTBEAT = "feat_heartbeat" FEAT_ERROR = "feat_error" FEAT_WARN = "feat_warn" QA_PROGRESS = "qa" QA_DONE = "qa_done" # ============================================================================= # WORKER BASE CLASS (DRY - common functionality) # ============================================================================= class BaseWorker(ABC): """Base class for pipeline workers with common functionality.""" def __init__( self, worker_id: int, status_queue: mp.Queue, worker_type: str ): self.worker_id = worker_id self.status_queue = status_queue self.worker_type = worker_type self.processed = 0 self.start_time = time.time() self.last_heartbeat = time.time() def send_heartbeat(self) -> None: """Send heartbeat if interval has passed.""" if time.time() - self.last_heartbeat > TimeoutConfig.HEARTBEAT_INTERVAL: msg_type = f"{self.worker_type}_heartbeat" self.status_queue.put((msg_type, self.worker_id, self.processed)) self.last_heartbeat = time.time() def send_ready(self) -> None: """Signal that worker is ready.""" msg_type = f"{self.worker_type}_ready" self.status_queue.put((msg_type, self.worker_id)) def send_progress(self, batch_rate: float) -> None: """Send progress update.""" elapsed = time.time() - self.start_time avg_rate = self.processed / elapsed if elapsed > 0 else 0 self.status_queue.put(( self.worker_type, self.worker_id, self.processed, avg_rate, batch_rate )) def send_done(self) -> None: """Signal worker completion.""" msg_type = f"{self.worker_type}_done" self.status_queue.put((msg_type, self.worker_id, self.processed)) def send_error(self, error: str) -> None: """Send error message.""" msg_type = f"{self.worker_type}_error" self.status_queue.put((msg_type, self.worker_id, error)) def send_warning(self, warning: str) -> None: """Send warning message.""" msg_type = f"{self.worker_type}_warn" self.status_queue.put((msg_type, self.worker_id, warning)) @abstractmethod def run(self) -> None: """Main worker loop - to be implemented by subclasses.""" pass # ============================================================================= # Q&A GENERATOR # ============================================================================= class QAGenerator: """Generates Q&A pairs from GROQ API.""" def __init__(self, config: PipelineConfig): self.config = config def generate_batch(self, request_id: int) -> List[Dict[str, str]]: """Generate a batch of Q&A pairs.""" headers = { "Authorization": f"Bearer {self.config.groq_api_key}", "Content-Type": "application/json" } prompt = f"""Generate {self.config.qa_per_request} unique Q&A pairs on diverse topics. Format: Q: [question] A: [answer] Questions 2-25 words, answers 1-3 sentences.""" for retry in range(3): try: response = requests.post( "https://api.groq.com/openai/v1/chat/completions", headers=headers, json={ "model": self.config.groq_model, "messages": [{"role": "user", "content": prompt}], "max_tokens": 8000, "temperature": 1.0 }, timeout=60 ) if response.status_code == 429: time.sleep(2 ** retry) continue response.raise_for_status() return self._parse_qa(response.json()["choices"][0]["message"]["content"]) except Exception: if retry < 2: time.sleep(1) return [] def _parse_qa(self, content: str) -> List[Dict[str, str]]: """Parse Q&A pairs from API response.""" pairs = [] content = content.replace("**", "") current_q, current_a = None, None for line in content.split("\n"): line = line.strip() if not line: continue qm = re.match(r"^[\d\.\)\-\*]*\s*[Qq][:\.]?\s*(.+)", line) am = re.match(r"^[\d\.\)\-\*]*\s*[Aa][:\.]?\s*(.+)", line) if qm: current_q = qm.group(1).strip() elif am: current_a = am.group(1).strip() if current_q and current_a: word_count = len(current_q.split()) if 2 <= word_count <= 25 and len(current_a) > 3: pairs.append({"q": current_q, "a": current_a}) current_q, current_a = None, None return pairs # ============================================================================= # SNAC TOKEN EXTRACTOR # ============================================================================= class SNACTokenExtractor: """Extracts SNAC tokens with correct position-based offsets.""" SNAC_BASE = 128266 TOKENS_PER_FRAME = 7 @classmethod def extract(cls, codes: List[torch.Tensor], idx: int) -> List[int]: """ Extract SNAC tokens for a single item. Each frame has 7 tokens: - Position 0: codebook 0 - Positions 1-2: codebook 1 - Positions 3-6: codebook 2 """ tokens = [] for j in range(codes[0].shape[-1]): # Position 0: codebook 0 token tokens.append(codes[0][idx, j].item() + cls.SNAC_BASE + 0 * 4096) # Positions 1-2: codebook 1 tokens if j * 2 + 1 < codes[1].shape[-1]: tokens.append(codes[1][idx, j * 2].item() + cls.SNAC_BASE + 1 * 4096) tokens.append(codes[1][idx, j * 2 + 1].item() + cls.SNAC_BASE + 2 * 4096) # Positions 3-6: codebook 2 tokens for k in range(4): if j * 4 + k < codes[2].shape[-1]: tokens.append(codes[2][idx, j * 4 + k].item() + cls.SNAC_BASE + (3 + k) * 4096) return tokens # ============================================================================= # WORD ALIGNMENT # ============================================================================= class WordAligner: """Generates word alignments for IST-LM interleaving.""" SNAC_FPS = 75 # SNAC frames per second SNAC_SAMPLES_PER_FRAME = 320 def __init__(self, tokenizer=None): self.tokenizer = tokenizer def align_proportional( self, audio_data: np.ndarray, text: str, sample_rate: int = 32000 ) -> List[Dict]: """Proportional word alignment based on character count.""" words = text.split() if not words: return [] # Calculate SNAC frame count audio_24k_samples = len(audio_data) * 24000 / sample_rate total_frames = int(audio_24k_samples / self.SNAC_SAMPLES_PER_FRAME) if total_frames == 0: return [] total_chars = sum(len(w) for w in words) if total_chars == 0: return [] # Pre-tokenize words word_tokens = self._tokenize_words(words) alignments = [] current_frame = 0 for i, word in enumerate(words): # Distribute frames by character count word_frames = int((len(word) / total_chars) * total_frames) if i == len(words) - 1: end_frame = total_frames else: end_frame = min(current_frame + max(1, word_frames), total_frames) start_frame = current_frame start_time = start_frame / self.SNAC_FPS end_time = end_frame / self.SNAC_FPS alignments.append({ 'word': word, 'start': start_time, 'end': end_time, 'start_frame': start_frame, 'end_frame': end_frame, 'tokens': word_tokens[i] }) current_frame = end_frame return alignments def _tokenize_words(self, words: List[str]) -> List[List[int]]: """Tokenize all words.""" if self.tokenizer is None: return [[] for _ in words] return [ self.tokenizer.encode(word, add_special_tokens=False) for word in words ] # ============================================================================= # ENVIRONMENT LOADER # ============================================================================= def load_dotenv(env_path: Optional[str] = None) -> None: """Load environment variables from .env file.""" if env_path is None: for path in [Path(".env"), Path(__file__).parent.parent / ".env"]: if path.exists(): env_path = str(path) break if env_path and Path(env_path).exists(): with open(env_path) as f: for line in f: line = line.strip() if line and not line.startswith("#") and "=" in line: key, value = line.split("=", 1) os.environ.setdefault(key.strip(), value.strip()) log(f"[ENV] Loaded from {env_path}") # ============================================================================= # WORKER FUNCTIONS (for multiprocessing) # ============================================================================= def qa_producer( config: PipelineConfig, target_count: int, tts_queue: mp.Queue, status_queue: mp.Queue, num_workers: int ) -> None: """Produces Q&A batches for TTS pipeline.""" generator = QAGenerator(config) seen = set() pairs = [] pending = [] batch_idx = 0 t0 = time.time() batch_size = config.batch_sizes.tts while len(pairs) < target_count: with ThreadPoolExecutor(max_workers=config.groq_parallel_requests) as ex: futures = [ ex.submit(generator.generate_batch, i) for i in range(config.groq_parallel_requests) ] for f in as_completed(futures): for p in f.result(): if len(pairs) >= target_count: break norm = p["q"].lower().strip().rstrip("?") if norm not in seen: seen.add(norm) pairs.append(p) pending.append(p) if len(pending) >= batch_size: tts_queue.put((batch_idx, pending[:batch_size])) batch_idx += 1 pending = pending[batch_size:] if len(pairs) >= target_count: break elapsed = time.time() - t0 rate = len(pairs) / elapsed if elapsed > 0 else 0 status_queue.put(("qa", len(pairs), target_count, rate)) if pending: tts_queue.put((batch_idx, pending)) batch_idx += 1 for _ in range(num_workers): tts_queue.put(None) status_queue.put(("qa_done", len(pairs), batch_idx)) def tts_worker( gpu_id: int, tts_queue: mp.Queue, feat_queue: mp.Queue, status_queue: mp.Queue, batch_sizes: BatchSizeConfig, num_gpus: int = 1, gpu_offset: int = 0 ) -> None: """TTS worker - converts text to speech.""" import torch _orig_load = torch.load torch.load = lambda *a, **kw: _orig_load(*a, **{**kw, 'weights_only': False}) # Setup GPU actual_gpu = gpu_offset + (gpu_id % num_gpus) torch.cuda.set_device(actual_gpu) print(f"[TTS-GPU{gpu_id}] Assigned to CUDA device {actual_gpu}") # Patch JSON encoder for torch.dtype import json _orig_default = json.JSONEncoder.default def _patched_default(self, obj): if isinstance(obj, torch.dtype): return str(obj).split('.')[-1] return _orig_default(self, obj) json.JSONEncoder.default = _patched_default # Load TTS model from soprano import SopranoTTS vram_gb = GPUManager.get_vram_gb(actual_gpu) gpu_name = GPUManager.get_device_name(actual_gpu) if GPUManager.supports_lmdeploy(gpu_name): try: dec_batch = 32 if vram_gb >= 80 else (16 if vram_gb >= 40 else 8) tts = SopranoTTS( backend="lmdeploy", device="cuda", cache_size_mb=4000 if vram_gb >= 24 else 2000, decoder_batch_size=dec_batch, ) print(f"[TTS-GPU{gpu_id}] Using lmdeploy backend") except Exception as e: print(f"[TTS-GPU{gpu_id}] lmdeploy failed ({e}), using transformers") tts = SopranoTTS(backend="transformers", device="cuda") else: print(f"[TTS-GPU{gpu_id}] Using transformers backend") tts = SopranoTTS(backend="transformers", device="cuda") status_queue.put(("tts_ready", gpu_id)) processed = 0 t_start = time.time() last_heartbeat = time.time() while True: # Heartbeat if time.time() - last_heartbeat > TimeoutConfig.HEARTBEAT_INTERVAL: status_queue.put(("tts_heartbeat", gpu_id, processed)) last_heartbeat = time.time() try: item = tts_queue.get(timeout=TimeoutConfig.QUEUE_GET) except: continue if item is None: break batch_idx, pairs = item all_results = [] try: questions = [p["q"] for p in pairs] answers = [p["a"] for p in pairs] combined = questions + answers all_audios = tts.infer_batch(combined) q_audios = all_audios[:len(questions)] a_audios = all_audios[len(questions):] for j, p in enumerate(pairs): q_np = q_audios[j].cpu().numpy() if hasattr(q_audios[j], 'numpy') else q_audios[j] a_np = a_audios[j].cpu().numpy() if hasattr(a_audios[j], 'numpy') else a_audios[j] all_results.append({ "question": p["q"], "answer": p["a"], "q_audio": np.asarray(q_np, dtype=np.float32), "a_audio": np.asarray(a_np, dtype=np.float32), }) except Exception as e: import traceback status_queue.put(("tts_error", gpu_id, f"{e}\n{traceback.format_exc()}")) continue if all_results: feat_queue.put((batch_idx, all_results)) processed += len(all_results) elapsed = time.time() - t_start batch_rate = len(all_results) / (time.time() - t_start) if elapsed > 0 else 0 status_queue.put(("tts", gpu_id, processed, processed/elapsed, batch_rate)) feat_queue.put(None) status_queue.put(("tts_done", gpu_id, processed)) def features_worker( gpu_id: int, feat_queue: mp.Queue, result_queue: mp.Queue, status_queue: mp.Queue, batch_sizes: BatchSizeConfig, num_gpus: int = 1, gpu_offset: int = 0 ) -> None: """Features worker - extracts Whisper features and SNAC tokens.""" import torch _orig = torch.load torch.load = lambda *a, **kw: _orig(*a, **{**kw, 'weights_only': False}) import torchaudio import snac from transformers import AutoTokenizer, WhisperModel, WhisperFeatureExtractor from huggingface_hub import login # Setup GPU actual_gpu = gpu_offset + (gpu_id % num_gpus) torch.cuda.set_device(actual_gpu) device = f"cuda:{actual_gpu}" print(f"[Features-GPU{gpu_id}] Assigned to CUDA device {actual_gpu}") # HuggingFace auth hf_token = os.environ.get("HF_TOKEN") if hf_token: try: login(token=hf_token) except Exception as e: print(f"[WARN] HuggingFace login failed: {e}") # Load models print("[Features] Loading Whisper with SDPA attention...") whisper_model = WhisperModel.from_pretrained( "openai/whisper-large-v3-turbo", torch_dtype=torch.float16, attn_implementation="sdpa" ).to(device).eval() whisper_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-large-v3-turbo") # torch.compile for faster encoder inference (1.5-2.5x speedup) try: whisper_model.encoder = torch.compile(whisper_model.encoder, mode="reduce-overhead") print(f"[Features-GPU{gpu_id}] torch.compile applied to Whisper encoder") except Exception as e: print(f"[Features-GPU{gpu_id}] torch.compile failed ({e}), using eager mode") print("[Features] Whisper loaded") snac_model = snac.SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval() # Load tokenizer tokenizer = None for model_path in ["canopylabs/orpheus-3b-0.1-pretrained", "meta-llama/Llama-3.2-3B"]: try: tokenizer = AutoTokenizer.from_pretrained(model_path, token=hf_token) print(f"[Features] Loaded tokenizer from {model_path}") break except Exception as e: print(f"[WARN] Failed to load tokenizer from {model_path}: {e}") aligner = WordAligner(tokenizer) snac_batch = batch_sizes.snac # Whisper batch size: maximize GPU utilization (encoder uses ~300MB per sample) vram = GPUManager.get_vram_gb(actual_gpu) if vram >= 24: whisper_batch = 32 elif vram >= 16: whisper_batch = 20 else: whisper_batch = 8 print(f"[Features-GPU{gpu_id}] Whisper batch={whisper_batch}, SNAC batch={snac_batch}") # Pre-compute resampling kernel on GPU for SNAC (32kHz -> 24kHz) _resample_fn = torchaudio.transforms.Resample(32000, 24000).to(device) # Pre-compute resampling kernel for Whisper (32kHz -> 16kHz) on GPU _resample_16k_fn = torchaudio.transforms.Resample(32000, 16000).to(device) status_queue.put(("feat_ready", gpu_id)) # Warmup torch.compile with a dummy forward pass try: dummy_mel = torch.randn(1, 128, 3000, device=device, dtype=torch.float16) with torch.no_grad(): _ = whisper_model.encoder(dummy_mel) del dummy_mel torch.cuda.empty_cache() print(f"[Features-GPU{gpu_id}] Warmup complete") except Exception as e: print(f"[Features-GPU{gpu_id}] Warmup failed: {e}") def process_whisper_batch(audio_list): """Process batch of audios with Whisper in single forward pass (optimized).""" max_samples = 480000 # 30s at 16kHz WHISPER_MEL_LENGTH = 3000 # Whisper expects exactly 3000 mel frames n_mels = whisper_extractor.feature_size # 128 for large-v3-turbo # Batch resample on GPU: stack all audios, resample together max_raw_len = max(a.shape[0] for a in audio_list) padded_raw = np.zeros((len(audio_list), max_raw_len), dtype=np.float32) for i, a in enumerate(audio_list): padded_raw[i, :a.shape[0]] = a raw_tensor = torch.from_numpy(padded_raw).to(device) resampled_16k = _resample_16k_fn(raw_tensor).cpu() # [B, T_16k] del raw_tensor # Truncate to max 30s if resampled_16k.shape[1] > max_samples: resampled_16k = resampled_16k[:, :max_samples] # Batch mel extraction using WhisperFeatureExtractor # Process all audios at once (the extractor supports batch input) audios_np = [resampled_16k[i].numpy() for i in range(len(audio_list))] inputs = whisper_extractor( audios_np, sampling_rate=16000, return_tensors="pt", padding=True ) mel_batch = inputs.input_features # [B, n_mels, T] # Pad or truncate to exactly 3000 frames T = mel_batch.shape[-1] if T < WHISPER_MEL_LENGTH: mel_batch = torch.nn.functional.pad(mel_batch, (0, WHISPER_MEL_LENGTH - T)) elif T > WHISPER_MEL_LENGTH: mel_batch = mel_batch[:, :, :WHISPER_MEL_LENGTH] # Forward pass input_features = mel_batch.to(device, dtype=torch.float16) with torch.no_grad(): encoder_outputs = whisper_model.encoder(input_features) # Split back to individual features features = encoder_outputs.last_hidden_state.cpu().half() return [features[i] for i in range(len(audio_list))] processed = 0 t_start = time.time() last_heartbeat = time.time() while True: # Heartbeat if time.time() - last_heartbeat > TimeoutConfig.HEARTBEAT_INTERVAL: status_queue.put(("feat_heartbeat", gpu_id, processed)) last_heartbeat = time.time() try: item = feat_queue.get(timeout=TimeoutConfig.QUEUE_GET) except: continue if item is None: break batch_idx, audio_batch = item t0 = time.time() try: # 1. Whisper encoding (batched) q_audios = [ad["q_audio"] for ad in audio_batch] whisper_features = [] # Process in mini-batches for memory efficiency for start in range(0, len(q_audios), whisper_batch): end = min(start + whisper_batch, len(q_audios)) batch_features = process_whisper_batch(q_audios[start:end]) whisper_features.extend(batch_features) # 2. SNAC encoding (GPU-accelerated resampling) a_audios = [ad["a_audio"] for ad in audio_batch] all_tokens = [] for start in range(0, len(a_audios), snac_batch): end = min(start + snac_batch, len(a_audios)) mini_audios = a_audios[start:end] max_len = max(a.shape[0] for a in mini_audios) padded = [np.pad(a, (0, max_len - len(a))) for a in mini_audios] # Batch resample on GPU (32kHz -> 24kHz) padded_tensor = torch.from_numpy(np.stack(padded)).to(device) audios_24k = _resample_fn(padded_tensor).cpu() del padded_tensor with torch.no_grad(): codes = snac_model.encode(audios_24k.unsqueeze(1).to(device)) for i in range(len(mini_audios)): all_tokens.append(SNACTokenExtractor.extract(codes, i)) torch.cuda.synchronize() # Memory cleanup if processed > 0 and processed % MemoryConfig.CLEANUP_INTERVAL_ITEMS == 0: GPUManager.clear_memory() # 3. Build results results = [] for i, ad in enumerate(audio_batch): result = { "whisper_features": whisper_features[i], "snac_tokens": torch.tensor(all_tokens[i], dtype=torch.long), "text": ad["question"], "answer": ad["answer"] } if tokenizer is not None: text_tokens = tokenizer.encode(ad["answer"], add_special_tokens=False) result["text_tokens"] = torch.tensor(text_tokens, dtype=torch.long) alignments = aligner.align_proportional(ad["a_audio"], ad["answer"]) if alignments: result["word_alignments"] = alignments results.append(result) batch_time = time.time() - t0 batch_rate = len(results) / batch_time if batch_time > 0 else 0 # Put results with retry for attempt in range(MemoryConfig.MAX_PUT_ATTEMPTS): try: result_queue.put((batch_idx, results), timeout=TimeoutConfig.QUEUE_PUT) break except Exception: if attempt < MemoryConfig.MAX_PUT_ATTEMPTS - 1: status_queue.put(("feat_warn", gpu_id, f"Queue full, retry {attempt+1}")) time.sleep(1) else: raise processed += len(results) elapsed = time.time() - t_start status_queue.put(("feat", gpu_id, processed, processed/elapsed, batch_rate)) # Cleanup del whisper_features, all_tokens, results except Exception as e: import traceback error_msg = str(e) status_queue.put(("feat_error", gpu_id, f"{error_msg}\n{traceback.format_exc()}")) GPUManager.clear_memory() if "out of memory" in error_msg.lower(): status_queue.put(("feat_warn", gpu_id, "OOM detected, clearing memory...")) time.sleep(2) continue status_queue.put(("feat_done", gpu_id, processed)) # ============================================================================= # PIPELINE MONITOR (with incremental saving) # ============================================================================= class PipelineMonitor: """ Monitors pipeline workers and saves results incrementally. Saves after EVERY batch - crash resilient. """ def __init__( self, config: PipelineConfig, status_queue: mp.Queue, result_queue: mp.Queue, num_tts_workers: int, num_feat_workers: int, saver: IncrementalSaver ): self.config = config self.status_queue = status_queue self.result_queue = result_queue self.num_tts_workers = num_tts_workers self.num_feat_workers = num_feat_workers self.saver = saver # Buffer for out-of-order batches self.pending_batches: Dict[int, List] = {} self.next_batch_to_save = 0 self.total_collected = 0 # Items collected (may not be saved yet) self.total_saved = 0 # Items saved to disk self.tts_done_count = 0 self.feat_done_count = 0 self.errors: List = [] self.last_result_time = time.time() self.last_heartbeat_time = { f"tts_{i}": time.time() for i in range(num_tts_workers) } self.last_heartbeat_time.update({ f"feat_{i}": time.time() for i in range(num_feat_workers) }) def run(self, target_count: int, workers: List) -> int: """ Main monitoring loop. Saves incrementally after each batch. Returns total saved count. """ t0 = time.time() stall_warning_shown = False start_count = self.saver.get_count() while True: # Process status messages self._process_status_messages() # Collect and save results saved_now = self._collect_and_save() if saved_now > 0: elapsed = time.time() - t0 total = self.saver.get_count() rate = (total - start_count) / elapsed if elapsed > 0 else 0 log(f"[Saved] {total}/{target_count} | {rate:.1f}/s") stall_warning_shown = False self.last_result_time = time.time() # Check exit conditions if self.saver.get_count() >= target_count: log(f"[Main] Target reached: {self.saver.get_count()}/{target_count}") break if self.feat_done_count >= self.num_feat_workers: log("[Main] All workers done, draining queue...") self._final_drain() break # Stall detection time_since_result = time.time() - self.last_result_time if time_since_result > TimeoutConfig.STALL_WARNING and not stall_warning_shown: stuck = self._get_stuck_workers() pending_info = f"pending batches: {len(self.pending_batches)}, next: {self.next_batch_to_save}" log(f"[WARN] No results for {int(time_since_result)}s ({pending_info})") if stuck: log(f"[WARN] Stuck workers: {stuck}") stall_warning_shown = True if time_since_result > TimeoutConfig.STALL_EXIT: if self._all_workers_stuck(): log(f"[WARN] All workers stuck, stopping with {self.saver.get_count()} items") break if time_since_result > TimeoutConfig.NO_PROGRESS_EXIT: log(f"[WARN] No progress for 10min, stopping with {self.saver.get_count()} items") break return self.saver.get_count() def _collect_and_save(self) -> int: """ Collect results and save contiguous batches immediately. Returns count of items saved this call. """ saved = 0 drain_start = time.time() # Collect available batches while time.time() - drain_start < TimeoutConfig.DRAIN_LOOP: try: batch_idx, items = self.result_queue.get_nowait() self.pending_batches[batch_idx] = items self.total_collected += len(items) except: time.sleep(0.05) try: batch_idx, items = self.result_queue.get_nowait() self.pending_batches[batch_idx] = items self.total_collected += len(items) except: break # Save contiguous batches in order while self.next_batch_to_save in self.pending_batches: batch = self.pending_batches.pop(self.next_batch_to_save) self.saver.add_batch(batch) saved += len(batch) self.next_batch_to_save += 1 self.total_saved += saved return saved def _final_drain(self) -> None: """Final drain of result queue and save all pending.""" drain_start = time.time() # Drain queue while time.time() - drain_start < 60.0: try: batch_idx, items = self.result_queue.get(timeout=0.5) self.pending_batches[batch_idx] = items except: break # Save all remaining in order for batch_idx in sorted(self.pending_batches.keys()): if batch_idx >= self.next_batch_to_save: batch = self.pending_batches.pop(batch_idx) self.saver.add_batch(batch) self.next_batch_to_save = batch_idx + 1 log(f"[Drain] Saved all pending, total: {self.saver.get_count()}") def _process_status_messages(self) -> None: """Process all pending status messages.""" for _ in range(100): try: msg = self.status_queue.get_nowait() msg_type = msg[0] if msg_type == "tts_ready": log(f"[TTS-GPU{msg[1]}] Ready") elif msg_type == "feat_ready": log(f"[Features-GPU{msg[1]}] Ready") elif msg_type == "qa": log(f"[Q&A] {msg[1]}/{msg[2]} | {msg[3]:.1f}/s") elif msg_type == "qa_done": log(f"[Q&A] Done: {msg[1]} pairs") elif msg_type == "tts": log(f"[TTS-GPU{msg[1]}] {msg[2]} items | avg {msg[3]:.1f}/s") self.last_heartbeat_time[f"tts_{msg[1]}"] = time.time() elif msg_type == "tts_done": self.tts_done_count += 1 log(f"[TTS-GPU{msg[1]}] Done: {msg[2]} items") elif msg_type == "feat": log(f"[Feat-GPU{msg[1]}] {msg[2]} items | avg {msg[3]:.1f}/s") self.last_heartbeat_time[f"feat_{msg[1]}"] = time.time() elif msg_type == "feat_done": self.feat_done_count += 1 log(f"[Features-GPU{msg[1]}] Done: {msg[2]} items") elif "heartbeat" in msg_type: worker_type = "tts" if "tts" in msg_type else "feat" self.last_heartbeat_time[f"{worker_type}_{msg[1]}"] = time.time() elif "error" in msg_type: log(f"[Error] {msg}") self.errors.append(msg) elif "warn" in msg_type: log(f"[WARN] {msg[2]}") except: break def _get_stuck_workers(self) -> List[str]: """Get list of stuck workers.""" stuck = [] for worker_id, last_hb in self.last_heartbeat_time.items(): if time.time() - last_hb > TimeoutConfig.STUCK_WORKER_THRESHOLD: stuck.append(worker_id) return stuck def _all_workers_stuck(self) -> bool: """Check if all workers are stuck.""" stuck = self._get_stuck_workers() return len(stuck) >= (self.num_tts_workers + self.num_feat_workers) # ============================================================================= # MAIN PIPELINE # ============================================================================= def create_config_from_args(args) -> PipelineConfig: """Create pipeline config from command line args.""" vram_gb = GPUManager.get_vram_gb() if torch.cuda.is_available() else 0 batch_sizes = BatchSizeConfig.from_vram(vram_gb) # Apply overrides if args.tts_batch: batch_sizes.tts = args.tts_batch if args.snac_batch: batch_sizes.snac = args.snac_batch if args.whisper_workers: batch_sizes.whisper_workers = args.whisper_workers return PipelineConfig( output_path=args.output, target_count=args.count, num_gpus=args.gpus, log_file=args.log_file or f"{args.output}.log", batch_sizes=batch_sizes, groq_api_key=os.environ.get("GROQ_API_KEY", ""), ) def main(): global _logger mp.set_start_method('spawn', force=True) # Increase file descriptors import resource soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) resource.setrlimit(resource.RLIMIT_NOFILE, (min(65536, hard), hard)) # Load environment load_dotenv() # Parse arguments import argparse parser = argparse.ArgumentParser(description="Dataset Generator v2") parser.add_argument("--count", "--num_samples", type=int, default=1000, dest="count") parser.add_argument("--output", type=str, default="./data/dataset.pt") parser.add_argument("--gpus", type=int, default=GPUManager.get_num_gpus() or 1) parser.add_argument("--log-file", type=str, default=None) # Note: No --resume flag needed, always resumes automatically from output file parser.add_argument("--tts-batch", type=int, default=None) parser.add_argument("--snac-batch", type=int, default=None) parser.add_argument("--whisper-workers", type=int, default=None) args = parser.parse_args() # Create config config = create_config_from_args(args) # Setup logging Path(config.output_path).parent.mkdir(parents=True, exist_ok=True) _logger = PipelineLogger(config.log_file) # Print config log("=" * 60) log("Dataset Generator v2 - Refactored Pipeline") log(f"Target: {config.target_count} items, GPUs: {config.num_gpus}") vram_gb = GPUManager.get_vram_gb() if torch.cuda.is_available() else 0 log(f"[Config] GPU VRAM: {vram_gb:.1f}GB") log(f"[Config] Batch sizes: TTS={config.batch_sizes.tts}, " f"Whisper={config.batch_sizes.whisper_workers}, " f"SNAC={config.batch_sizes.snac}") log("=" * 60) # Incremental saver - saves after every batch saver = IncrementalSaver(config) # Resume: load existing items from output file start_count = saver.load_existing() if start_count >= config.target_count: log(f"[Resume] Already have {start_count} items, done!") return remaining_count = config.target_count - start_count log(f"[Main] Need {remaining_count} more items (have {start_count}/{config.target_count})") total_start = time.time() # Create queues tts_queue = mp.Queue() feat_queue = mp.Queue() result_queue = mp.Queue() status_queue = mp.Queue() # Determine worker counts actual_num_gpus = GPUManager.get_num_gpus() or 1 if actual_num_gpus >= 6: # Features is the bottleneck: allocate more GPUs to it (2 TTS + rest Features) tts_gpus = 2 feat_gpus = actual_num_gpus - tts_gpus tts_gpu_offset = 0 feat_gpu_offset = tts_gpus log(f"[Main] {actual_num_gpus} GPUs - TTS: 0-{tts_gpus-1}, Features: {feat_gpu_offset}-{actual_num_gpus-1}") elif actual_num_gpus >= 4: tts_gpus = actual_num_gpus // 2 feat_gpus = actual_num_gpus - tts_gpus tts_gpu_offset = 0 feat_gpu_offset = tts_gpus log(f"[Main] {actual_num_gpus} GPUs - TTS: 0-{tts_gpus-1}, Features: {feat_gpu_offset}-{actual_num_gpus-1}") else: tts_gpus = actual_num_gpus feat_gpus = actual_num_gpus tts_gpu_offset = 0 feat_gpu_offset = 0 log(f"[Main] {actual_num_gpus} GPUs, shared between TTS and Features") num_tts_workers = min(config.num_gpus, tts_gpus) num_feat_workers = min(config.num_gpus, feat_gpus) log(f"[Main] Spawning {num_tts_workers} TTS, {num_feat_workers} Features workers") # Start workers workers = [] qa_proc = mp.Process( target=qa_producer, args=(config, remaining_count, tts_queue, status_queue, num_tts_workers) ) qa_proc.start() workers.append(qa_proc) for gpu_id in range(num_tts_workers): p = mp.Process( target=tts_worker, args=(gpu_id, tts_queue, feat_queue, status_queue, config.batch_sizes, tts_gpus, tts_gpu_offset) ) p.start() workers.append(p) for gpu_id in range(num_feat_workers): p = mp.Process( target=features_worker, args=(gpu_id, feat_queue, result_queue, status_queue, config.batch_sizes, feat_gpus, feat_gpu_offset) ) p.start() workers.append(p) log("[Pipeline] All workers started, saving incrementally...") # Monitor pipeline - saves after every batch monitor = PipelineMonitor( config, status_queue, result_queue, num_tts_workers, num_feat_workers, saver ) final_count = monitor.run(config.target_count, workers) # Cleanup workers log("[Main] Waiting for workers...") for p in workers: p.join(timeout=TimeoutConfig.WORKER_JOIN) if p.is_alive(): p.terminate() # Finalize final_count = saver.finalize() # Summary total_time = time.time() - total_start new_count = final_count - start_count log("\n" + "=" * 60) log(f"COMPLETE: {final_count} items in {config.output_path}") if start_count > 0: log(f" (resumed from {start_count}, added {new_count} new)") log(f"Total time: {total_time:.1f}s ({total_time/60:.1f}m)") throughput = new_count / total_time if total_time > 0 else 0 log(f"Throughput: {throughput:.2f} items/s") log("=" * 60) if __name__ == "__main__": main()