datasets/create_dataset_v2.py

#!/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()