Initial commit: Speech-to-Speech dataset generation pipeline

- passo0_setup.py: Environment setup for GPU training
- passo1_convertdataset.py: Convert raw data to InstructS2S format
- passo2_finetune_stage1.py: Stage 1 training (adapter only)
- passo3_finetune_stage2.py: Stage 2 training (adapter + LoRA)
- datasets/generate_dataset.py: Generate Q&A dataset with TTS/STT
- services/: Modular services for Groq, WhisperX, Kokoro TTS, Llama

All API keys use environment variables for security.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,27 @@
+# OS
+.DS_Store
+
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+dist/
+build/
+
+# Data and checkpoints
+data/
+checkpoints/
+logs/
+*.pt
+*.gguf
+*.onnx
+*.bin
+
+# Temp
+*.tmp
+temp_*/
+
+# IDE
+.vscode/
+.idea/

claude.md

@@ -0,0 +1,122 @@
+# claude.md - Contexto do Projeto
+
+## Objetivo
+Modelo **Speech-to-Speech** end-to-end: Áudio entrada → LLM → Áudio resposta (sem texto intermediário).
+
+## Arquitetura
+
+```
+Áudio Usuário → Whisper → SpeechAdapter → Orpheus+LoRA → Interleaved Output → SNAC → Áudio
+                           (Stage 1+2)    (Stage 2 only)
+```
+
+## Componentes
+
+| Componente | Estrutura | Parâmetros | Treina em |
+|------------|-----------|------------|-----------|
+| **Whisper** | whisper-large-v3 (1280 dim) | Frozen | - |
+| **SpeechAdapter** | 5× downsample + FFN(6400→2048→3072) + LayerNorm | ~19M | Stage 1, Stage 2 |
+| **Orpheus** | LLaMA 3B (orpheus-3b-es-it) | Frozen (Stage 1), LoRA (Stage 2) | Stage 2 |
+| **LoRA** | r=16, alpha=32, all projections | ~50M | Stage 2 |
+
+## SNAC Token Offsets (CRÍTICO)
+
+Tokens SNAC usam offsets position-based para Orpheus:
+```python
+SNAC_BASE = 128266
+EOS_TOKEN = 128009
+# 7 tokens por frame: pos 0-6 têm offsets diferentes
+# Offset = 128266 + (pos % 7) * 4096
+```
+
+## Pipeline de Training (IST-LM + LLaMA-Omni 2)
+
+**IMPORTANTE: NÃO crie arquivos temporários, checkpoints, logs ou datasets neste diretório.**
+Use subdiretórios como `./checkpoints`, `./data`, `./logs` para manter este diretório limpo.
+
+Pipeline de 4 passos com **Scheduled Interleaved Speech-Text**:
+- Começa com 90% texto, diminui 0.1 a cada 300 steps até 100% áudio
+
+### Passo 0: Setup do Ambiente
+```bash
+python passo0_setup.py
+```
+- Instala todas as dependências (torch, transformers, snac, etc.)
+- Verifica CUDA/GPU
+- Cria diretórios necessários
+- Otimiza automaticamente batch size com base na VRAM
+
+### Passo 1: ConvertDataset
+```bash
+python passo1_convertdataset.py --input_dir ./data/raw --output ./data/processed/data.pt
+```
+- Converte datasets brutos para formato InstructS2S
+- Extrai Whisper features e SNAC tokens
+- Gerenciamento automático de memória (evita OOM)
+- Output: `data.pt`
+
+### Passo 2: Stage 1 - Adapter Only (LLM Frozen)
+```bash
+python passo2_finetune_stage1.py --data ./data/processed/data.pt --epochs 2 --lr 5e-5 --output_dir ./checkpoints
+```
+- Treina: SpeechAdapter only
+- Frozen: Orpheus (completamente)
+- Output: Interleaved text+audio tokens
+- Epochs: 1-2 (warmup)
+- Output: `checkpoints/stage1_best.pt`
+
+### Passo 3: Stage 2 - Adapter + LoRA Together
+```bash
+python passo3_finetune_stage2.py --data ./data/processed/data.pt --stage1_ckpt ./checkpoints/stage1_best.pt --epochs 3 --lr 5e-5 --output_dir ./checkpoints
+```
+- Treina: SpeechAdapter + LoRA (juntos)
+- Output: Interleaved text+audio tokens
+- Epochs: 3+
+- Output: `checkpoints/stage2_best.pt`
+
+## Scheduled Interleaving (IST-LM Paper)
+
+```python
+def get_text_ratio(global_step, decay_steps=300, initial_ratio=0.9):
+    num_decays = global_step // decay_steps
+    text_ratio = initial_ratio - (num_decays * 0.1)
+    return max(0.0, text_ratio)
+```
+
+Pattern por text_ratio:
+- 0.9: 1 text token + 3 audio frames (21 tokens)
+- 0.7: 1 text token + 5 audio frames
+- 0.5: 1 text token + 7 audio frames
+- 0.3: 1 text token + 10 audio frames
+- 0.0: Pure audio
+
+## Otimizações Comuns
+
+- **GPU Auto-Detection**: B200/A100/RTX4090 → batch/dtype automático
+- **Label Smoothing**: `--label_smoothing 0.1`
+- **Dynamic Loading**: Datasets carregam em background via mmap
+- **Async Checkpoints**: Saves não bloqueiam training
+
+## Dataset
+
+**ATENÇÃO: Dataset InstructS2S-200K localizado em `/workspace/dataset/InstructS2S-200K/`**
+- Arquivos de origem (en_part_00 a en_part_32): 33 arquivos de ~10GB cada
+- **NÃO DELETE** os arquivos em `/workspace/dataset/InstructS2S-200K/`
+- Arquivos extraídos serão salvos em `/root/InstructS2S-200K/extracted_data/`
+
+Formato InstructS2S:
+```python
+{
+    "whisper_features": [seq_len, 1280],  # Whisper do áudio USUÁRIO
+    "snac_tokens": [num_tokens],          # SNAC do áudio RESPOSTA
+    "text": str,                          # Texto da pergunta
+    "answer": str                         # Texto da resposta (para interleaving)
+}
+```
+
+## Regras
+
+- **Whisper**: Usar `whisper-large-v3` (1280 dim) - NUNCA whisper-small
+- **HF Token**: `$HF_TOKEN`
+- **SNAC**: Truncar para múltiplo de 7 (frames completos)
+- **Interleaved**: Sempre usar scheduled interleaving, nunca pure text ou pure audio fixo

datasets/generate_dataset.py

@@ -0,0 +1,848 @@
+#!/usr/bin/env python3
+"""
+Full GPU pipeline with Kokoro TTS - MULTI-GPU VERSION.
+
+Features:
+- Auto-detects all available GPUs
+- Distributes workload across GPUs
+- Each GPU runs independent pipeline
+- Merges results at the end
+- Single WhisperX model for BOTH alignment AND feature extraction (no duplicate Whisper!)
+- SNAC with torch.compile optimization
+
+Usage:
+    python generate_kokoro_gpu.py --count 60000 --output-dir ./qa_dataset
+
+With 1 GPU at ~5 items/s = 60k in ~3.3 hours
+"""
+
+import os
+import sys
+
+# CRITICAL: Set LD_LIBRARY_PATH for ONNX CUDA before any imports
+# This must happen before importing onnxruntime or kokoro_onnx
+nvidia_libs = []
+for pyver in ['python3.12', 'python3.11', 'python3.10']:
+    nvidia_base = f'/usr/local/lib/{pyver}/dist-packages/nvidia'
+    for subdir in ['cublas', 'cudnn', 'cufft', 'curand', 'cusolver', 'cusparse',
+                   'cuda_runtime', 'cuda_nvrtc', 'cuda_cupti', 'nvjitlink', 'nccl', 'nvtx', 'cusparselt', 'cufile']:
+        lib_path = f"{nvidia_base}/{subdir}/lib"
+        if os.path.exists(lib_path):
+            nvidia_libs.append(lib_path)
+
+os.environ['LD_LIBRARY_PATH'] = ':'.join(nvidia_libs) + ':' + os.environ.get('LD_LIBRARY_PATH', '')
+
+# CRITICAL: Force Kokoro TTS to use CUDA provider (not CPU)
+os.environ['ONNX_PROVIDER'] = 'CUDAExecutionProvider'
+
+import json
+import time
+import argparse
+import threading
+import queue
+import multiprocessing as mp
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from pathlib import Path
+from typing import List, Dict, Tuple, Optional
+import numpy as np
+import torch
+import torchaudio
+import soundfile as sf
+
+# Configuration
+GROQ_API_KEY = os.environ.get('GROQ_API_KEY', '')
+GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
+MODEL = "llama-3.1-8b-instant"
+API_BATCH_SIZE = 20  # QA pairs per LLM request (15-20 works best)
+
+# Local Llama configuration
+LLAMA_MODEL_PATH = "/tmp/llama-3.2-3b-instruct-q4_k_m.gguf"
+USE_LOCAL_LLAMA = False  # Set via --use-llama flag
+LLAMA_GPU_ID = 1  # Use GPU 1 for Llama (GPU 0 for TTS/WhisperX)
+
+KOKORO_MODEL_PATH = '/tmp/kokoro-v1.0.onnx'
+KOKORO_VOICES_PATH = '/tmp/voices-v1.0.bin'
+SAMPLE_RATE = 24000
+
+WHISPERX_MODEL = "large-v3-turbo"  # Single model for alignment AND feature extraction
+WHISPERX_COMPUTE_TYPE = "float16"  # float16 for RTX 50xx
+SNAC_MODEL = "hubertsiuzdak/snac_24khz"
+
+BATCH_SIZE = 50
+NUM_TTS_WORKERS = 1  # Single TTS model to avoid GPU contention
+SILENCE_BETWEEN_PAIRS = 0.5
+IO_WORKERS = 4
+
+
+class PipelineStats:
+    def __init__(self, target_count: int, gpu_id: int = 0):
+        self.lock = threading.Lock()
+        self.target_count = target_count
+        self.gpu_id = gpu_id
+        self.api_pairs_sent = 0
+        self.tts_batches = 0
+        self.feature_batches = 0
+        self.total_items = 0
+        self.start_time = time.time()
+        self.stop_event = threading.Event()
+
+    def add_api_pairs(self, count: int):
+        with self.lock:
+            self.api_pairs_sent += count
+
+    def get_api_pairs_sent(self):
+        with self.lock:
+            return self.api_pairs_sent
+
+    def update(self, stage: str, count: int = 0):
+        with self.lock:
+            if stage == 'tts':
+                self.tts_batches += 1
+            elif stage == 'features':
+                self.feature_batches += 1
+                self.total_items += count
+                if self.total_items >= self.target_count:
+                    self.stop_event.set()
+
+    def should_stop(self):
+        return self.stop_event.is_set()
+
+    def get_total_items(self):
+        with self.lock:
+            return self.total_items
+
+    def print_status(self):
+        with self.lock:
+            elapsed = time.time() - self.start_time
+            rate = self.total_items / elapsed if elapsed > 0 else 0
+            remaining = self.target_count - self.total_items
+            eta_seconds = remaining / rate if rate > 0 else 0
+            eta_min = int(eta_seconds // 60)
+            eta_sec = int(eta_seconds % 60)
+            print(f"\r[GPU {self.gpu_id}] {self.total_items}/{self.target_count} | {rate:.1f}/s | "
+                  f"ETA: {eta_min}m{eta_sec:02d}s", end='', flush=True)
+
+
+def generate_qa_batch_groq(batch_size: int) -> List[Dict[str, str]]:
+    """Generate Q&A pairs using Groq API."""
+    import requests
+    headers = {
+        "Authorization": f"Bearer {GROQ_API_KEY}",
+        "Content-Type": "application/json"
+    }
+    prompt = f"""Generate {batch_size} unique question-answer pairs about general knowledge.
+REQUIREMENTS:
+- Questions: 5 to 10 words
+- Answers: 5 to 10 words
+Format: Q: [question]
+A: [answer]
+Return exactly {batch_size} pairs."""
+
+    payload = {
+        "messages": [{"role": "user", "content": prompt}],
+        "model": MODEL,
+        "temperature": 0.7,
+        "max_tokens": 2000
+    }
+    try:
+        response = requests.post(GROQ_API_URL, headers=headers, json=payload, timeout=30)
+        response.raise_for_status()
+        content = response.json()['choices'][0]['message']['content']
+        return parse_qa_pairs(content)
+    except Exception as e:
+        return []
+
+
+def generate_qa_batch_llama(batch_size: int, llm) -> List[Dict[str, str]]:
+    """Generate Q&A pairs using local Llama model."""
+    prompt = f"""<|start_header_id|>system<|end_header_id|>
+Generate concise Q&A pairs. Keep answers SHORT (5-10 words max).<|eot_id|><|start_header_id|>user<|end_header_id|>
+Generate {batch_size} Q&A pairs about general knowledge.
+STRICT FORMAT - each answer must be 5-10 words only:
+Q: What is the capital of France?
+A: Paris is the capital of France.
+Q: Who invented the telephone?
+A: Alexander Graham Bell invented the telephone.<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+"""
+    try:
+        # ~35 tokens per QA pair (Q + A + formatting)
+        response = llm(prompt, max_tokens=batch_size * 35, temperature=0.8, stop=["<|eot_id|>"])
+        content = response['choices'][0]['text']
+        return parse_qa_pairs(content)
+    except Exception as e:
+        print(f"[LLM] Error: {e}", flush=True)
+        return []
+
+
+def generate_qa_batch(batch_size: int, llm=None) -> List[Dict[str, str]]:
+    """Generate Q&A pairs using either Groq API or local Llama."""
+    if USE_LOCAL_LLAMA and llm is not None:
+        return generate_qa_batch_llama(batch_size, llm)
+    else:
+        return generate_qa_batch_groq(batch_size)
+
+
+def parse_qa_pairs(content: str) -> List[Dict[str, str]]:
+    pairs = []
+    lines = content.split('\n')
+    current_q, current_a = None, None
+    for line in lines:
+        line = line.strip()
+        if line.lower().startswith('q:'):
+            current_q = line[2:].strip()
+        elif line.lower().startswith('a:'):
+            current_a = line[2:].strip()
+            if current_q and current_a:
+                # Accept 3-20 words for both Q and A (more lenient for local LLM)
+                q_words = len(current_q.split())
+                a_words = len(current_a.split())
+                if 3 <= q_words <= 20 and 3 <= a_words <= 20:
+                    pairs.append({'question': current_q, 'answer': current_a})
+                current_q, current_a = None, None
+    return pairs
+
+
+def api_worker(output_queue: queue.Queue, total_needed: int, stats: PipelineStats):
+    collected = []
+    batches_sent = 0
+    total_batches_needed = (total_needed + BATCH_SIZE - 1) // BATCH_SIZE + 1
+
+    # Load Llama model if using local LLM
+    llm = None
+    if USE_LOCAL_LLAMA:
+        try:
+            from llama_cpp import Llama
+            print(f"[LLM] Loading Llama model on GPU {LLAMA_GPU_ID}...", flush=True)
+
+            # Set CUDA device for Llama
+            os.environ['CUDA_VISIBLE_DEVICES'] = str(LLAMA_GPU_ID)
+
+            llm = Llama(
+                model_path=LLAMA_MODEL_PATH,
+                n_gpu_layers=-1,
+                n_ctx=2048,
+                n_batch=512,
+                verbose=False
+            )
+            # Warmup
+            llm("Hello", max_tokens=10)
+            print(f"[LLM] Llama loaded on GPU {LLAMA_GPU_ID}", flush=True)
+        except Exception as e:
+            print(f"[LLM] Failed to load Llama: {e}, falling back to Groq API", flush=True)
+            llm = None
+
+    source = "LLM" if llm else "API"
+    print(f"[{source}] Starting, need {total_batches_needed} batches", flush=True)
+
+    while batches_sent < total_batches_needed and not stats.should_stop():
+        pairs_in_queue = stats.get_api_pairs_sent() - stats.get_total_items()
+        if pairs_in_queue > total_needed * 0.3:
+            time.sleep(0.5)
+            continue
+
+        needed = min(API_BATCH_SIZE, total_needed - len(collected) - stats.get_api_pairs_sent() + batches_sent * BATCH_SIZE)
+        if needed <= 0:
+            break
+
+        print(f"[{source}] Generating {needed} QA pairs...", flush=True)
+        pairs = generate_qa_batch(needed, llm)
+        print(f"[{source}] Got {len(pairs)} pairs", flush=True)
+        collected.extend(pairs)
+
+        while len(collected) >= BATCH_SIZE and batches_sent < total_batches_needed:
+            batch = collected[:BATCH_SIZE]
+            collected = collected[BATCH_SIZE:]
+            output_queue.put(batch)
+            stats.add_api_pairs(len(batch))
+            batches_sent += 1
+
+            if stats.should_stop():
+                break
+
+    if collected and not stats.should_stop():
+        remaining_needed = total_needed - stats.get_api_pairs_sent()
+        if remaining_needed > 0:
+            batch = collected[:remaining_needed]
+            output_queue.put(batch)
+            stats.add_api_pairs(len(batch))
+
+    output_queue.put(None)
+
+
+def tts_worker(input_queue: queue.Queue, output_queue: queue.Queue, stats: PipelineStats):
+    from kokoro_onnx import Kokoro
+
+    print(f"[TTS] Loading {NUM_TTS_WORKERS} Kokoro models...", flush=True)
+
+    # Kokoro will use CUDA via ONNX_PROVIDER env var set at top of script
+    tts_models = [Kokoro(KOKORO_MODEL_PATH, KOKORO_VOICES_PATH) for _ in range(NUM_TTS_WORKERS)]
+
+    # Log providers to confirm CUDA is being used
+    providers = tts_models[0].sess.get_providers()
+    print(f"[TTS] ONNX Providers: {providers}", flush=True)
+    if 'CUDAExecutionProvider' not in providers:
+        print("[TTS] WARNING: CUDA not available, running on CPU (will be slow!)", flush=True)
+
+    for i, tts in enumerate(tts_models):
+        tts.create("warmup", voice='af_heart', speed=1.0)
+    print(f"[TTS] Models loaded and warmed up", flush=True)
+
+    def process_single_pair(args):
+        idx, pair, tts_model = args
+        try:
+            q_samples, _ = tts_model.create(pair['question'], voice='af_heart', speed=1.0)
+            a_samples, _ = tts_model.create(pair['answer'], voice='af_heart', speed=1.0)
+            return {
+                'question': pair['question'],
+                'answer': pair['answer'],
+                'q_audio': q_samples,
+                'a_audio': a_samples
+            }
+        except Exception as e:
+            print(f"[TTS] Error: {e}", flush=True)
+            return None
+
+    batch_idx = 0
+
+    while not stats.should_stop():
+        try:
+            batch = input_queue.get(timeout=1)
+        except queue.Empty:
+            continue
+
+        if batch is None:
+            output_queue.put(None)
+            break
+
+        batch_idx += 1
+        print(f"[TTS] Processing batch {batch_idx} with {len(batch)} items...", flush=True)
+        audio_data = []
+
+        with ThreadPoolExecutor(max_workers=NUM_TTS_WORKERS) as executor:
+            tasks = [(i, pair, tts_models[i % NUM_TTS_WORKERS]) for i, pair in enumerate(batch)]
+            futures = {executor.submit(process_single_pair, task): task for task in tasks}
+
+            for future in as_completed(futures):
+                if stats.should_stop():
+                    break
+                result = future.result()
+                if result:
+                    audio_data.append(result)
+
+        if stats.should_stop():
+            output_queue.put(None)
+            break
+
+        print(f"[TTS] Batch {batch_idx} done: {len(audio_data)} items", flush=True)
+        stats.update('tts')
+        output_queue.put((batch_idx, audio_data))
+
+
+def batch_extract_whisper_features(
+    audios: List[np.ndarray],
+    whisper_encoder,  # WhisperX internal model (FasterWhisper)
+    device: str
+) -> List[torch.Tensor]:
+    """Extract encoder features using WhisperX's internal faster-whisper model."""
+    import ctranslate2
+
+    if not audios:
+        return []
+
+    all_features = []
+
+    for audio in audios:
+        # Resample to 16kHz for Whisper
+        audio_16k = torchaudio.functional.resample(
+            torch.from_numpy(audio), 24000, 16000
+        ).numpy().astype(np.float32)
+
+        # Use WhisperX's internal feature extractor and encoder
+        mel_features = whisper_encoder.feature_extractor(audio_16k)
+        encoded = whisper_encoder.encode(mel_features)  # Returns ctranslate2 StorageView
+
+        # Convert ctranslate2 StorageView to numpy (must move to CPU first)
+        cpu_view = encoded.to_device(ctranslate2.Device.cpu)
+        features_np = np.array(cpu_view, copy=True)  # [1, seq_len, 1280]
+        features = torch.from_numpy(features_np).squeeze(0).float()  # [seq_len, 1280]
+        all_features.append(features)
+
+    return all_features
+
+
+def batch_extract_snac_tokens(audios: List[np.ndarray], snac_model, device: str) -> List[List[int]]:
+    if not audios:
+        return []
+
+    all_tokens = []
+
+    with torch.no_grad():
+        for audio in audios:
+            audio_tensor = torch.from_numpy(audio).unsqueeze(0).unsqueeze(0).to(device)
+            snac_codes = snac_model.encode(audio_tensor)
+
+            min_len = snac_codes[0].shape[-1]
+            snac_tokens = []
+
+            for i in range(min_len):
+                snac_tokens.append(snac_codes[0][0, i].item() + 128266)
+                if i * 2 + 1 < snac_codes[1].shape[-1]:
+                    snac_tokens.append(snac_codes[1][0, i*2].item() + 128266 + 4096)
+                    snac_tokens.append(snac_codes[1][0, i*2+1].item() + 128266 + 4096)
+                for k in range(4):
+                    if i*4+k < snac_codes[2].shape[-1]:
+                        snac_tokens.append(snac_codes[2][0, i*4+k].item() + 128266 + 2*4096)
+
+            all_tokens.append(snac_tokens)
+
+    return all_tokens
+
+
+def feature_worker(
+    input_queue: queue.Queue,
+    output_queue: queue.Queue,
+    whisperx_model,
+    align_model,
+    align_metadata,
+    whisper_encoder,  # WhisperX internal model for feature extraction
+    snac_model,
+    device: str,
+    output_dir: Path,
+    stats: PipelineStats
+):
+    import whisperx
+
+    while not stats.should_stop():
+        try:
+            data = input_queue.get(timeout=1)
+        except queue.Empty:
+            continue
+
+        if data is None:
+            break
+
+        batch_idx, audio_data = data
+
+        items_needed = stats.target_count - stats.get_total_items()
+        if items_needed <= 0:
+            break
+
+        audio_data = audio_data[:items_needed]
+        if not audio_data:
+            continue
+
+        silence = np.zeros(int(SILENCE_BETWEEN_PAIRS * SAMPLE_RATE), dtype=np.float32)
+        all_answer_audio = []
+        segment_boundaries = []
+        current_pos = 0
+
+        for item in audio_data:
+            a_audio = item['a_audio']
+            start_pos = current_pos
+            all_answer_audio.append(a_audio)
+            all_answer_audio.append(silence)
+            current_pos += len(a_audio) + len(silence)
+            segment_boundaries.append((start_pos, start_pos + len(a_audio), item))
+
+        if not all_answer_audio:
+            continue
+
+        concatenated = np.concatenate(all_answer_audio)
+        temp_path = output_dir / f'temp_batch_{batch_idx}_{stats.gpu_id}.wav'
+        sf.write(temp_path, concatenated, SAMPLE_RATE)
+
+        try:
+            result = whisperx_model.transcribe(str(temp_path), batch_size=16)
+            result = whisperx.align(result["segments"], align_model, align_metadata,
+                                   str(temp_path), device)
+            words = result.get("word_segments", [])
+        except Exception as e:
+            temp_path.unlink(missing_ok=True)
+            continue
+
+        temp_path.unlink(missing_ok=True)
+
+        processed_items = []
+
+        for start_sample, end_sample, item in segment_boundaries:
+            if stats.get_total_items() + len(processed_items) >= stats.target_count:
+                break
+
+            start_time_s = start_sample / SAMPLE_RATE
+            end_time_s = end_sample / SAMPLE_RATE
+
+            segment_words = [w for w in words
+                           if w.get('start', 0) >= start_time_s - 0.1
+                           and w.get('end', 0) <= end_time_s + 0.1]
+
+            adjusted_words = []
+            for w in segment_words:
+                adjusted_words.append({
+                    'word': w.get('word', ''),
+                    'start': w.get('start', 0) - start_time_s,
+                    'end': w.get('end', 0) - start_time_s,
+                    'start_frame': int((w.get('start', 0) - start_time_s) * 75),
+                    'end_frame': int((w.get('end', 0) - start_time_s) * 75)
+                })
+
+            processed_items.append({
+                'item': item,
+                'word_alignments': adjusted_words
+            })
+
+        if not processed_items:
+            continue
+
+        q_audios = [pi['item']['q_audio'] for pi in processed_items]
+        batch_features = batch_extract_whisper_features(
+            q_audios, whisper_encoder, device
+        )
+
+        a_audios = [pi['item']['a_audio'] for pi in processed_items]
+        snac_results = batch_extract_snac_tokens(a_audios, snac_model, device)
+
+        final_items = []
+        for i, pi in enumerate(processed_items):
+            final_items.append({
+                'whisper_features': batch_features[i],
+                'snac_tokens': torch.tensor(snac_results[i], dtype=torch.long),
+                'text': pi['item']['question'],
+                'answer': pi['item']['answer'],
+                'word_alignments': pi['word_alignments'],
+                'q_audio': pi['item']['q_audio'],
+                'a_audio': pi['item']['a_audio']
+            })
+
+        if final_items:
+            output_queue.put(final_items)
+            stats.update('features', len(final_items))
+            stats.print_status()
+
+
+def saver_worker(
+    input_queue: queue.Queue,
+    output_dir: Path,
+    stats: PipelineStats,
+    start_idx: int = 0
+):
+    all_items = []
+    questions_dir = output_dir / 'questions'
+    answers_dir = output_dir / 'answers'
+    questions_dir.mkdir(parents=True, exist_ok=True)
+    answers_dir.mkdir(parents=True, exist_ok=True)
+
+    global_idx = start_idx
+    last_save = 0
+    save_interval = 5000
+
+    io_executor = ThreadPoolExecutor(max_workers=IO_WORKERS)
+    pending_futures = []
+
+    def save_wav(path, audio, sr):
+        sf.write(path, audio, sr)
+
+    while True:
+        try:
+            data = input_queue.get(timeout=2)
+        except queue.Empty:
+            if stats.should_stop() or stats.get_total_items() >= stats.target_count:
+                break
+            continue
+
+        if data is None:
+            break
+
+        for item in data:
+            global_idx += 1
+
+            q_path = questions_dir / f'q_{global_idx:06d}.wav'
+            a_path = answers_dir / f'a_{global_idx:06d}.wav'
+
+            pending_futures.append(io_executor.submit(save_wav, q_path, item['q_audio'], SAMPLE_RATE))
+            pending_futures.append(io_executor.submit(save_wav, a_path, item['a_audio'], SAMPLE_RATE))
+
+            all_items.append({
+                'whisper_features': item['whisper_features'],
+                'snac_tokens': item['snac_tokens'],
+                'text': item['text'],
+                'answer': item['answer'],
+                'word_alignments': item['word_alignments']
+            })
+
+            if len(pending_futures) > 100:
+                pending_futures = [f for f in pending_futures if not f.done()]
+
+        if len(all_items) - last_save >= save_interval:
+            for f in pending_futures:
+                f.result()
+            pending_futures = []
+
+            checkpoint_path = output_dir / f'checkpoint_{len(all_items)}.pt'
+            torch.save(all_items, checkpoint_path)
+            last_save = len(all_items)
+            print(f"\n[GPU {stats.gpu_id}] Checkpoint: {len(all_items)} items")
+
+    for f in pending_futures:
+        try:
+            f.result()
+        except:
+            pass
+
+    io_executor.shutdown(wait=True)
+
+    # Save partial dataset for this GPU
+    partial_path = output_dir / f'dataset_gpu{stats.gpu_id}.pt'
+    torch.save(all_items, partial_path)
+    print(f"\n[GPU {stats.gpu_id}] Saved {len(all_items)} items to {partial_path}")
+
+    return len(all_items)
+
+
+def run_gpu_pipeline(gpu_id: int, count: int, output_dir: Path, start_idx: int = 0) -> int:
+    """Run complete pipeline on a single GPU. Returns number of items generated."""
+
+    # Set GPU - use "cuda" as device since CUDA_VISIBLE_DEVICES makes only one GPU visible
+    os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_id)
+    device = "cuda"  # Not "cuda:0" - WhisperX/ctranslate2 expects just "cuda"
+
+    torch.cuda.set_device(0)
+
+    print(f"\n[GPU {gpu_id}] Starting pipeline for {count} items...")
+    print(f"[GPU {gpu_id}] Device: {torch.cuda.get_device_name(0)}")
+
+    # Load WhisperX - single model for BOTH alignment AND feature extraction
+    import whisperx
+    whisperx_model = whisperx.load_model(WHISPERX_MODEL, device, compute_type=WHISPERX_COMPUTE_TYPE,
+                                         language="en", vad_method="silero")
+    align_model, align_metadata = whisperx.load_align_model(language_code="en", device=device)
+
+    # Get internal model for feature extraction (no need for separate HuggingFace model!)
+    whisper_encoder = whisperx_model.model  # FasterWhisper model with encode() method
+
+    import snac
+    snac_model = snac.SNAC.from_pretrained(SNAC_MODEL).to(device)
+    snac_model.eval()
+
+    # Compile SNAC for faster inference (requires PyTorch 2.0+)
+    try:
+        snac_model = torch.compile(snac_model, mode="reduce-overhead")
+        print(f"[GPU {gpu_id}] SNAC compiled with torch.compile")
+    except Exception as e:
+        print(f"[GPU {gpu_id}] torch.compile not available: {e}")
+
+    print(f"[GPU {gpu_id}] Models loaded. VRAM: {torch.cuda.memory_allocated() / 1e9:.1f}GB")
+
+    # Queues
+    api_to_tts = queue.Queue(maxsize=3)
+    tts_to_features = queue.Queue(maxsize=3)
+    features_to_saver = queue.Queue(maxsize=5)
+
+    stats = PipelineStats(target_count=count, gpu_id=gpu_id)
+
+    # Create output directory for this GPU
+    gpu_output_dir = output_dir / f'gpu{gpu_id}'
+    gpu_output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Start workers
+    api_thread = threading.Thread(target=api_worker, args=(api_to_tts, count, stats))
+    tts_thread = threading.Thread(target=tts_worker, args=(api_to_tts, tts_to_features, stats))
+    feature_thread = threading.Thread(
+        target=feature_worker,
+        args=(tts_to_features, features_to_saver, whisperx_model,
+              align_model, align_metadata, whisper_encoder,
+              snac_model, device, gpu_output_dir, stats)
+    )
+    saver_thread = threading.Thread(
+        target=saver_worker,
+        args=(features_to_saver, gpu_output_dir, stats, start_idx)
+    )
+
+    start_time = time.time()
+
+    api_thread.start()
+    tts_thread.start()
+    feature_thread.start()
+    saver_thread.start()
+
+    api_thread.join()
+    tts_thread.join()
+    tts_to_features.put(None)
+    feature_thread.join()
+    features_to_saver.put(None)
+    saver_thread.join()
+
+    total_time = time.time() - start_time
+    print(f"\n[GPU {gpu_id}] Complete: {stats.total_items} items in {total_time:.1f}s ({stats.total_items/total_time:.2f}/s)")
+
+    return stats.total_items
+
+
+def gpu_worker_process(gpu_id: int, count: int, output_dir: str, start_idx: int, result_queue: mp.Queue):
+    """Wrapper for multiprocessing."""
+    try:
+        output_path = Path(output_dir)
+        items_generated = run_gpu_pipeline(gpu_id, count, output_path, start_idx)
+        result_queue.put((gpu_id, items_generated, None))
+    except Exception as e:
+        result_queue.put((gpu_id, 0, str(e)))
+
+
+def merge_datasets(output_dir: Path, num_gpus: int) -> int:
+    """Merge partial datasets from all GPUs into one."""
+    print("\n" + "="*60)
+    print("MERGING DATASETS")
+    print("="*60)
+
+    all_items = []
+
+    for gpu_id in range(num_gpus):
+        partial_path = output_dir / f'gpu{gpu_id}' / f'dataset_gpu{gpu_id}.pt'
+        if partial_path.exists():
+            items = torch.load(partial_path, weights_only=False)
+            print(f"  GPU {gpu_id}: {len(items)} items")
+            all_items.extend(items)
+
+    # Save merged dataset
+    merged_path = output_dir / 'dataset.pt'
+    torch.save(all_items, merged_path)
+    print(f"\nMerged: {len(all_items)} total items -> {merged_path}")
+
+    # Clean up partial files
+    for gpu_id in range(num_gpus):
+        partial_path = output_dir / f'gpu{gpu_id}' / f'dataset_gpu{gpu_id}.pt'
+        if partial_path.exists():
+            partial_path.unlink()
+        # Remove checkpoints
+        gpu_dir = output_dir / f'gpu{gpu_id}'
+        for ckpt in gpu_dir.glob('checkpoint_*.pt'):
+            ckpt.unlink()
+
+    return len(all_items)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Multi-GPU Dataset Generator")
+    parser.add_argument('--count', type=int, default=1000, help='Total items to generate')
+    parser.add_argument('--batch-size', type=int, default=50, help='TTS batch size')
+    parser.add_argument('--output-dir', type=str, default='./qa_dataset', help='Output directory')
+    parser.add_argument('--tts-workers', type=int, default=4, help='TTS workers per GPU')
+    parser.add_argument('--io-workers', type=int, default=4, help='I/O workers per GPU')
+    parser.add_argument('--gpus', type=str, default=None, help='Comma-separated GPU IDs (default: all)')
+    parser.add_argument('--use-llama', action='store_true', help='Use local Llama 3.2 instead of Groq API')
+    parser.add_argument('--llama-gpu', type=int, default=1, help='GPU ID for Llama model (default: 1)')
+    args = parser.parse_args()
+
+    global BATCH_SIZE, NUM_TTS_WORKERS, IO_WORKERS, USE_LOCAL_LLAMA, LLAMA_GPU_ID
+    BATCH_SIZE = args.batch_size
+    NUM_TTS_WORKERS = args.tts_workers
+    IO_WORKERS = args.io_workers
+    USE_LOCAL_LLAMA = args.use_llama
+    LLAMA_GPU_ID = args.llama_gpu
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Detect GPUs
+    if args.gpus:
+        gpu_ids = [int(g.strip()) for g in args.gpus.split(',')]
+    else:
+        num_gpus = torch.cuda.device_count()
+        gpu_ids = list(range(num_gpus))
+
+    num_gpus = len(gpu_ids)
+
+    if num_gpus == 0:
+        print("ERROR: No GPUs detected!")
+        print("This script requires CUDA GPUs to run.")
+        print("Please run on a machine with NVIDIA GPUs.")
+        sys.exit(1)
+
+    print("="*60)
+    print("MULTI-GPU DATASET GENERATOR")
+    print("="*60)
+    print(f"Total items: {args.count}")
+    print(f"GPUs detected: {num_gpus}")
+    for gpu_id in gpu_ids:
+        torch.cuda.set_device(gpu_id)
+        print(f"  GPU {gpu_id}: {torch.cuda.get_device_name(gpu_id)}")
+    print(f"Items per GPU: ~{args.count // num_gpus}")
+    print(f"Q&A source: {'Local Llama 3.2 (GPU ' + str(LLAMA_GPU_ID) + ')' if USE_LOCAL_LLAMA else 'Groq API'}")
+    print(f"Output: {output_dir}")
+    print("="*60)
+
+    if num_gpus == 1:
+        # Single GPU - run directly
+        print("\nSingle GPU mode")
+        start_time = time.time()
+        total_items = run_gpu_pipeline(gpu_ids[0], args.count, output_dir)
+
+        # Rename partial to final
+        partial = output_dir / f'gpu{gpu_ids[0]}' / f'dataset_gpu{gpu_ids[0]}.pt'
+        if partial.exists():
+            items = torch.load(partial, weights_only=False)
+            torch.save(items, output_dir / 'dataset.pt')
+            partial.unlink()
+
+        total_time = time.time() - start_time
+        print(f"\n{'='*60}")
+        print(f"COMPLETE: {total_items} items in {total_time:.1f}s")
+        print(f"Rate: {total_items/total_time:.2f} items/s")
+        print(f"Output: {output_dir / 'dataset.pt'}")
+        print("="*60)
+    else:
+        # Multi-GPU - spawn processes
+        print(f"\nMulti-GPU mode: {num_gpus} GPUs")
+
+        # Split work across GPUs
+        items_per_gpu = args.count // num_gpus
+        remainder = args.count % num_gpus
+
+        result_queue = mp.Queue()
+        processes = []
+
+        start_time = time.time()
+        start_idx = 0
+
+        for i, gpu_id in enumerate(gpu_ids):
+            # Last GPU gets remainder
+            gpu_count = items_per_gpu + (remainder if i == len(gpu_ids) - 1 else 0)
+
+            p = mp.Process(
+                target=gpu_worker_process,
+                args=(gpu_id, gpu_count, str(output_dir), start_idx, result_queue)
+            )
+            p.start()
+            processes.append(p)
+            start_idx += gpu_count
+
+        # Wait for all processes
+        results = []
+        for _ in processes:
+            result = result_queue.get()
+            results.append(result)
+
+        for p in processes:
+            p.join()
+
+        # Check results
+        total_items = 0
+        for gpu_id, items, error in results:
+            if error:
+                print(f"[GPU {gpu_id}] ERROR: {error}")
+            else:
+                total_items += items
+
+        # Merge datasets
+        total_items = merge_datasets(output_dir, num_gpus)
+
+        total_time = time.time() - start_time
+        print(f"\n{'='*60}")
+        print(f"COMPLETE: {total_items} items in {total_time:.1f}s")
+        print(f"Rate: {total_items/total_time:.2f} items/s")
+        print(f"Output: {output_dir / 'dataset.pt'}")
+        print("="*60)
+
+
+if __name__ == '__main__':
+    mp.set_start_method('spawn', force=True)
+    main()

passo0_setup.py

@@ -0,0 +1,434 @@
+#!/usr/bin/env python3
+"""
+Passo 0: Setup do Ambiente
+Instala todas as dependências e configura o ambiente para:
+- Geração de dataset (Kokoro TTS GPU, WhisperX, SNAC)
+- Treinamento do modelo Speech-to-Speech
+
+Usage:
+    python passo0_setup.py [--dataset_only] [--training_only]
+"""
+
+import os
+import sys
+import subprocess
+import argparse
+import shutil
+
+def log(msg):
+    print(f"[SETUP] {msg}")
+    sys.stdout.flush()
+
+def run_command(cmd, check=True, capture=True):
+    """Execute command and return output."""
+    log(f"Running: {cmd}")
+    if capture:
+        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
+        if check and result.returncode != 0:
+            log(f"ERROR: {result.stderr}")
+            return False
+        return True
+    else:
+        return subprocess.run(cmd, shell=True).returncode == 0
+
+def check_python_version():
+    """Check Python version."""
+    version = sys.version_info
+    log(f"Python version: {version.major}.{version.minor}.{version.micro}")
+
+    if version.major < 3 or (version.major == 3 and version.minor < 10):
+        log("ERROR: Python 3.10+ required")
+        sys.exit(1)
+
+    if version.minor > 12:
+        log("WARNING: Python 3.13+ may have compatibility issues")
+
+    log("Python version OK")
+    return True
+
+def check_nvidia_smi():
+    """Check if nvidia-smi is available."""
+    if shutil.which('nvidia-smi') is None:
+        log("ERROR: nvidia-smi not found. NVIDIA GPU required.")
+        return False
+
+    result = subprocess.run(['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader'],
+                          capture_output=True, text=True)
+    if result.returncode == 0:
+        log(f"GPU: {result.stdout.strip()}")
+        return True
+    return False
+
+def check_cuda():
+    """Check CUDA availability via PyTorch."""
+    try:
+        import torch
+        cuda_available = torch.cuda.is_available()
+        if cuda_available:
+            gpu_count = torch.cuda.device_count()
+            gpu_name = torch.cuda.get_device_name(0)
+            vram = torch.cuda.get_device_properties(0).total_memory / 1024**3
+            log(f"CUDA available: {gpu_count}x {gpu_name} ({vram:.1f}GB)")
+        else:
+            log("WARNING: CUDA not available - will be very slow!")
+        return cuda_available
+    except Exception as e:
+        log(f"Could not check CUDA: {e}")
+        return False
+
+def upgrade_pip():
+    """Upgrade pip to latest version."""
+    log("\nUpgrading pip...")
+    run_command("pip install --upgrade pip -q")
+    log("pip upgraded")
+
+def install_pytorch():
+    """Install PyTorch with CUDA 12.8 (latest, supports H100/H200)."""
+    log("\nInstalling PyTorch with CUDA 12.8...")
+
+    # Always reinstall to ensure correct CUDA version
+    # PyTorch 2.8+ with CUDA 12.8 for best compatibility with modern GPUs
+    cmd = "pip install --force-reinstall torch==2.8.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128"
+    if run_command(cmd, capture=False):
+        log("PyTorch 2.8.0 installed with CUDA 12.8")
+        # Verify
+        try:
+            import importlib
+            import torch
+            importlib.reload(torch)
+            if torch.cuda.is_available():
+                log(f"Verified: PyTorch {torch.__version__} with CUDA {torch.version.cuda}")
+                return True
+        except:
+            pass
+        return True
+    else:
+        log("Failed to install PyTorch")
+        return False
+
+def install_onnxruntime_gpu():
+    """Install onnxruntime-gpu for Kokoro TTS with CUDA support."""
+    log("\nInstalling onnxruntime-gpu with CUDA support...")
+
+    # Remove any existing onnxruntime
+    run_command("pip uninstall onnxruntime onnxruntime-gpu -y -q", check=False)
+
+    # Install onnxruntime-gpu 1.20.1 - known to work with CUDA 12.x
+    # Newer versions (1.23+) may not include CUDA provider
+    cmd = "pip install onnxruntime-gpu==1.20.1"
+    if run_command(cmd, capture=False):
+        # Verify - need to restart Python to pick up new module
+        import subprocess
+        result = subprocess.run(
+            ['python3', '-c', 'import onnxruntime as ort; print(ort.get_available_providers())'],
+            capture_output=True, text=True
+        )
+        providers = result.stdout.strip()
+        log(f"onnxruntime providers: {providers}")
+        if 'CUDAExecutionProvider' in providers:
+            log("onnxruntime-gpu OK with CUDA")
+            return True
+        else:
+            log("WARNING: CUDAExecutionProvider not available")
+            return False
+    return False
+
+def install_dataset_packages():
+    """Install packages for dataset generation."""
+    log("\nInstalling dataset generation packages...")
+
+    packages = [
+        "numpy",                 # Required by many packages
+        "kokoro-onnx",           # TTS (GPU via ONNX CUDA)
+        "whisperx",              # Speech alignment (GPU)
+        "ctranslate2",           # Required by WhisperX for encoder features
+        "snac",                  # Audio codec (GPU)
+        "transformers>=4.40.0",  # Whisper encoder
+        "soundfile",             # Audio I/O
+        "groq",                  # Q&A generation API
+        "requests",              # HTTP requests
+        "tqdm",                  # Progress bars
+    ]
+
+    for pkg in packages:
+        log(f"  Installing {pkg}...")
+        run_command(f"pip install {pkg} -q", check=False)
+
+    log("Dataset packages installed")
+
+def install_training_packages():
+    """Install packages for model training."""
+    log("\nInstalling training packages...")
+
+    packages = [
+        "transformers>=4.40.0",
+        "peft>=0.6.0",
+        "accelerate>=0.24.0",
+        "pandas>=1.5.0",
+        "huggingface-hub>=0.17.0",
+        "psutil>=5.9.0",
+        "tensorboard",
+    ]
+
+    for pkg in packages:
+        log(f"  Installing {pkg}...")
+        run_command(f"pip install {pkg} -q", check=False)
+
+    log("Training packages installed")
+
+def download_kokoro_models():
+    """Download Kokoro TTS model files."""
+    log("\nDownloading Kokoro TTS models...")
+
+    kokoro_dir = "/tmp"
+    model_url = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/kokoro-v1.0.onnx"
+    voices_url = "https://github.com/thewh1teagle/kokoro-onnx/releases/download/model-files-v1.0/voices-v1.0.bin"
+
+    model_path = os.path.join(kokoro_dir, "kokoro-v1.0.onnx")
+    voices_path = os.path.join(kokoro_dir, "voices-v1.0.bin")
+
+    if not os.path.exists(model_path):
+        log("  Downloading kokoro-v1.0.onnx (310MB)...")
+        run_command(f"wget -q --show-progress -O {model_path} {model_url}", capture=False)
+    else:
+        log("  kokoro-v1.0.onnx already exists")
+
+    if not os.path.exists(voices_path):
+        log("  Downloading voices-v1.0.bin (27MB)...")
+        run_command(f"wget -q --show-progress -O {voices_path} {voices_url}", capture=False)
+    else:
+        log("  voices-v1.0.bin already exists")
+
+    log("Kokoro models ready")
+
+def setup_ld_library_path():
+    """Create script to set LD_LIBRARY_PATH for ONNX CUDA."""
+    log("\nCreating environment setup script...")
+
+    # Critical libraries for onnxruntime-gpu CUDA provider
+    cuda_libs = [
+        'cublas', 'cudnn', 'cufft', 'curand', 'cusolver', 'cusparse',
+        'cuda_runtime', 'cuda_nvrtc', 'cuda_cupti', 'nvjitlink', 'nccl',
+        'nvtx', 'cusparselt', 'cufile'
+    ]
+
+    script_content = f'''#!/bin/bash
+# Source this before running dataset generation
+# Usage: source /tmp/setup_env.sh
+
+# Set LD_LIBRARY_PATH for ONNX CUDA - CRITICAL for GPU inference
+NVIDIA_LIBS=""
+
+# Check all Python versions
+for pyver in python3.12 python3.11 python3.10; do
+    NVIDIA_BASE="/usr/local/lib/$pyver/dist-packages/nvidia"
+    for subdir in {' '.join(cuda_libs)}; do
+        lib_path="$NVIDIA_BASE/$subdir/lib"
+        if [ -d "$lib_path" ]; then
+            NVIDIA_LIBS="$NVIDIA_LIBS:$lib_path"
+        fi
+    done
+done
+
+export LD_LIBRARY_PATH="${{NVIDIA_LIBS:1}}:$LD_LIBRARY_PATH"
+
+# CRITICAL: Force Kokoro to use CUDA provider
+export ONNX_PROVIDER="CUDAExecutionProvider"
+
+echo "Environment configured for GPU inference"
+echo "LD_LIBRARY_PATH set with NVIDIA CUDA libraries"
+echo "ONNX_PROVIDER=$ONNX_PROVIDER"
+'''
+
+    with open("/tmp/setup_env.sh", "w") as f:
+        f.write(script_content)
+    os.chmod("/tmp/setup_env.sh", 0o755)
+
+    log("Created /tmp/setup_env.sh")
+
+def create_directories():
+    """Create necessary directories."""
+    log("\nCreating directories...")
+
+    dirs = [
+        "./data",
+        "./data/raw",
+        "./data/processed",
+        "./checkpoints",
+        "./logs",
+        "./datasets",
+    ]
+
+    for d in dirs:
+        os.makedirs(d, exist_ok=True)
+
+    log("Directories created")
+
+def login_huggingface():
+    """Login to HuggingFace."""
+    log("\nHuggingFace login...")
+
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        log("WARNING: HF_TOKEN not set. Set it with: export HF_TOKEN=your_token")
+        return False
+
+    try:
+        from huggingface_hub import login
+        login(token=token)
+        log("HuggingFace login successful")
+        return True
+    except Exception as e:
+        log(f"HuggingFace login failed: {e}")
+        return False
+
+def test_kokoro_gpu():
+    """Test Kokoro TTS on GPU."""
+    log("\nTesting Kokoro TTS on GPU...")
+    log("NOTE: LD_LIBRARY_PATH must be set BEFORE Python starts for CUDA to work.")
+    log("Run this test via: source /tmp/setup_env.sh && python3 -c 'from passo0_setup import test_kokoro_gpu; test_kokoro_gpu()'")
+
+    # Check if ONNX_PROVIDER is set
+    if os.environ.get('ONNX_PROVIDER') != 'CUDAExecutionProvider':
+        log("Setting ONNX_PROVIDER=CUDAExecutionProvider")
+        os.environ['ONNX_PROVIDER'] = 'CUDAExecutionProvider'
+
+    try:
+        import time
+        from kokoro_onnx import Kokoro
+
+        k = Kokoro('/tmp/kokoro-v1.0.onnx', '/tmp/voices-v1.0.bin')
+        providers = k.sess.get_providers()
+        log(f"  Providers: {providers}")
+
+        if 'CUDAExecutionProvider' in providers:
+            # Warmup
+            k.create("warmup", voice='af_heart', speed=1.0)
+
+            # Benchmark
+            start = time.time()
+            for _ in range(5):
+                samples, sr = k.create("Hello, testing GPU acceleration.", voice='af_heart', speed=1.0)
+            elapsed = time.time() - start
+            calls_per_sec = 5 / elapsed
+
+            log(f"  Speed: {calls_per_sec:.1f} calls/s")
+            if calls_per_sec > 5:
+                log("Kokoro GPU test PASSED")
+                return True
+            else:
+                log("WARNING: GPU detected but speed is slow")
+                return True
+        else:
+            log("WARNING: Kokoro running on CPU, not GPU")
+            log("Make sure to run: source /tmp/setup_env.sh before running Python")
+            return False
+    except Exception as e:
+        log(f"Kokoro test failed: {e}")
+        return False
+
+def verify_installation():
+    """Verify all packages are installed."""
+    log("\nVerifying installation...")
+
+    packages = {
+        "torch": "PyTorch",
+        "torchaudio": "TorchAudio",
+        "onnxruntime": "ONNX Runtime",
+        "kokoro_onnx": "Kokoro TTS",
+        "whisperx": "WhisperX",
+        "ctranslate2": "CTranslate2",
+        "snac": "SNAC",
+        "transformers": "Transformers",
+        "peft": "PEFT",
+        "soundfile": "SoundFile",
+        "groq": "Groq API",
+    }
+
+    all_ok = True
+    for pkg, name in packages.items():
+        try:
+            __import__(pkg)
+            log(f"  {name}: OK")
+        except ImportError:
+            log(f"  {name}: MISSING")
+            all_ok = False
+
+    return all_ok
+
+def main():
+    parser = argparse.ArgumentParser(description="Setup environment for Speech-to-Speech")
+    parser.add_argument("--dataset_only", action="store_true", help="Only install dataset generation packages")
+    parser.add_argument("--training_only", action="store_true", help="Only install training packages")
+    parser.add_argument("--skip_kokoro_download", action="store_true", help="Skip Kokoro model download")
+    parser.add_argument("--skip_test", action="store_true", help="Skip Kokoro GPU test")
+
+    args = parser.parse_args()
+
+    log("="*60)
+    log("PASSO 0: SETUP DO AMBIENTE")
+    log("="*60)
+
+    # Check prerequisites
+    check_python_version()
+
+    if not check_nvidia_smi():
+        log("ERROR: NVIDIA GPU required")
+        sys.exit(1)
+
+    # Upgrade pip
+    upgrade_pip()
+
+    # Install PyTorch
+    install_pytorch()
+    check_cuda()
+
+    # Install packages based on mode
+    if args.training_only:
+        install_training_packages()
+    elif args.dataset_only:
+        install_onnxruntime_gpu()
+        install_dataset_packages()
+    else:
+        # Install everything
+        install_onnxruntime_gpu()
+        install_dataset_packages()
+        install_training_packages()
+
+    # Download Kokoro models
+    if not args.training_only and not args.skip_kokoro_download:
+        download_kokoro_models()
+
+    # Setup environment
+    setup_ld_library_path()
+
+    # Create directories
+    create_directories()
+
+    # HuggingFace login
+    login_huggingface()
+
+    # Verify
+    if not verify_installation():
+        log("\nWARNING: Some packages missing")
+
+    # Test Kokoro GPU
+    if not args.training_only and not args.skip_test:
+        test_kokoro_gpu()
+
+    log("\n" + "="*60)
+    log("SETUP COMPLETO!")
+    log("="*60)
+
+    log("\nPara gerar dataset:")
+    log("  1. source /tmp/setup_env.sh")
+    log("  2. python datasets/generate_dataset.py --count 1000 --output-dir ./data/processed")
+    log("")
+    log("Para treinar:")
+    log("  1. python passo2_finetune_stage1.py --data ./data/processed/dataset.pt")
+    log("  2. python passo3_finetune_stage2.py --data ./data/processed/dataset.pt --stage1_ckpt ./checkpoints/stage1_best.pt")
+
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())

passo2_finetune_stage1.py

@@ -0,0 +1,817 @@
+#!/usr/bin/env python3
+"""
+Stage 1: Train Adapter Only (LLM Frozen) with Scheduled Interleaved Output
+
+Based on IST-LM paper combined with LLaMA-Omni 2 staging:
+- Adapter learns to map audio → LLM embedding space
+- Output is interleaved text+audio (90% text initially)
+- LLM is completely frozen (adapter gets a "head start")
+
+Usage:
+    python finetune_stage1.py --data data.pt --epochs 2
+
+Next: Stage 2 trains Adapter + LoRA together
+"""
+
+import os
+import sys
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader, DistributedSampler, ConcatDataset
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import login
+from tqdm import tqdm
+import threading
+
+# ============================================================
+# Config
+# ============================================================
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
+torch.set_float32_matmul_precision('high')
+
+# SNAC token offsets for Orpheus
+SNAC_BASE_OFFSET = 128266
+EOS_TOKEN = 128009
+
+
+def log(msg):
+    print(msg)
+    sys.stdout.flush()
+
+
+def apply_snac_offset(token_idx, position):
+    """Apply position-based offset to SNAC token.
+    If token is already offset (>= SNAC_BASE_OFFSET), return as-is.
+    """
+    if int(token_idx) >= SNAC_BASE_OFFSET:
+        # Already has offset applied
+        return int(token_idx)
+    offset = SNAC_BASE_OFFSET + (position % 7) * 4096
+    return int(token_idx) + offset
+
+
+def get_text_ratio(global_step, decay_steps=300, initial_ratio=0.9, min_ratio=0.0):
+    """
+    IST-LM: Start with 90% text, decrease by 0.1 every 300 steps.
+    """
+    num_decays = global_step // decay_steps
+    text_ratio = initial_ratio - (num_decays * 0.1)
+    return max(min_ratio, text_ratio)
+
+
+# ============================================================
+# Async Checkpoint Saving
+# ============================================================
+_save_threads = []
+
+def save_checkpoint_async(state_dict, path, is_main=True):
+    global _save_threads
+    _save_threads = [t for t in _save_threads if t.is_alive()]
+
+    def copy_to_cpu(obj):
+        if isinstance(obj, torch.Tensor):
+            return obj.detach().cpu().clone()
+        elif isinstance(obj, dict):
+            return {k: copy_to_cpu(v) for k, v in obj.items()}
+        return obj
+
+    state_copy = copy_to_cpu(state_dict)
+
+    def _save():
+        try:
+            torch.save(state_copy, path)
+            if is_main:
+                log(f"[ASYNC] Saved: {path}")
+        except Exception as e:
+            if is_main:
+                log(f"[ASYNC] Error: {e}")
+
+    thread = threading.Thread(target=_save, daemon=True)
+    thread.start()
+    _save_threads.append(thread)
+
+
+def wait_for_checkpoints():
+    global _save_threads
+    for t in _save_threads:
+        t.join()
+    _save_threads = []
+
+
+# ============================================================
+# GPU Auto-Detection
+# ============================================================
+def auto_detect_gpu_config():
+    try:
+        import subprocess
+        result = subprocess.run(
+            ['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader,nounits'],
+            capture_output=True, text=True
+        )
+        lines = result.stdout.strip().split('\n')
+        gpu_name, vram_mb = lines[0].split(', ')
+        vram_gb = int(vram_mb) // 1024
+
+        if vram_gb >= 150:
+            return {"name": "B200", "batch_size": 8, "grad_accum": 4, "vram_gb": vram_gb}
+        elif vram_gb >= 80:
+            return {"name": "A100-80GB", "batch_size": 6, "grad_accum": 5, "vram_gb": vram_gb}
+        elif vram_gb >= 35:
+            return {"name": "A100-40GB", "batch_size": 4, "grad_accum": 8, "vram_gb": vram_gb}
+        else:
+            return {"name": "RTX4090", "batch_size": 2, "grad_accum": 16, "vram_gb": vram_gb}
+    except:
+        return {"name": "Unknown", "batch_size": 2, "grad_accum": 16, "vram_gb": 24}
+
+def get_ram_info():
+    """Get RAM info in GB."""
+    try:
+        import psutil
+        total = psutil.virtual_memory().total / 1024**3
+        available = psutil.virtual_memory().available / 1024**3
+        return total, available
+    except ImportError:
+        try:
+            import subprocess
+            result = subprocess.run(
+                ['free', '-g', '--output=SIZE,AVAILABLE'],
+                capture_output=True, text=True
+            )
+            lines = result.stdout.strip().split('\n')
+            if len(lines) >= 2:
+                total, available = map(float, lines[1].split())
+                return total, available
+        except:
+            pass
+    except:
+        pass
+    return 0, 0
+
+def limit_ram_usage(max_ram_gb):
+    """Limit RAM usage by setting resource limits."""
+    try:
+        import resource
+        max_bytes = int(max_ram_gb * 1024**3)
+        resource.setrlimit(resource.RLIMIT_AS, (max_bytes, max_bytes))
+    except:
+        pass
+
+def log_memory_usage():
+    """Log current memory usage."""
+    msg = []
+    if torch.cuda.is_available():
+        used = torch.cuda.memory_allocated() / 1024**3
+        reserved = torch.cuda.memory_reserved() / 1024**3
+        msg.append(f"GPU: {used:.2f}GB / {reserved:.2f}GB")
+    try:
+        import psutil
+        ram_used = psutil.virtual_memory().used / 1024**3
+        ram_total = psutil.virtual_memory().total / 1024**3
+        msg.append(f"RAM: {ram_used:.1f}GB / {ram_total:.1f}GB")
+    except:
+        pass
+    return " | ".join(msg)
+
+
+# ============================================================
+# Speech Adapter (LLaMA-Omni 2 Style)
+# ============================================================
+class SpeechAdapter(nn.Module):
+    """
+    5× downsampling + FFN with intermediate dim 2048
+    """
+    def __init__(self, whisper_dim=1280, llm_dim=3072, downsample=5, intermediate_dim=2048):
+        super().__init__()
+        self.downsample = downsample
+        concat_dim = whisper_dim * downsample
+
+        self.ffn = nn.Sequential(
+            nn.Linear(concat_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, llm_dim),
+            nn.LayerNorm(llm_dim)
+        )
+
+    def forward(self, x):
+        B, T, D = x.shape
+        T_new = (T // self.downsample) * self.downsample
+        x = x[:, :T_new]
+        x = x.reshape(B, T_new // self.downsample, D * self.downsample)
+        return self.ffn(x)
+
+
+# ============================================================
+# Scheduled Interleaved Sequence Creation with Word Alignment
+# ============================================================
+def create_interleaved_sequence(text_tokens, snac_tokens, text_ratio=0.9, word_alignments=None, tokenizer=None, answer_text=None):
+    """
+    Create interleaved sequence based on text_ratio with word-level alignment.
+    - text_ratio=0.9 means 90% of words are replaced by text tokens
+    - text_ratio=0.0 means 100% audio (no text replacement)
+
+    With word_alignments: replaces aligned audio spans with corresponding text tokens
+    Without word_alignments: falls back to positional interleaving
+    """
+    interleaved = []
+    is_audio_mask = []
+
+    if len(snac_tokens) == 0:
+        return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
+
+    # Group SNAC into frames of 7
+    frames = []
+    for i in range(0, len(snac_tokens), 7):
+        frame = snac_tokens[i:i+7]
+        if len(frame) == 7:
+            frames.append(frame)
+
+    if len(frames) == 0:
+        return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
+
+    total_frames = len(frames)
+
+    # If we have word alignments, use semantic interleaving
+    if word_alignments and tokenizer and answer_text and text_ratio > 0:
+        import random
+
+        # Decide which words to replace with text based on text_ratio
+        num_words = len(word_alignments)
+        num_text_words = int(num_words * text_ratio)
+
+        # Randomly select which word indices to replace with text
+        word_indices = list(range(num_words))
+        random.shuffle(word_indices)
+        text_word_indices = set(word_indices[:num_text_words])
+
+        # Build interleaved sequence frame by frame
+        frame_idx = 0
+        snac_position = 0
+
+        for word_idx, alignment in enumerate(word_alignments):
+            word = alignment['word']
+            start_frame = alignment['start_frame']
+            end_frame = min(alignment['end_frame'], total_frames)
+
+            if word_idx in text_word_indices:
+                # Replace this word's audio with text tokens
+                word_tokens = tokenizer.encode(word, add_special_tokens=False)
+                for tok in word_tokens:
+                    interleaved.append(tok)
+                    is_audio_mask.append(False)
+                # Skip the audio frames for this word
+                snac_position = end_frame * 7
+            else:
+                # Keep audio for this word
+                for f_idx in range(start_frame, end_frame):
+                    if f_idx < total_frames:
+                        frame = frames[f_idx]
+                        for tok in frame:
+                            interleaved.append(apply_snac_offset(tok, snac_position))
+                            is_audio_mask.append(True)
+                            snac_position += 1
+
+            frame_idx = end_frame
+
+        # Add any remaining frames after the last word
+        while frame_idx < total_frames:
+            frame = frames[frame_idx]
+            for tok in frame:
+                interleaved.append(apply_snac_offset(tok, snac_position))
+                is_audio_mask.append(True)
+                snac_position += 1
+            frame_idx += 1
+
+    else:
+        # Fallback: positional interleaving (original behavior)
+        total_text = len(text_tokens)
+
+        # Determine interleaving pattern based on text_ratio
+        if text_ratio >= 0.9:
+            text_per_chunk, frames_per_chunk = 1, 3
+        elif text_ratio >= 0.7:
+            text_per_chunk, frames_per_chunk = 1, 5
+        elif text_ratio >= 0.5:
+            text_per_chunk, frames_per_chunk = 1, 7
+        elif text_ratio >= 0.3:
+            text_per_chunk, frames_per_chunk = 1, 10
+        else:
+            text_per_chunk, frames_per_chunk = 0, 1
+
+        text_idx = 0
+        frame_idx = 0
+        snac_position = 0
+
+        while frame_idx < total_frames:
+            if text_per_chunk > 0 and text_idx < total_text:
+                for _ in range(text_per_chunk):
+                    if text_idx < total_text:
+                        interleaved.append(text_tokens[text_idx])
+                        is_audio_mask.append(False)
+                        text_idx += 1
+
+            for _ in range(frames_per_chunk):
+                if frame_idx < total_frames:
+                    frame = frames[frame_idx]
+                    for tok in frame:
+                        interleaved.append(apply_snac_offset(tok, snac_position))
+                        is_audio_mask.append(True)
+                        snac_position += 1
+                    frame_idx += 1
+
+        while text_idx < total_text:
+            interleaved.append(text_tokens[text_idx])
+            is_audio_mask.append(False)
+            text_idx += 1
+
+    # Add EOS
+    interleaved.append(EOS_TOKEN)
+    is_audio_mask.append(False)
+
+    return interleaved, is_audio_mask
+
+
+# ============================================================
+# Dataset
+# ============================================================
+class InterleavedDataset(Dataset):
+    def __init__(self, data, tokenizer, max_audio_len=500, max_seq_len=2048):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_audio = max_audio_len * 5
+        self.max_seq_len = max_seq_len
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+
+        # Whisper features
+        whisper = item["whisper_features"][:self.max_audio]
+
+        # Text tokens - use pre-computed if available, otherwise tokenize
+        if "text_tokens" in item and len(item["text_tokens"]) > 0:
+            tt = item["text_tokens"]
+            text_tokens = tt.tolist() if hasattr(tt, 'tolist') else list(tt)
+        else:
+            text = item.get("answer", item.get("text", ""))
+            if isinstance(text, str) and len(text) > 0:
+                text_tokens = self.tokenizer.encode(text, add_special_tokens=False)
+            else:
+                text_tokens = []
+
+        # SNAC tokens
+        snac = item["snac_tokens"]
+        snac_len = (len(snac) // 7) * 7
+        snac = snac[:snac_len] if snac_len > 0 else snac[:7]
+        snac_list = snac.tolist() if hasattr(snac, 'tolist') else list(snac)
+
+        # Word alignments (if available)
+        word_alignments = item.get("word_alignments", None)
+        answer_text = item.get("answer", "")
+
+        return {
+            "whisper": whisper,
+            "text_tokens": text_tokens,
+            "snac_tokens": snac_list,
+            "word_alignments": word_alignments,
+            "answer_text": answer_text
+        }
+
+
+def collate_fn(batch, text_ratio=0.9, tokenizer=None):
+    """Collate with dynamic interleaving based on text_ratio and word alignments."""
+    max_w = max(b["whisper"].shape[0] for b in batch)
+    max_w = ((max_w + 4) // 5) * 5
+
+    whisper_batch = []
+    interleaved_batch = []
+    is_audio_batch = []
+
+    max_seq = 0
+    sequences = []
+
+    for b in batch:
+        interleaved, is_audio = create_interleaved_sequence(
+            b["text_tokens"],
+            b["snac_tokens"],
+            text_ratio,
+            word_alignments=b.get("word_alignments"),
+            tokenizer=tokenizer,
+            answer_text=b.get("answer_text")
+        )
+        sequences.append((interleaved, is_audio))
+        max_seq = max(max_seq, len(interleaved))
+
+    for i, b in enumerate(batch):
+        w = b["whisper"]
+        w_pad = F.pad(w, (0, 0, 0, max_w - w.shape[0]))
+        whisper_batch.append(w_pad)
+
+        interleaved, is_audio = sequences[i]
+        seq_tensor = torch.tensor(interleaved, dtype=torch.long)
+        mask_tensor = torch.tensor(is_audio, dtype=torch.bool)
+
+        seq_pad = F.pad(seq_tensor, (0, max_seq - len(interleaved)), value=-100)
+        mask_pad = F.pad(mask_tensor, (0, max_seq - len(is_audio)), value=False)
+
+        interleaved_batch.append(seq_pad)
+        is_audio_batch.append(mask_pad)
+
+    return {
+        "whisper": torch.stack(whisper_batch),
+        "interleaved": torch.stack(interleaved_batch),
+        "is_audio_mask": torch.stack(is_audio_batch)
+    }
+
+
+# ============================================================
+# Arguments
+# ============================================================
+def parse_args():
+    parser = argparse.ArgumentParser(description="Stage 1: Adapter Only with Interleaved Output")
+    parser.add_argument("--data", type=str, required=True)
+    parser.add_argument("--output_dir", type=str, default="./checkpoints")
+    parser.add_argument("--lr", type=float, default=5e-5)
+    parser.add_argument("--epochs", type=int, default=2, help="1-2 epochs for adapter warmup")
+    parser.add_argument("--batch_size", type=int, default=None)
+    parser.add_argument("--grad_accum", type=int, default=None)
+    parser.add_argument("--warmup_ratio", type=float, default=0.03)
+    parser.add_argument("--max_grad_norm", type=float, default=1.0)
+    parser.add_argument("--save_steps", type=int, default=200)
+    parser.add_argument("--label_smoothing", type=float, default=0.1)
+    # Scheduled interleaving
+    parser.add_argument("--initial_text_ratio", type=float, default=0.9)
+    parser.add_argument("--decay_steps", type=int, default=300)
+    # Model
+    parser.add_argument("--model_path", type=str, default="canopylabs/3b-es_it-ft-research_release")
+    parser.add_argument("--resume", type=str, default=None)
+    # Memory limits
+    parser.add_argument("--vram_fraction", type=float, default=0.80, help="VRAM fraction to use (default 0.80)")
+    parser.add_argument("--ram_limit_gb", type=float, default=None, help="RAM limit in GB (auto if not specified)")
+    # Modes
+    parser.add_argument("--demo", action="store_true")
+    parser.add_argument("--test", action="store_true")
+    return parser.parse_args()
+
+
+# ============================================================
+# Main
+# ============================================================
+def main():
+    args = parse_args()
+
+    # DDP setup
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+
+    if world_size > 1:
+        dist.init_process_group("nccl")
+        torch.cuda.set_device(local_rank)
+
+    device = torch.device(f"cuda:{local_rank}")
+    is_main = local_rank == 0
+
+    # GPU config
+    gpu_config = auto_detect_gpu_config()
+    if args.batch_size is None:
+        args.batch_size = gpu_config["batch_size"]
+    if args.grad_accum is None:
+        args.grad_accum = gpu_config["grad_accum"]
+
+    torch_dtype = torch.bfloat16
+
+    # Get RAM info and set limits
+    ram_total, ram_available = get_ram_info()
+    if args.ram_limit_gb is None:
+        args.ram_limit_gb = ram_total * 0.80  # Default to 80% of RAM
+    limit_ram_usage(args.ram_limit_gb)
+
+    if is_main:
+        log("=" * 60)
+        log("STAGE 1: Adapter Only (LLM Frozen) + Interleaved Output")
+        log("=" * 60)
+        log(f"GPU: {gpu_config['name']} ({gpu_config['vram_gb']}GB)")
+        log(f"RAM: {ram_total:.1f}GB total, {ram_available:.1f}GB available")
+        log(f"Batch: {args.batch_size}, Grad accum: {args.grad_accum}")
+        log(f"LR: {args.lr}, Epochs: {args.epochs}")
+        log(f"VRAM limit: {args.vram_fraction*100:.0f}%")
+        log(f"RAM limit: {args.ram_limit_gb:.1f}GB (80% of total)")
+        log(f"Initial text ratio: {args.initial_text_ratio}")
+        log(f"Decay steps: {args.decay_steps}")
+
+    # Apply VRAM and RAM limits
+    if str(device).startswith('cuda'):
+        torch.cuda.set_per_process_memory_fraction(args.vram_fraction)
+        torch.cuda.empty_cache()
+        torch.backends.cudnn.benchmark = True
+        torch.set_float32_matmul_precision('high')
+        if is_main:
+            log(f"[MEMORY] VRAM limited to {args.vram_fraction*100:.0f}%")
+            log(f"[MEMORY] RAM limited to {args.ram_limit_gb:.1f}GB")
+
+    # HuggingFace login
+    hf_token = os.environ.get("HF_TOKEN")
+    if hf_token:
+        login(token=hf_token)
+
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Load datasets
+    data_paths = [p.strip() for p in args.data.split(",")]
+    all_datasets = []
+
+    if is_main:
+        log("\nLoading datasets...")
+
+    for path in data_paths:
+        if os.path.exists(path):
+            data = torch.load(path, weights_only=False, mmap=True)
+            dataset = InterleavedDataset(data, tokenizer)
+            all_datasets.append(dataset)
+            if is_main:
+                log(f"  {os.path.basename(path)}: {len(data):,} samples")
+
+    if len(all_datasets) == 0:
+        raise ValueError("No datasets loaded!")
+
+    combined_dataset = ConcatDataset(all_datasets) if len(all_datasets) > 1 else all_datasets[0]
+
+    # Demo/Test mode
+    if args.test:
+        combined_dataset = torch.utils.data.Subset(combined_dataset, range(min(5, len(combined_dataset))))
+        args.batch_size = min(args.batch_size, len(combined_dataset))
+        args.grad_accum = 1
+    elif args.demo:
+        combined_dataset = torch.utils.data.Subset(combined_dataset, range(min(1000, len(combined_dataset))))
+        args.batch_size = min(4, args.batch_size)
+        args.grad_accum = max(8, args.grad_accum)
+
+    if is_main:
+        log(f"Total samples: {len(combined_dataset):,}")
+
+    # Load LLM (FROZEN)
+    if is_main:
+        log(f"\nLoading LLM (FROZEN): {args.model_path}")
+
+    llm = AutoModelForCausalLM.from_pretrained(
+        args.model_path,
+        torch_dtype=torch_dtype,
+        attn_implementation="sdpa",
+        device_map={"": device}
+    )
+
+    # Freeze LLM completely
+    for p in llm.parameters():
+        p.requires_grad = False
+    llm.eval()
+
+    # Create adapter (TRAINABLE)
+    adapter = SpeechAdapter(
+        whisper_dim=1280,
+        llm_dim=3072,
+        downsample=5,
+        intermediate_dim=2048
+    ).to(device, dtype=torch_dtype)
+
+    adapter_params = sum(p.numel() for p in adapter.parameters())
+    if is_main:
+        log(f"\nTrainable: Adapter only ({adapter_params:,} = {adapter_params/1e6:.1f}M params)")
+        log("LLM: FROZEN")
+
+    # DDP
+    if world_size > 1:
+        adapter = DDP(adapter, device_ids=[local_rank])
+        dist.barrier()
+
+    # Optimizer (only adapter)
+    optimizer = torch.optim.AdamW(adapter.parameters(), lr=args.lr, weight_decay=0.01)
+
+    # We need to create dataloader with dynamic collate
+    # For simplicity, we'll update text_ratio per epoch
+    global_step = 0
+    start_epoch = 0
+    best_loss = float("inf")
+    current_text_ratio = args.initial_text_ratio
+
+    # Resume
+    if args.resume and os.path.exists(args.resume):
+        if is_main:
+            log(f"\nResuming from: {args.resume}")
+        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
+        adapter_module = adapter.module if hasattr(adapter, "module") else adapter
+        if "adapter" in ckpt:
+            adapter_module.load_state_dict(ckpt["adapter"])
+        if "optimizer" in ckpt:
+            optimizer.load_state_dict(ckpt["optimizer"])
+        if "step" in ckpt:
+            global_step = ckpt["step"]
+        if "epoch" in ckpt:
+            start_epoch = ckpt["epoch"]
+        if "text_ratio" in ckpt:
+            current_text_ratio = ckpt["text_ratio"]
+
+    # Training
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Create dataloader first to calculate steps
+    def collate_with_ratio(batch):
+        current_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
+        return collate_fn(batch, current_ratio, tokenizer=tokenizer)
+
+    if world_size > 1:
+        sampler = DistributedSampler(combined_dataset, shuffle=True)
+        sampler.set_epoch(0)
+    else:
+        sampler = None
+
+    train_loader = DataLoader(
+        combined_dataset,
+        batch_size=args.batch_size,
+        shuffle=(sampler is None),
+        sampler=sampler,
+        collate_fn=collate_with_ratio,
+        num_workers=4,
+        pin_memory=True
+    )
+
+    steps_per_epoch = max(1, len(train_loader) // args.grad_accum)
+    total_steps = max(1, steps_per_epoch * args.epochs)
+    warmup_steps = int(total_steps * args.warmup_ratio)
+
+    scheduler = CosineAnnealingLR(optimizer, T_max=max(1, total_steps - warmup_steps), eta_min=1e-6)
+
+    if is_main:
+        log(f"Steps per epoch: {steps_per_epoch}, Total: {total_steps}, Warmup: {warmup_steps}")
+
+    if is_main:
+        log("\n" + "=" * 60)
+        log("STARTING STAGE 1 TRAINING")
+        log("=" * 60)
+
+    for epoch in range(start_epoch, args.epochs):
+        # Update text ratio based on global step
+        current_text_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
+
+        # Update sampler for this epoch
+        if world_size > 1:
+            sampler.set_epoch(epoch)
+
+        # Recreate dataloader with current text_ratio
+        def collate_with_ratio(batch):
+            return collate_fn(batch, current_text_ratio, tokenizer=tokenizer)
+
+        train_loader = DataLoader(
+            combined_dataset,
+            batch_size=args.batch_size,
+            shuffle=(sampler is None),
+            sampler=sampler,
+            collate_fn=collate_with_ratio,
+            num_workers=4,
+            pin_memory=True
+        )
+
+        adapter.train()
+        epoch_loss = 0
+        accum_loss = 0
+
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}", disable=not is_main)
+
+        for batch_idx, batch in enumerate(pbar):
+            # Update text ratio dynamically
+            current_text_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
+
+            whisper = batch["whisper"].to(device, dtype=torch_dtype, non_blocking=True)
+            interleaved = batch["interleaved"].to(device, non_blocking=True)
+
+            # Forward through adapter
+            adapter_module = adapter.module if hasattr(adapter, "module") else adapter
+            audio_embeds = adapter_module(whisper)
+
+            # Get token embeddings for interleaved sequence (teacher forcing)
+            input_tokens = interleaved[:, :-1].clamp(min=0)
+            with torch.no_grad():
+                token_embeds = llm.model.embed_tokens(input_tokens)
+
+            # Combine: [audio_embeds] + [token_embeds]
+            combined = torch.cat([audio_embeds, token_embeds], dim=1)
+
+            # Forward through frozen LLM (no torch.no_grad - need gradients to flow back to adapter)
+            # LLM params are frozen (requires_grad=False), so they won't be updated
+            outputs = llm(inputs_embeds=combined, use_cache=False)
+            logits = outputs.logits
+
+            # Loss: predict interleaved tokens after audio prefix
+            # logits[:, audio_len-1] predicts first token of interleaved
+            # logits[:, audio_len-1+i] predicts interleaved[i]
+            audio_len = audio_embeds.shape[1]
+            seq_len = interleaved.shape[1]
+
+            # Get logits that predict the interleaved sequence
+            # Position audio_len-1 predicts interleaved[0], etc.
+            seq_logits = logits[:, audio_len-1:audio_len-1+seq_len]
+            targets = interleaved
+
+            loss = F.cross_entropy(
+                seq_logits.reshape(-1, logits.size(-1)),
+                targets.reshape(-1),
+                ignore_index=-100,
+                label_smoothing=args.label_smoothing
+            )
+
+            loss = loss / args.grad_accum
+            loss.backward()
+            accum_loss += loss.item() * args.grad_accum
+
+            # Update
+            if (batch_idx + 1) % args.grad_accum == 0 or (batch_idx + 1) == len(train_loader):
+                torch.nn.utils.clip_grad_norm_(adapter.parameters(), args.max_grad_norm)
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+                if global_step < warmup_steps:
+                    lr_scale = (global_step + 1) / warmup_steps
+                    for pg in optimizer.param_groups:
+                        pg["lr"] = args.lr * lr_scale
+                else:
+                    scheduler.step()
+
+                global_step += 1
+                epoch_loss += accum_loss
+
+                pbar.set_postfix(
+                    loss=f"{accum_loss:.4f}",
+                    text_ratio=f"{current_text_ratio:.1f}",
+                    lr=f"{optimizer.param_groups[0]['lr']:.2e}"
+                )
+
+                # Save checkpoint
+                if global_step % args.save_steps == 0 and is_main:
+                    adapter_to_save = adapter.module if hasattr(adapter, "module") else adapter
+                    ckpt_path = os.path.join(args.output_dir, f"stage1_step{global_step}.pt")
+                    save_checkpoint_async({
+                        "adapter": adapter_to_save.state_dict(),
+                        "optimizer": optimizer.state_dict(),
+                        "step": global_step,
+                        "epoch": epoch,
+                        "loss": accum_loss,
+                        "text_ratio": current_text_ratio
+                    }, ckpt_path, is_main)
+
+                accum_loss = 0
+
+        # Epoch end
+        avg_loss = epoch_loss / max(1, steps_per_epoch)
+
+        if is_main:
+            log(f"Epoch {epoch+1} avg loss: {avg_loss:.4f}, text_ratio: {current_text_ratio:.1f}")
+
+            adapter_to_save = adapter.module if hasattr(adapter, "module") else adapter
+            ckpt_path = os.path.join(args.output_dir, f"stage1_epoch{epoch+1}.pt")
+            save_checkpoint_async({
+                "adapter": adapter_to_save.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "step": global_step,
+                "epoch": epoch + 1,
+                "loss": avg_loss,
+                "text_ratio": current_text_ratio
+            }, ckpt_path, is_main)
+
+            if avg_loss < best_loss:
+                best_loss = avg_loss
+                best_path = os.path.join(args.output_dir, "stage1_best.pt")
+                save_checkpoint_async({
+                    "adapter": adapter_to_save.state_dict(),
+                    "step": global_step,
+                    "epoch": epoch + 1,
+                    "loss": best_loss,
+                    "text_ratio": current_text_ratio
+                }, best_path, is_main)
+                log(f"Best model saved! Loss: {best_loss:.4f}")
+
+    # Finish
+    if is_main:
+        wait_for_checkpoints()
+        log("\n" + "=" * 60)
+        log("STAGE 1 COMPLETE!")
+        log(f"Best loss: {best_loss:.4f}")
+        log(f"Final text_ratio: {current_text_ratio:.1f}")
+        log("Next: Stage 2 (Adapter + LoRA)")
+        log("=" * 60)
+
+    if world_size > 1:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

passo3_finetune_stage2.py

@@ -0,0 +1,876 @@
+#!/usr/bin/env python3
+"""
+Stage 2: Train Adapter + LoRA Together with Scheduled Interleaved Output
+
+Continues from Stage 1 checkpoint:
+- Loads pre-trained adapter from Stage 1
+- Adds LoRA to LLM
+- Trains both Adapter + LoRA together
+- Continues scheduled interleaving (90% text -> 0% text)
+
+Usage:
+    python finetune_stage2.py --data data.pt --stage1_ckpt checkpoints/stage1_best.pt --epochs 3
+
+Based on IST-LM paper + LLaMA-Omni 2 staging approach.
+"""
+
+import os
+import sys
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+from torch.utils.data import Dataset, DataLoader, DistributedSampler, ConcatDataset
+from torch.nn.parallel import DistributedDataParallel as DDP
+from torch.optim.lr_scheduler import CosineAnnealingLR
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import LoraConfig, get_peft_model, TaskType
+from huggingface_hub import login
+from tqdm import tqdm
+import threading
+
+# ============================================================
+# Config
+# ============================================================
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.backends.cudnn.allow_tf32 = True
+torch.backends.cudnn.benchmark = True
+torch.set_float32_matmul_precision('high')
+
+# SNAC token offsets for Orpheus
+SNAC_BASE_OFFSET = 128266
+EOS_TOKEN = 128009
+
+
+def log(msg):
+    print(msg)
+    sys.stdout.flush()
+
+
+def apply_snac_offset(token_idx, position):
+    """Apply position-based offset to SNAC token.
+    If token is already offset (>= SNAC_BASE_OFFSET), return as-is.
+    """
+    if int(token_idx) >= SNAC_BASE_OFFSET:
+        # Already has offset applied
+        return int(token_idx)
+    offset = SNAC_BASE_OFFSET + (position % 7) * 4096
+    return int(token_idx) + offset
+
+
+def get_text_ratio(global_step, decay_steps=300, initial_ratio=0.9, min_ratio=0.0):
+    """
+    IST-LM: Start with 90% text, decrease by 0.1 every 300 steps.
+    """
+    num_decays = global_step // decay_steps
+    text_ratio = initial_ratio - (num_decays * 0.1)
+    return max(min_ratio, text_ratio)
+
+
+# ============================================================
+# Async Checkpoint Saving
+# ============================================================
+_save_threads = []
+
+def save_checkpoint_async(state_dict, path, is_main=True):
+    global _save_threads
+    _save_threads = [t for t in _save_threads if t.is_alive()]
+
+    def copy_to_cpu(obj):
+        if isinstance(obj, torch.Tensor):
+            return obj.detach().cpu().clone()
+        elif isinstance(obj, dict):
+            return {k: copy_to_cpu(v) for k, v in obj.items()}
+        return obj
+
+    state_copy = copy_to_cpu(state_dict)
+
+    def _save():
+        try:
+            torch.save(state_copy, path)
+            if is_main:
+                log(f"[ASYNC] Saved: {path}")
+        except Exception as e:
+            if is_main:
+                log(f"[ASYNC] Error: {e}")
+
+    thread = threading.Thread(target=_save, daemon=True)
+    thread.start()
+    _save_threads.append(thread)
+
+
+def wait_for_checkpoints():
+    global _save_threads
+    for t in _save_threads:
+        t.join()
+    _save_threads = []
+
+
+# ============================================================
+# GPU Auto-Detection
+# ============================================================
+def auto_detect_gpu_config():
+    try:
+        import subprocess
+        result = subprocess.run(
+            ['nvidia-smi', '--query-gpu=name,memory.total', '--format=csv,noheader,nounits'],
+            capture_output=True, text=True
+        )
+        lines = result.stdout.strip().split('\n')
+        gpu_name, vram_mb = lines[0].split(', ')
+        vram_gb = int(vram_mb) // 1024
+
+        if vram_gb >= 150:
+            return {"name": "B200", "batch_size": 8, "grad_accum": 4, "vram_gb": vram_gb}
+        elif vram_gb >= 80:
+            return {"name": "A100-80GB", "batch_size": 6, "grad_accum": 5, "vram_gb": vram_gb}
+        elif vram_gb >= 35:
+            return {"name": "A100-40GB", "batch_size": 4, "grad_accum": 8, "vram_gb": vram_gb}
+        else:
+            return {"name": "RTX4090", "batch_size": 2, "grad_accum": 16, "vram_gb": vram_gb}
+    except:
+        return {"name": "Unknown", "batch_size": 2, "grad_accum": 16, "vram_gb": 24}
+
+def get_ram_info():
+    """Get RAM info in GB."""
+    try:
+        import psutil
+        total = psutil.virtual_memory().total / 1024**3
+        available = psutil.virtual_memory().available / 1024**3
+        return total, available
+    except ImportError:
+        try:
+            import subprocess
+            result = subprocess.run(
+                ['free', '-g', '--output=SIZE,AVAILABLE'],
+                capture_output=True, text=True
+            )
+            lines = result.stdout.strip().split('\n')
+            if len(lines) >= 2:
+                total, available = map(float, lines[1].split())
+                return total, available
+        except:
+            pass
+    except:
+        pass
+    return 0, 0
+
+def limit_ram_usage(max_ram_gb):
+    """Limit RAM usage by setting resource limits."""
+    try:
+        import resource
+        max_bytes = int(max_ram_gb * 1024**3)
+        resource.setrlimit(resource.RLIMIT_AS, (max_bytes, max_bytes))
+    except:
+        pass
+
+def log_memory_usage():
+    """Log current memory usage."""
+    msg = []
+    if torch.cuda.is_available():
+        used = torch.cuda.memory_allocated() / 1024**3
+        reserved = torch.cuda.memory_reserved() / 1024**3
+        msg.append(f"GPU: {used:.2f}GB / {reserved:.2f}GB")
+    try:
+        import psutil
+        ram_used = psutil.virtual_memory().used / 1024**3
+        ram_total = psutil.virtual_memory().total / 1024**3
+        msg.append(f"RAM: {ram_used:.1f}GB / {ram_total:.1f}GB")
+    except:
+        pass
+    return " | ".join(msg)
+
+
+# ============================================================
+# Speech Adapter (Same as Stage 1 - LLaMA-Omni 2 Style)
+# ============================================================
+class SpeechAdapter(nn.Module):
+    """
+    5x downsampling + FFN with intermediate dim 2048
+    MUST match Stage 1 architecture exactly for checkpoint loading.
+    """
+    def __init__(self, whisper_dim=1280, llm_dim=3072, downsample=5, intermediate_dim=2048):
+        super().__init__()
+        self.downsample = downsample
+        concat_dim = whisper_dim * downsample
+
+        self.ffn = nn.Sequential(
+            nn.Linear(concat_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, llm_dim),
+            nn.LayerNorm(llm_dim)
+        )
+
+    def forward(self, x):
+        B, T, D = x.shape
+        T_new = (T // self.downsample) * self.downsample
+        x = x[:, :T_new]
+        x = x.reshape(B, T_new // self.downsample, D * self.downsample)
+        return self.ffn(x)
+
+
+# ============================================================
+# Scheduled Interleaved Sequence Creation with Word Alignment
+# ============================================================
+def create_interleaved_sequence(text_tokens, snac_tokens, text_ratio=0.9, word_alignments=None, tokenizer=None, answer_text=None):
+    """
+    Create interleaved sequence based on text_ratio with word-level alignment.
+    - text_ratio=0.9 means 90% of words are replaced by text tokens
+    - text_ratio=0.0 means 100% audio (no text replacement)
+
+    With word_alignments: replaces aligned audio spans with corresponding text tokens
+    Without word_alignments: falls back to positional interleaving
+    """
+    interleaved = []
+    is_audio_mask = []
+
+    if len(snac_tokens) == 0:
+        return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
+
+    # Group SNAC into frames of 7
+    frames = []
+    for i in range(0, len(snac_tokens), 7):
+        frame = snac_tokens[i:i+7]
+        if len(frame) == 7:
+            frames.append(frame)
+
+    if len(frames) == 0:
+        return text_tokens + [EOS_TOKEN], [False] * (len(text_tokens) + 1)
+
+    total_frames = len(frames)
+
+    # If we have word alignments, use semantic interleaving
+    if word_alignments and tokenizer and answer_text and text_ratio > 0:
+        import random
+
+        # Decide which words to replace with text based on text_ratio
+        num_words = len(word_alignments)
+        num_text_words = int(num_words * text_ratio)
+
+        # Randomly select which word indices to replace with text
+        word_indices = list(range(num_words))
+        random.shuffle(word_indices)
+        text_word_indices = set(word_indices[:num_text_words])
+
+        # Build interleaved sequence frame by frame
+        frame_idx = 0
+        snac_position = 0
+
+        for word_idx, alignment in enumerate(word_alignments):
+            word = alignment['word']
+            start_frame = alignment['start_frame']
+            end_frame = min(alignment['end_frame'], total_frames)
+
+            if word_idx in text_word_indices:
+                # Replace this word's audio with text tokens
+                word_tokens = tokenizer.encode(word, add_special_tokens=False)
+                for tok in word_tokens:
+                    interleaved.append(tok)
+                    is_audio_mask.append(False)
+                # Skip the audio frames for this word
+                snac_position = end_frame * 7
+            else:
+                # Keep audio for this word
+                for f_idx in range(start_frame, end_frame):
+                    if f_idx < total_frames:
+                        frame = frames[f_idx]
+                        for tok in frame:
+                            interleaved.append(apply_snac_offset(tok, snac_position))
+                            is_audio_mask.append(True)
+                            snac_position += 1
+
+            frame_idx = end_frame
+
+        # Add any remaining frames after the last word
+        while frame_idx < total_frames:
+            frame = frames[frame_idx]
+            for tok in frame:
+                interleaved.append(apply_snac_offset(tok, snac_position))
+                is_audio_mask.append(True)
+                snac_position += 1
+            frame_idx += 1
+
+    else:
+        # Fallback: positional interleaving (original behavior)
+        total_text = len(text_tokens)
+
+        # Determine interleaving pattern based on text_ratio
+        if text_ratio >= 0.9:
+            text_per_chunk, frames_per_chunk = 1, 3
+        elif text_ratio >= 0.7:
+            text_per_chunk, frames_per_chunk = 1, 5
+        elif text_ratio >= 0.5:
+            text_per_chunk, frames_per_chunk = 1, 7
+        elif text_ratio >= 0.3:
+            text_per_chunk, frames_per_chunk = 1, 10
+        else:
+            text_per_chunk, frames_per_chunk = 0, 1
+
+        text_idx = 0
+        frame_idx = 0
+        snac_position = 0
+
+        while frame_idx < total_frames:
+            if text_per_chunk > 0 and text_idx < total_text:
+                for _ in range(text_per_chunk):
+                    if text_idx < total_text:
+                        interleaved.append(text_tokens[text_idx])
+                        is_audio_mask.append(False)
+                        text_idx += 1
+
+            for _ in range(frames_per_chunk):
+                if frame_idx < total_frames:
+                    frame = frames[frame_idx]
+                    for tok in frame:
+                        interleaved.append(apply_snac_offset(tok, snac_position))
+                        is_audio_mask.append(True)
+                        snac_position += 1
+                    frame_idx += 1
+
+        while text_idx < total_text:
+            interleaved.append(text_tokens[text_idx])
+            is_audio_mask.append(False)
+            text_idx += 1
+
+    # Add EOS
+    interleaved.append(EOS_TOKEN)
+    is_audio_mask.append(False)
+
+    return interleaved, is_audio_mask
+
+
+# ============================================================
+# Dataset
+# ============================================================
+class InterleavedDataset(Dataset):
+    def __init__(self, data, tokenizer, max_audio_len=500, max_seq_len=2048):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_audio = max_audio_len * 5
+        self.max_seq_len = max_seq_len
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        item = self.data[idx]
+
+        # Whisper features
+        whisper = item["whisper_features"][:self.max_audio]
+
+        # Text tokens - use pre-computed if available, otherwise tokenize
+        if "text_tokens" in item and len(item["text_tokens"]) > 0:
+            tt = item["text_tokens"]
+            text_tokens = tt.tolist() if hasattr(tt, 'tolist') else list(tt)
+        else:
+            text = item.get("answer", item.get("text", ""))
+            if isinstance(text, str) and len(text) > 0:
+                text_tokens = self.tokenizer.encode(text, add_special_tokens=False)
+            else:
+                text_tokens = []
+
+        # SNAC tokens
+        snac = item["snac_tokens"]
+        snac_len = (len(snac) // 7) * 7
+        snac = snac[:snac_len] if snac_len > 0 else snac[:7]
+        snac_list = snac.tolist() if hasattr(snac, 'tolist') else list(snac)
+
+        # Word alignments (if available)
+        word_alignments = item.get("word_alignments", None)
+        answer_text = item.get("answer", "")
+
+        return {
+            "whisper": whisper,
+            "text_tokens": text_tokens,
+            "snac_tokens": snac_list,
+            "word_alignments": word_alignments,
+            "answer_text": answer_text
+        }
+
+
+def collate_fn(batch, text_ratio=0.9, tokenizer=None):
+    """Collate with dynamic interleaving based on text_ratio and word alignments."""
+    max_w = max(b["whisper"].shape[0] for b in batch)
+    max_w = ((max_w + 4) // 5) * 5
+
+    whisper_batch = []
+    interleaved_batch = []
+    is_audio_batch = []
+
+    max_seq = 0
+    sequences = []
+
+    for b in batch:
+        interleaved, is_audio = create_interleaved_sequence(
+            b["text_tokens"],
+            b["snac_tokens"],
+            text_ratio,
+            word_alignments=b.get("word_alignments"),
+            tokenizer=tokenizer,
+            answer_text=b.get("answer_text")
+        )
+        sequences.append((interleaved, is_audio))
+        max_seq = max(max_seq, len(interleaved))
+
+    for i, b in enumerate(batch):
+        w = b["whisper"]
+        w_pad = F.pad(w, (0, 0, 0, max_w - w.shape[0]))
+        whisper_batch.append(w_pad)
+
+        interleaved, is_audio = sequences[i]
+        seq_tensor = torch.tensor(interleaved, dtype=torch.long)
+        mask_tensor = torch.tensor(is_audio, dtype=torch.bool)
+
+        seq_pad = F.pad(seq_tensor, (0, max_seq - len(interleaved)), value=-100)
+        mask_pad = F.pad(mask_tensor, (0, max_seq - len(is_audio)), value=False)
+
+        interleaved_batch.append(seq_pad)
+        is_audio_batch.append(mask_pad)
+
+    return {
+        "whisper": torch.stack(whisper_batch),
+        "interleaved": torch.stack(interleaved_batch),
+        "is_audio_mask": torch.stack(is_audio_batch)
+    }
+
+
+# ============================================================
+# Arguments
+# ============================================================
+def parse_args():
+    parser = argparse.ArgumentParser(description="Stage 2: Adapter + LoRA with Interleaved Output")
+    parser.add_argument("--data", type=str, required=True)
+    parser.add_argument("--output_dir", type=str, default="./checkpoints")
+    parser.add_argument("--stage1_ckpt", type=str, default=None,
+                        help="Path to Stage 1 adapter checkpoint (required for proper init)")
+    parser.add_argument("--lr", type=float, default=5e-5)
+    parser.add_argument("--epochs", type=int, default=3, help="3+ epochs recommended")
+    parser.add_argument("--batch_size", type=int, default=None)
+    parser.add_argument("--grad_accum", type=int, default=None)
+    parser.add_argument("--warmup_ratio", type=float, default=0.03)
+    parser.add_argument("--max_grad_norm", type=float, default=1.0)
+    parser.add_argument("--save_steps", type=int, default=200)
+    parser.add_argument("--label_smoothing", type=float, default=0.1)
+    # Scheduled interleaving
+    parser.add_argument("--initial_text_ratio", type=float, default=0.9)
+    parser.add_argument("--decay_steps", type=int, default=300)
+    # LoRA config
+    parser.add_argument("--lora_r", type=int, default=16)
+    parser.add_argument("--lora_alpha", type=int, default=32)
+    parser.add_argument("--lora_dropout", type=float, default=0.05)
+    # Model
+    parser.add_argument("--model_path", type=str, default="canopylabs/3b-es_it-ft-research_release")
+    parser.add_argument("--resume", type=str, default=None)
+    # Memory limits
+    parser.add_argument("--vram_fraction", type=float, default=0.80, help="VRAM fraction to use (default 0.80)")
+    parser.add_argument("--ram_limit_gb", type=float, default=None, help="RAM limit in GB (auto if not specified)")
+    # Modes
+    parser.add_argument("--demo", action="store_true")
+    parser.add_argument("--test", action="store_true")
+    return parser.parse_args()
+
+
+# ============================================================
+# Main
+# ============================================================
+def main():
+    args = parse_args()
+
+    # DDP setup
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+
+    if world_size > 1:
+        dist.init_process_group("nccl")
+        torch.cuda.set_device(local_rank)
+
+    device = torch.device(f"cuda:{local_rank}")
+    is_main = local_rank == 0
+
+    # GPU config
+    gpu_config = auto_detect_gpu_config()
+    if args.batch_size is None:
+        args.batch_size = gpu_config["batch_size"]
+    if args.grad_accum is None:
+        args.grad_accum = gpu_config["grad_accum"]
+
+    torch_dtype = torch.bfloat16
+
+    # Get RAM info and set limits
+    ram_total, ram_available = get_ram_info()
+    if args.ram_limit_gb is None:
+        args.ram_limit_gb = ram_total * 0.80  # Default to 80% of RAM
+    limit_ram_usage(args.ram_limit_gb)
+
+    if is_main:
+        log("=" * 60)
+        log("STAGE 2: Adapter + LoRA (Both Trainable) + Interleaved Output")
+        log("=" * 60)
+        log(f"GPU: {gpu_config['name']} ({gpu_config['vram_gb']}GB)")
+        log(f"RAM: {ram_total:.1f}GB total, {ram_available:.1f}GB available")
+        log(f"Batch: {args.batch_size}, Grad accum: {args.grad_accum}")
+        log(f"LR: {args.lr}, Epochs: {args.epochs}")
+        log(f"LoRA: r={args.lora_r}, alpha={args.lora_alpha}")
+        log(f"VRAM limit: {args.vram_fraction*100:.0f}%")
+        log(f"RAM limit: {args.ram_limit_gb:.1f}GB (80% of total)")
+        log(f"Initial text ratio: {args.initial_text_ratio}")
+        log(f"Decay steps: {args.decay_steps}")
+
+    # Apply VRAM and RAM limits
+    if str(device).startswith('cuda'):
+        torch.cuda.set_per_process_memory_fraction(args.vram_fraction)
+        torch.cuda.empty_cache()
+        torch.backends.cudnn.benchmark = True
+        torch.set_float32_matmul_precision('high')
+        if is_main:
+            log(f"[MEMORY] VRAM limited to {args.vram_fraction*100:.0f}%")
+            log(f"[MEMORY] RAM limited to {args.ram_limit_gb:.1f}GB")
+
+    # HuggingFace login
+    hf_token = os.environ.get("HF_TOKEN")
+    if hf_token:
+        login(token=hf_token)
+
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Load datasets
+    data_paths = [p.strip() for p in args.data.split(",")]
+    all_datasets = []
+
+    if is_main:
+        log("\nLoading datasets...")
+
+    for path in data_paths:
+        if os.path.exists(path):
+            data = torch.load(path, weights_only=False, mmap=True)
+            dataset = InterleavedDataset(data, tokenizer)
+            all_datasets.append(dataset)
+            if is_main:
+                log(f"  {os.path.basename(path)}: {len(data):,} samples")
+
+    if len(all_datasets) == 0:
+        raise ValueError("No datasets loaded!")
+
+    combined_dataset = ConcatDataset(all_datasets) if len(all_datasets) > 1 else all_datasets[0]
+
+    # Demo/Test mode
+    if args.test:
+        combined_dataset = torch.utils.data.Subset(combined_dataset, range(min(5, len(combined_dataset))))
+        args.batch_size = min(args.batch_size, len(combined_dataset))
+        args.grad_accum = 1
+    elif args.demo:
+        combined_dataset = torch.utils.data.Subset(combined_dataset, range(min(1000, len(combined_dataset))))
+        args.batch_size = min(4, args.batch_size)
+        args.grad_accum = max(8, args.grad_accum)
+
+    if is_main:
+        log(f"Total samples: {len(combined_dataset):,}")
+
+    # Load LLM with LoRA
+    if is_main:
+        log(f"\nLoading LLM with LoRA: {args.model_path}")
+
+    llm = AutoModelForCausalLM.from_pretrained(
+        args.model_path,
+        torch_dtype=torch_dtype,
+        attn_implementation="sdpa",
+        device_map={"": device}
+    )
+
+    # Apply LoRA
+    lora_config = LoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        task_type=TaskType.CAUSAL_LM
+    )
+    llm = get_peft_model(llm, lora_config)
+
+    if is_main:
+        llm.print_trainable_parameters()
+
+    # Create adapter (TRAINABLE - same architecture as Stage 1)
+    adapter = SpeechAdapter(
+        whisper_dim=1280,
+        llm_dim=3072,
+        downsample=5,
+        intermediate_dim=2048
+    ).to(device, dtype=torch_dtype)
+
+    # Load Stage 1 adapter checkpoint
+    if args.stage1_ckpt and os.path.exists(args.stage1_ckpt):
+        if is_main:
+            log(f"\nLoading Stage 1 adapter from: {args.stage1_ckpt}")
+        ckpt = torch.load(args.stage1_ckpt, map_location=device, weights_only=False)
+        if "adapter" in ckpt:
+            adapter.load_state_dict(ckpt["adapter"])
+            if is_main:
+                log("  Adapter weights loaded successfully!")
+                if "loss" in ckpt:
+                    log(f"  Stage 1 final loss: {ckpt['loss']:.4f}")
+                if "text_ratio" in ckpt:
+                    log(f"  Stage 1 final text_ratio: {ckpt['text_ratio']:.1f}")
+        else:
+            if is_main:
+                log("  WARNING: No 'adapter' key in checkpoint, using random init")
+    else:
+        if is_main:
+            log("\nWARNING: No Stage 1 checkpoint provided, adapter starting from random init")
+            log("  Recommended: --stage1_ckpt checkpoints/stage1_best.pt")
+
+    adapter_params = sum(p.numel() for p in adapter.parameters())
+    lora_params = sum(p.numel() for p in llm.parameters() if p.requires_grad)
+
+    if is_main:
+        log(f"\nTrainable parameters:")
+        log(f"  Adapter: {adapter_params:,} ({adapter_params/1e6:.1f}M)")
+        log(f"  LoRA: {lora_params:,} ({lora_params/1e6:.1f}M)")
+        log(f"  Total: {adapter_params + lora_params:,} ({(adapter_params + lora_params)/1e6:.1f}M)")
+
+    # DDP
+    if world_size > 1:
+        adapter = DDP(adapter, device_ids=[local_rank])
+        llm = DDP(llm, device_ids=[local_rank])
+        dist.barrier()
+
+    # Optimizer (Adapter + LoRA together)
+    all_params = list(adapter.parameters()) + [p for p in llm.parameters() if p.requires_grad]
+    optimizer = torch.optim.AdamW(all_params, lr=args.lr, weight_decay=0.01)
+
+    # Training state
+    global_step = 0
+    start_epoch = 0
+    best_loss = float("inf")
+    current_text_ratio = args.initial_text_ratio
+
+    # Resume from Stage 2 checkpoint
+    if args.resume and os.path.exists(args.resume):
+        if is_main:
+            log(f"\nResuming from: {args.resume}")
+        ckpt = torch.load(args.resume, map_location=device, weights_only=False)
+        adapter_module = adapter.module if hasattr(adapter, "module") else adapter
+        llm_module = llm.module if hasattr(llm, "module") else llm
+        if "adapter" in ckpt:
+            adapter_module.load_state_dict(ckpt["adapter"])
+        if "lora" in ckpt:
+            llm_module.load_state_dict(ckpt["lora"], strict=False)
+        if "optimizer" in ckpt:
+            optimizer.load_state_dict(ckpt["optimizer"])
+        if "step" in ckpt:
+            global_step = ckpt["step"]
+        if "epoch" in ckpt:
+            start_epoch = ckpt["epoch"]
+        if "text_ratio" in ckpt:
+            current_text_ratio = ckpt["text_ratio"]
+        if "loss" in ckpt:
+            best_loss = ckpt["loss"]
+
+    # Training
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Create dataloader first to calculate steps
+    def collate_with_ratio(batch):
+        current_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
+        return collate_fn(batch, current_ratio, tokenizer=tokenizer)
+
+    if world_size > 1:
+        sampler = DistributedSampler(combined_dataset, shuffle=True)
+        sampler.set_epoch(0)
+    else:
+        sampler = None
+
+    train_loader = DataLoader(
+        combined_dataset,
+        batch_size=args.batch_size,
+        shuffle=(sampler is None),
+        sampler=sampler,
+        collate_fn=collate_with_ratio,
+        num_workers=4,
+        pin_memory=True
+    )
+
+    steps_per_epoch = max(1, len(train_loader) // args.grad_accum)
+    total_steps = max(1, steps_per_epoch * args.epochs)
+    warmup_steps = int(total_steps * args.warmup_ratio)
+
+    scheduler = CosineAnnealingLR(optimizer, T_max=max(1, total_steps - warmup_steps), eta_min=1e-6)
+
+    if is_main:
+        log(f"Steps per epoch: {steps_per_epoch}, Total: {total_steps}, Warmup: {warmup_steps}")
+
+    if is_main:
+        log("\n" + "=" * 60)
+        log("STARTING STAGE 2 TRAINING")
+        log("=" * 60)
+
+    for epoch in range(start_epoch, args.epochs):
+        # Update text ratio based on global step
+        current_text_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
+
+        # Update sampler for this epoch
+        if world_size > 1:
+            sampler.set_epoch(epoch)
+
+        # Recreate dataloader with current text_ratio
+        def collate_with_ratio(batch):
+            return collate_fn(batch, current_text_ratio, tokenizer=tokenizer)
+
+        train_loader = DataLoader(
+            combined_dataset,
+            batch_size=args.batch_size,
+            shuffle=(sampler is None),
+            sampler=sampler,
+            collate_fn=collate_with_ratio,
+            num_workers=4,
+            pin_memory=True
+        )
+
+        adapter.train()
+        llm.train()
+        epoch_loss = 0
+        accum_loss = 0
+
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs}", disable=not is_main)
+
+        for batch_idx, batch in enumerate(pbar):
+            # Update text ratio dynamically
+            current_text_ratio = get_text_ratio(global_step, args.decay_steps, args.initial_text_ratio)
+
+            whisper = batch["whisper"].to(device, dtype=torch_dtype, non_blocking=True)
+            interleaved = batch["interleaved"].to(device, non_blocking=True)
+
+            # Forward through adapter
+            adapter_module = adapter.module if hasattr(adapter, "module") else adapter
+            llm_module = llm.module if hasattr(llm, "module") else llm
+            audio_embeds = adapter_module(whisper)
+
+            # Get token embeddings for interleaved sequence (teacher forcing)
+            input_tokens = interleaved[:, :-1].clamp(min=0)
+            with torch.no_grad():
+                # Access base model's embed_tokens through PEFT wrapper
+                token_embeds = llm_module.get_base_model().model.embed_tokens(input_tokens)
+
+            # Combine: [audio_embeds] + [token_embeds]
+            combined = torch.cat([audio_embeds, token_embeds], dim=1)
+
+            # Forward through LLM with LoRA (trainable)
+            outputs = llm(inputs_embeds=combined, use_cache=False)
+            logits = outputs.logits
+
+            # Loss: predict interleaved tokens after audio prefix
+            # logits[:, audio_len-1] predicts first token of interleaved
+            # logits[:, audio_len-1+i] predicts interleaved[i]
+            audio_len = audio_embeds.shape[1]
+            seq_len = interleaved.shape[1]
+
+            # Get logits that predict the interleaved sequence
+            # Position audio_len-1 predicts interleaved[0], etc.
+            seq_logits = logits[:, audio_len-1:audio_len-1+seq_len]
+            targets = interleaved
+
+            loss = F.cross_entropy(
+                seq_logits.reshape(-1, logits.size(-1)),
+                targets.reshape(-1),
+                ignore_index=-100,
+                label_smoothing=args.label_smoothing
+            )
+
+            loss = loss / args.grad_accum
+            loss.backward()
+            accum_loss += loss.item() * args.grad_accum
+
+            # Update
+            if (batch_idx + 1) % args.grad_accum == 0 or (batch_idx + 1) == len(train_loader):
+                torch.nn.utils.clip_grad_norm_(all_params, args.max_grad_norm)
+                optimizer.step()
+                optimizer.zero_grad(set_to_none=True)
+
+                if global_step < warmup_steps:
+                    lr_scale = (global_step + 1) / warmup_steps
+                    for pg in optimizer.param_groups:
+                        pg["lr"] = args.lr * lr_scale
+                else:
+                    scheduler.step()
+
+                global_step += 1
+                epoch_loss += accum_loss
+
+                pbar.set_postfix(
+                    loss=f"{accum_loss:.4f}",
+                    text_ratio=f"{current_text_ratio:.1f}",
+                    lr=f"{optimizer.param_groups[0]['lr']:.2e}"
+                )
+
+                # Save checkpoint
+                if global_step % args.save_steps == 0 and is_main:
+                    adapter_to_save = adapter.module if hasattr(adapter, "module") else adapter
+                    llm_to_save = llm.module if hasattr(llm, "module") else llm
+                    ckpt_path = os.path.join(args.output_dir, f"stage2_step{global_step}.pt")
+                    save_checkpoint_async({
+                        "adapter": adapter_to_save.state_dict(),
+                        "lora": llm_to_save.state_dict(),
+                        "optimizer": optimizer.state_dict(),
+                        "step": global_step,
+                        "epoch": epoch,
+                        "loss": accum_loss,
+                        "text_ratio": current_text_ratio
+                    }, ckpt_path, is_main)
+
+                accum_loss = 0
+
+        # Epoch end
+        avg_loss = epoch_loss / max(1, steps_per_epoch)
+
+        if is_main:
+            log(f"Epoch {epoch+1} avg loss: {avg_loss:.4f}, text_ratio: {current_text_ratio:.1f}")
+
+            adapter_to_save = adapter.module if hasattr(adapter, "module") else adapter
+            llm_to_save = llm.module if hasattr(llm, "module") else llm
+            ckpt_path = os.path.join(args.output_dir, f"stage2_epoch{epoch+1}.pt")
+            save_checkpoint_async({
+                "adapter": adapter_to_save.state_dict(),
+                "lora": llm_to_save.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "step": global_step,
+                "epoch": epoch + 1,
+                "loss": avg_loss,
+                "text_ratio": current_text_ratio
+            }, ckpt_path, is_main)
+
+            if avg_loss < best_loss:
+                best_loss = avg_loss
+                best_path = os.path.join(args.output_dir, "stage2_best.pt")
+                save_checkpoint_async({
+                    "adapter": adapter_to_save.state_dict(),
+                    "lora": llm_to_save.state_dict(),
+                    "step": global_step,
+                    "epoch": epoch + 1,
+                    "loss": best_loss,
+                    "text_ratio": current_text_ratio
+                }, best_path, is_main)
+                log(f"Best model saved! Loss: {best_loss:.4f}")
+
+    # Finish
+    if is_main:
+        wait_for_checkpoints()
+        log("\n" + "=" * 60)
+        log("STAGE 2 COMPLETE!")
+        log(f"Best loss: {best_loss:.4f}")
+        log(f"Final text_ratio: {current_text_ratio:.1f}")
+        log("=" * 60)
+        log("\nCheckpoints saved:")
+        log(f"  Best: {args.output_dir}/stage2_best.pt")
+        log(f"  Last: {args.output_dir}/stage2_epoch{args.epochs}.pt")
+
+    if world_size > 1:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

passo4_inference.py

@@ -0,0 +1,239 @@
+#!/usr/bin/env python3
+"""
+Passo 4: Inference - Speech-to-Speech
+Load trained model and generate audio responses from audio input.
+"""
+
+import os
+import sys
+import argparse
+import torch
+import torchaudio
+import numpy as np
+from pathlib import Path
+
+# SNAC token configuration
+SNAC_BASE = 128266
+EOS_TOKEN = 128009
+
+def load_models(checkpoint_path: str, device: str = "cuda"):
+    """Load all models for inference."""
+    from transformers import WhisperModel, WhisperFeatureExtractor, AutoTokenizer
+    from peft import PeftModel
+    import snac
+
+    print("Loading models...")
+
+    # Load Whisper encoder
+    print("  Whisper encoder...")
+    whisper_model = WhisperModel.from_pretrained("openai/whisper-large-v3").to(device)
+    whisper_model.eval()
+    feature_extractor = WhisperFeatureExtractor.from_pretrained("openai/whisper-large-v3")
+
+    # Load checkpoint
+    print(f"  Loading checkpoint: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+
+    # Load LLM with LoRA
+    print("  Loading LLM with LoRA...")
+    from transformers import AutoModelForCausalLM
+
+    llm = AutoModelForCausalLM.from_pretrained(
+        "canopylabs/3b-es_it-ft-research_release",
+        torch_dtype=torch.bfloat16,
+        device_map=device
+    )
+
+    # Load LoRA weights if present
+    if 'lora' in checkpoint or 'lora_state_dict' in checkpoint:
+        print("  Loading LoRA weights...")
+        from peft import LoraConfig, get_peft_model
+        lora_config = LoraConfig(
+            r=16,
+            lora_alpha=32,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+            lora_dropout=0.05,
+            bias="none",
+            task_type="CAUSAL_LM"
+        )
+        llm = get_peft_model(llm, lora_config)
+        lora_key = 'lora' if 'lora' in checkpoint else 'lora_state_dict'
+        llm.load_state_dict(checkpoint[lora_key], strict=False)
+
+    llm.eval()
+
+    # Load Speech Adapter
+    print("  Loading Speech Adapter...")
+    from passo2_finetune_stage1 import SpeechAdapter
+    adapter = SpeechAdapter(
+        whisper_dim=1280,
+        llm_dim=3072,
+        downsample=5
+    ).to(device)
+
+    adapter_key = 'adapter' if 'adapter' in checkpoint else 'adapter_state_dict'
+    if adapter_key in checkpoint:
+        adapter.load_state_dict(checkpoint[adapter_key])
+    adapter.eval()
+
+    # Load SNAC decoder
+    print("  SNAC decoder...")
+    snac_model = snac.SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device)
+    snac_model.eval()
+
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("canopylabs/3b-es_it-ft-research_release")
+
+    print("Models loaded!")
+    return whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer
+
+
+def encode_audio(audio_path: str, whisper_model, feature_extractor, adapter, device: str):
+    """Encode input audio to embeddings."""
+    # Load audio
+    waveform, sr = torchaudio.load(audio_path)
+    if sr != 16000:
+        waveform = torchaudio.functional.resample(waveform, sr, 16000)
+
+    # Extract Whisper features
+    with torch.no_grad():
+        inputs = feature_extractor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt")
+        whisper_features = whisper_model.encoder(
+            inputs.input_features.to(device)
+        ).last_hidden_state
+
+        # Adapt to LLM space
+        adapted = adapter(whisper_features)
+
+    return adapted.to(torch.bfloat16)
+
+
+def decode_snac_tokens(tokens: list, snac_model, device: str):
+    """Decode SNAC tokens to audio waveform."""
+    # Remove offsets and reshape
+    # SNAC has 3 layers with 1:2:4 ratio
+    layer0_tokens = []
+    layer1_tokens = []
+    layer2_tokens = []
+
+    # Parse interleaved tokens: 1 from layer0, 2 from layer1, 4 from layer2 per frame
+    i = 0
+    while i < len(tokens):
+        if i < len(tokens):
+            layer0_tokens.append(tokens[i] - SNAC_BASE)
+            i += 1
+        if i < len(tokens):
+            layer1_tokens.append(tokens[i] - SNAC_BASE - 4096)
+            i += 1
+        if i < len(tokens):
+            layer1_tokens.append(tokens[i] - SNAC_BASE - 4096)
+            i += 1
+        for _ in range(4):
+            if i < len(tokens):
+                layer2_tokens.append(tokens[i] - SNAC_BASE - 8192)
+                i += 1
+
+    # Create tensors
+    min_len = min(len(layer0_tokens), len(layer1_tokens) // 2, len(layer2_tokens) // 4)
+    if min_len == 0:
+        return np.zeros(24000, dtype=np.float32)  # 1 second of silence
+
+    codes = [
+        torch.tensor([layer0_tokens[:min_len]], dtype=torch.long, device=device),
+        torch.tensor([layer1_tokens[:min_len * 2]], dtype=torch.long, device=device),
+        torch.tensor([layer2_tokens[:min_len * 4]], dtype=torch.long, device=device)
+    ]
+
+    # Decode
+    with torch.no_grad():
+        audio = snac_model.decode(codes)
+
+    return audio.squeeze().cpu().numpy()
+
+
+def generate_response(
+    audio_input: str,
+    whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer,
+    device: str,
+    max_new_tokens: int = 500,
+    temperature: float = 0.7
+):
+    """Generate speech response from audio input."""
+
+    # Encode input audio
+    print("Encoding input audio...")
+    audio_embeddings = encode_audio(audio_input, whisper_model, feature_extractor, adapter, device)
+
+    # Generate with LLM
+    print("Generating response...")
+
+    # Create input embeds by concatenating audio embeddings with start tokens
+    with torch.no_grad():
+        # Get embeddings layer
+        embed_layer = llm.get_input_embeddings()
+
+        # Add start token
+        start_tokens = tokenizer.encode("<|start|>", add_special_tokens=False, return_tensors="pt").to(device)
+        start_embeds = embed_layer(start_tokens)
+
+        # Concatenate: audio + start (ensure same dtype)
+        input_embeds = torch.cat([audio_embeddings.to(torch.bfloat16), start_embeds.to(torch.bfloat16)], dim=1)
+
+        # Generate
+        outputs = llm.generate(
+            inputs_embeds=input_embeds,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            do_sample=True,
+            top_p=0.9,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=EOS_TOKEN
+        )
+
+    # Extract SNAC tokens from output
+    generated_tokens = outputs[0].tolist()
+    snac_tokens = [t for t in generated_tokens if SNAC_BASE <= t < SNAC_BASE + 3 * 4096]
+
+    print(f"Generated {len(snac_tokens)} SNAC tokens")
+
+    # Decode to audio
+    print("Decoding to audio...")
+    audio_output = decode_snac_tokens(snac_tokens, snac_model, device)
+
+    return audio_output
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Speech-to-Speech Inference")
+    parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint")
+    parser.add_argument("--input", type=str, required=True, help="Input audio file")
+    parser.add_argument("--output", type=str, default="output.wav", help="Output audio file")
+    parser.add_argument("--max_tokens", type=int, default=500, help="Max tokens to generate")
+    parser.add_argument("--temperature", type=float, default=0.7, help="Sampling temperature")
+    args = parser.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Device: {device}")
+
+    # Load models
+    whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer = load_models(
+        args.checkpoint, device
+    )
+
+    # Generate
+    audio_output = generate_response(
+        args.input,
+        whisper_model, feature_extractor, llm, adapter, snac_model, tokenizer,
+        device,
+        max_new_tokens=args.max_tokens,
+        temperature=args.temperature
+    )
+
+    # Save output
+    import soundfile as sf
+    sf.write(args.output, audio_output, 24000)
+    print(f"Saved output to: {args.output}")
+
+
+if __name__ == "__main__":
+    main()

services/__init__.py

@@ -0,0 +1,9 @@
+"""
+Services for dataset generation.
+"""
+from .groq_service import GroqService
+from .whisperx_service import WhisperXService
+from .kokoro_service import KokoroService
+from .llama_service import LlamaService
+
+__all__ = ['GroqService', 'WhisperXService', 'KokoroService', 'LlamaService']

services/groq_service.py

@@ -0,0 +1,106 @@
+"""
+Groq API Service for Q&A generation.
+"""
+import os
+import requests
+from typing import List, Dict, Optional
+
+GROQ_API_KEY = os.environ.get('GROQ_API_KEY', '')
+GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
+DEFAULT_MODEL = "llama-3.1-8b-instant"
+
+
+class GroqService:
+    """Service for generating Q&A pairs using Groq API."""
+
+    def __init__(self, api_key: Optional[str] = None, model: str = DEFAULT_MODEL):
+        self.api_key = api_key or GROQ_API_KEY
+        self.model = model
+        self.headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
+
+    def generate_qa_pairs(self, count: int, timeout: int = 30) -> List[Dict[str, str]]:
+        """Generate Q&A pairs.
+
+        Args:
+            count: Number of Q&A pairs to generate.
+            timeout: Request timeout in seconds.
+
+        Returns:
+            List of dicts with 'question' and 'answer' keys.
+        """
+        prompt = f"""Generate {count} unique question-answer pairs about general knowledge.
+REQUIREMENTS:
+- Questions: 5 to 10 words
+- Answers: 5 to 10 words
+Format: Q: [question]
+A: [answer]
+Return exactly {count} pairs."""
+
+        payload = {
+            "messages": [{"role": "user", "content": prompt}],
+            "model": self.model,
+            "temperature": 0.7,
+            "max_tokens": 2000
+        }
+
+        try:
+            response = requests.post(
+                GROQ_API_URL,
+                headers=self.headers,
+                json=payload,
+                timeout=timeout
+            )
+            response.raise_for_status()
+            content = response.json()['choices'][0]['message']['content']
+            return self._parse_qa_pairs(content)
+        except Exception as e:
+            print(f"[GroqService] Error: {e}")
+            return []
+
+    def _parse_qa_pairs(self, content: str) -> List[Dict[str, str]]:
+        """Parse Q&A pairs from API response."""
+        pairs = []
+        lines = content.split('\n')
+        current_q, current_a = None, None
+
+        for line in lines:
+            line = line.strip()
+            if line.lower().startswith('q:'):
+                current_q = line[2:].strip()
+            elif line.lower().startswith('a:'):
+                current_a = line[2:].strip()
+                if current_q and current_a:
+                    word_count_q = len(current_q.split())
+                    word_count_a = len(current_a.split())
+                    if 3 <= word_count_q <= 15 and 3 <= word_count_a <= 15:
+                        pairs.append({'question': current_q, 'answer': current_a})
+                    current_q, current_a = None, None
+
+        return pairs
+
+    def generate_batch(self, total_count: int, batch_size: int = 15) -> List[Dict[str, str]]:
+        """Generate multiple batches of Q&A pairs.
+
+        Args:
+            total_count: Total number of pairs needed.
+            batch_size: Pairs per API call.
+
+        Returns:
+            List of all generated Q&A pairs.
+        """
+        all_pairs = []
+        remaining = total_count
+
+        while remaining > 0:
+            count = min(batch_size, remaining)
+            pairs = self.generate_qa_pairs(count)
+            all_pairs.extend(pairs)
+            remaining -= len(pairs)
+
+            if not pairs:
+                break
+
+        return all_pairs[:total_count]

services/kokoro_service.py

@@ -0,0 +1,238 @@
+"""
+Kokoro TTS Service for text-to-speech synthesis.
+
+IMPORTANT: For GPU acceleration, set these environment variables BEFORE importing:
+    export LD_LIBRARY_PATH="/usr/local/lib/python3.12/dist-packages/nvidia/cublas/lib:..."
+    export ONNX_PROVIDER=CUDAExecutionProvider
+
+Or use: source /tmp/setup_env.sh
+"""
+import os
+import numpy as np
+from typing import Tuple, List, Optional
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+# Default paths
+DEFAULT_MODEL_PATH = '/tmp/kokoro-v1.0.onnx'
+DEFAULT_VOICES_PATH = '/tmp/voices-v1.0.bin'
+DEFAULT_SAMPLE_RATE = 24000
+
+
+class KokoroService:
+    """Service for Kokoro TTS synthesis."""
+
+    def __init__(
+        self,
+        model_path: str = DEFAULT_MODEL_PATH,
+        voices_path: str = DEFAULT_VOICES_PATH,
+        num_workers: int = 1
+    ):
+        self.model_path = model_path
+        self.voices_path = voices_path
+        self.num_workers = num_workers
+        self.sample_rate = DEFAULT_SAMPLE_RATE
+
+        self._models = None
+
+    def load(self) -> 'KokoroService':
+        """Load Kokoro TTS models."""
+        from kokoro_onnx import Kokoro
+
+        print(f"[KokoroTTS] Loading {self.num_workers} model(s)...")
+
+        # Load models
+        self._models = [
+            Kokoro(self.model_path, self.voices_path)
+            for _ in range(self.num_workers)
+        ]
+
+        # Log provider info
+        providers = self._models[0].sess.get_providers()
+        print(f"[KokoroTTS] ONNX Providers: {providers}")
+
+        if 'CUDAExecutionProvider' not in providers:
+            print("[KokoroTTS] WARNING: Running on CPU! Set ONNX_PROVIDER=CUDAExecutionProvider")
+
+        # Warmup
+        for model in self._models:
+            model.create("warmup", voice='af_heart', speed=1.0)
+
+        print(f"[KokoroTTS] Models loaded and warmed up")
+        return self
+
+    def synthesize(
+        self,
+        text: str,
+        voice: str = 'af_heart',
+        speed: float = 1.0
+    ) -> Tuple[np.ndarray, int]:
+        """Synthesize speech from text.
+
+        Args:
+            text: Text to synthesize.
+            voice: Voice ID to use.
+            speed: Speech speed multiplier.
+
+        Returns:
+            Tuple of (audio samples, sample rate).
+        """
+        if self._models is None:
+            self.load()
+
+        samples, sr = self._models[0].create(text, voice=voice, speed=speed)
+        return samples, sr
+
+    def synthesize_batch(
+        self,
+        texts: List[str],
+        voice: str = 'af_heart',
+        speed: float = 1.0
+    ) -> List[Tuple[np.ndarray, int]]:
+        """Synthesize multiple texts in parallel.
+
+        Args:
+            texts: List of texts to synthesize.
+            voice: Voice ID to use.
+            speed: Speech speed multiplier.
+
+        Returns:
+            List of (audio samples, sample rate) tuples.
+        """
+        if self._models is None:
+            self.load()
+
+        def process_text(args):
+            idx, text, model = args
+            try:
+                samples, sr = model.create(text, voice=voice, speed=speed)
+                return idx, samples, sr
+            except Exception as e:
+                print(f"[KokoroTTS] Error synthesizing: {e}")
+                return idx, None, None
+
+        results = [None] * len(texts)
+
+        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
+            tasks = [
+                (i, text, self._models[i % self.num_workers])
+                for i, text in enumerate(texts)
+            ]
+            futures = {executor.submit(process_text, task): task for task in tasks}
+
+            for future in as_completed(futures):
+                idx, samples, sr = future.result()
+                if samples is not None:
+                    results[idx] = (samples, sr)
+
+        return [r for r in results if r is not None]
+
+    def synthesize_qa_pair(
+        self,
+        question: str,
+        answer: str,
+        voice: str = 'af_heart',
+        speed: float = 1.0
+    ) -> dict:
+        """Synthesize a Q&A pair.
+
+        Args:
+            question: Question text.
+            answer: Answer text.
+            voice: Voice ID to use.
+            speed: Speech speed multiplier.
+
+        Returns:
+            Dict with 'q_audio', 'a_audio', 'question', 'answer' keys.
+        """
+        if self._models is None:
+            self.load()
+
+        model = self._models[0]
+        q_samples, _ = model.create(question, voice=voice, speed=speed)
+        a_samples, _ = model.create(answer, voice=voice, speed=speed)
+
+        return {
+            'question': question,
+            'answer': answer,
+            'q_audio': q_samples,
+            'a_audio': a_samples
+        }
+
+    def synthesize_qa_batch(
+        self,
+        qa_pairs: List[dict],
+        voice: str = 'af_heart',
+        speed: float = 1.0
+    ) -> List[dict]:
+        """Synthesize multiple Q&A pairs.
+
+        Args:
+            qa_pairs: List of dicts with 'question' and 'answer' keys.
+            voice: Voice ID to use.
+            speed: Speech speed multiplier.
+
+        Returns:
+            List of dicts with audio added.
+        """
+        if self._models is None:
+            self.load()
+
+        def process_pair(args):
+            idx, pair, model = args
+            try:
+                q_samples, _ = model.create(pair['question'], voice=voice, speed=speed)
+                a_samples, _ = model.create(pair['answer'], voice=voice, speed=speed)
+                return {
+                    'question': pair['question'],
+                    'answer': pair['answer'],
+                    'q_audio': q_samples,
+                    'a_audio': a_samples
+                }
+            except Exception as e:
+                print(f"[KokoroTTS] Error: {e}")
+                return None
+
+        results = []
+
+        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
+            tasks = [
+                (i, pair, self._models[i % self.num_workers])
+                for i, pair in enumerate(qa_pairs)
+            ]
+            futures = {executor.submit(process_pair, task): task for task in tasks}
+
+            for future in as_completed(futures):
+                result = future.result()
+                if result:
+                    results.append(result)
+
+        return results
+
+    @property
+    def is_gpu_enabled(self) -> bool:
+        """Check if GPU acceleration is enabled."""
+        if self._models is None:
+            self.load()
+        providers = self._models[0].sess.get_providers()
+        return 'CUDAExecutionProvider' in providers
+
+    @staticmethod
+    def setup_gpu_environment():
+        """Set up environment variables for GPU acceleration.
+
+        Call this BEFORE importing kokoro_onnx.
+        """
+        nvidia_libs = []
+        for pyver in ['python3.12', 'python3.11', 'python3.10']:
+            nvidia_base = f'/usr/local/lib/{pyver}/dist-packages/nvidia'
+            for subdir in ['cublas', 'cudnn', 'cufft', 'curand', 'cusolver', 'cusparse',
+                           'cuda_runtime', 'cuda_nvrtc', 'cuda_cupti', 'nvjitlink']:
+                lib_path = f"{nvidia_base}/{subdir}/lib"
+                if os.path.exists(lib_path):
+                    nvidia_libs.append(lib_path)
+
+        if nvidia_libs:
+            os.environ['LD_LIBRARY_PATH'] = ':'.join(nvidia_libs) + ':' + os.environ.get('LD_LIBRARY_PATH', '')
+            os.environ['ONNX_PROVIDER'] = 'CUDAExecutionProvider'
+            return True
+        return False

services/llama_service.py

@@ -0,0 +1,212 @@
+"""
+Llama 3.2 Service for Q&A generation using llama.cpp.
+
+Requires: pip install llama-cpp-python
+For GPU: CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python --force-reinstall --no-cache-dir
+"""
+import os
+from typing import List, Dict, Optional
+from pathlib import Path
+
+# Default model path
+DEFAULT_MODEL_PATH = "/tmp/llama-3.2-3b-instruct-q4_k_m.gguf"
+DEFAULT_MODEL_URL = "https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/Llama-3.2-3B-Instruct-Q4_K_M.gguf"
+
+
+class LlamaService:
+    """Service for generating Q&A pairs using Llama 3.2 via llama.cpp."""
+
+    def __init__(
+        self,
+        model_path: str = DEFAULT_MODEL_PATH,
+        n_gpu_layers: int = -1,  # -1 = all layers on GPU
+        n_ctx: int = 2048,
+        n_batch: int = 512,
+        verbose: bool = False
+    ):
+        self.model_path = model_path
+        self.n_gpu_layers = n_gpu_layers
+        self.n_ctx = n_ctx
+        self.n_batch = n_batch
+        self.verbose = verbose
+        self._llm = None
+
+    def download_model(self) -> bool:
+        """Download the model if not present."""
+        if os.path.exists(self.model_path):
+            print(f"[LlamaService] Model already exists: {self.model_path}")
+            return True
+
+        print(f"[LlamaService] Downloading model to {self.model_path}...")
+        import subprocess
+        result = subprocess.run(
+            ["wget", "-q", "--show-progress", "-O", self.model_path, DEFAULT_MODEL_URL],
+            capture_output=False
+        )
+        return result.returncode == 0
+
+    def load(self) -> 'LlamaService':
+        """Load the Llama model."""
+        from llama_cpp import Llama
+
+        if not os.path.exists(self.model_path):
+            self.download_model()
+
+        print(f"[LlamaService] Loading model from {self.model_path}...")
+        print(f"[LlamaService] GPU layers: {self.n_gpu_layers}, Context: {self.n_ctx}")
+
+        self._llm = Llama(
+            model_path=self.model_path,
+            n_gpu_layers=self.n_gpu_layers,
+            n_ctx=self.n_ctx,
+            n_batch=self.n_batch,
+            verbose=self.verbose
+        )
+
+        print(f"[LlamaService] Model loaded successfully")
+        return self
+
+    def generate(
+        self,
+        prompt: str,
+        max_tokens: int = 512,
+        temperature: float = 0.7,
+        stop: Optional[List[str]] = None
+    ) -> str:
+        """Generate text from prompt.
+
+        Args:
+            prompt: The prompt text.
+            max_tokens: Maximum tokens to generate.
+            temperature: Sampling temperature.
+            stop: Stop sequences.
+
+        Returns:
+            Generated text.
+        """
+        if self._llm is None:
+            self.load()
+
+        response = self._llm(
+            prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            stop=stop or [],
+            echo=False
+        )
+
+        return response['choices'][0]['text'].strip()
+
+    def generate_qa_pairs(self, count: int) -> List[Dict[str, str]]:
+        """Generate Q&A pairs.
+
+        Args:
+            count: Number of Q&A pairs to generate.
+
+        Returns:
+            List of dicts with 'question' and 'answer' keys.
+        """
+        if self._llm is None:
+            self.load()
+
+        prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
+You are a helpful assistant that generates question-answer pairs about general knowledge.<|eot_id|><|start_header_id|>user<|end_header_id|>
+Generate {count} unique question-answer pairs about general knowledge.
+REQUIREMENTS:
+- Questions: 5 to 10 words
+- Answers: 5 to 10 words
+Format each pair as:
+Q: [question]
+A: [answer]
+
+Generate exactly {count} pairs now:<|eot_id|><|start_header_id|>assistant<|end_header_id|>
+"""
+
+        response = self._llm(
+            prompt,
+            max_tokens=count * 50,  # ~50 tokens per QA pair
+            temperature=0.7,
+            stop=["<|eot_id|>"],
+            echo=False
+        )
+
+        content = response['choices'][0]['text']
+        return self._parse_qa_pairs(content)
+
+    def _parse_qa_pairs(self, content: str) -> List[Dict[str, str]]:
+        """Parse Q&A pairs from generated text."""
+        pairs = []
+        lines = content.split('\n')
+        current_q, current_a = None, None
+
+        for line in lines:
+            line = line.strip()
+            if line.lower().startswith('q:'):
+                current_q = line[2:].strip()
+            elif line.lower().startswith('a:'):
+                current_a = line[2:].strip()
+                if current_q and current_a:
+                    word_count_q = len(current_q.split())
+                    word_count_a = len(current_a.split())
+                    if 3 <= word_count_q <= 15 and 3 <= word_count_a <= 15:
+                        pairs.append({'question': current_q, 'answer': current_a})
+                    current_q, current_a = None, None
+
+        return pairs
+
+    def generate_batch(self, total_count: int, batch_size: int = 15) -> List[Dict[str, str]]:
+        """Generate multiple batches of Q&A pairs.
+
+        Args:
+            total_count: Total number of pairs needed.
+            batch_size: Pairs per generation call.
+
+        Returns:
+            List of all generated Q&A pairs.
+        """
+        all_pairs = []
+        remaining = total_count
+
+        while remaining > 0:
+            count = min(batch_size, remaining)
+            pairs = self.generate_qa_pairs(count)
+            all_pairs.extend(pairs)
+            remaining -= len(pairs)
+
+            if not pairs:
+                # Retry once if no pairs generated
+                pairs = self.generate_qa_pairs(count)
+                all_pairs.extend(pairs)
+                remaining -= len(pairs)
+                if not pairs:
+                    break
+
+        return all_pairs[:total_count]
+
+    @staticmethod
+    def install_with_cuda():
+        """Install llama-cpp-python with CUDA support."""
+        import subprocess
+        print("[LlamaService] Installing llama-cpp-python with CUDA...")
+        cmd = 'CMAKE_ARGS="-DGGML_CUDA=on" pip install llama-cpp-python --force-reinstall --no-cache-dir'
+        return subprocess.run(cmd, shell=True).returncode == 0
+
+
+# Convenience function for quick testing
+def test_llama_service():
+    """Test the Llama service."""
+    service = LlamaService()
+    service.load()
+
+    print("\nGenerating 5 Q&A pairs...")
+    pairs = service.generate_qa_pairs(5)
+
+    for i, pair in enumerate(pairs, 1):
+        print(f"\n{i}. Q: {pair['question']}")
+        print(f"   A: {pair['answer']}")
+
+    return pairs
+
+
+if __name__ == "__main__":
+    test_llama_service()

services/whisperx_service.py

@@ -0,0 +1,210 @@
+"""
+WhisperX Service for speech recognition, alignment, and feature extraction.
+"""
+import os
+import numpy as np
+import torch
+import torchaudio
+from typing import List, Dict, Optional, Tuple, Any
+from pathlib import Path
+
+# Default configuration
+DEFAULT_MODEL = "large-v3-turbo"
+DEFAULT_COMPUTE_TYPE = "float16"
+DEFAULT_LANGUAGE = "en"
+
+
+class WhisperXService:
+    """Service for WhisperX speech recognition and feature extraction."""
+
+    def __init__(
+        self,
+        model_name: str = DEFAULT_MODEL,
+        device: str = "cuda",
+        compute_type: str = DEFAULT_COMPUTE_TYPE,
+        language: str = DEFAULT_LANGUAGE
+    ):
+        self.model_name = model_name
+        self.device = device
+        self.compute_type = compute_type
+        self.language = language
+
+        self._model = None
+        self._align_model = None
+        self._align_metadata = None
+        self._encoder = None
+
+    def load(self) -> 'WhisperXService':
+        """Load WhisperX models."""
+        import whisperx
+
+        print(f"[WhisperX] Loading model {self.model_name}...")
+
+        # Load main model
+        self._model = whisperx.load_model(
+            self.model_name,
+            self.device,
+            compute_type=self.compute_type,
+            language=self.language,
+            vad_method="silero"
+        )
+
+        # Load alignment model
+        self._align_model, self._align_metadata = whisperx.load_align_model(
+            language_code=self.language,
+            device=self.device
+        )
+
+        # Get internal encoder for feature extraction
+        self._encoder = self._model.model
+
+        print(f"[WhisperX] Models loaded on {self.device}")
+        return self
+
+    def transcribe(self, audio_path: str, batch_size: int = 16) -> Dict:
+        """Transcribe audio file.
+
+        Args:
+            audio_path: Path to audio file.
+            batch_size: Batch size for transcription.
+
+        Returns:
+            Transcription result with segments.
+        """
+        if self._model is None:
+            self.load()
+
+        return self._model.transcribe(audio_path, batch_size=batch_size)
+
+    def align(self, segments: List[Dict], audio_path: str) -> Dict:
+        """Align transcription to audio.
+
+        Args:
+            segments: Transcription segments.
+            audio_path: Path to audio file.
+
+        Returns:
+            Aligned result with word-level timestamps.
+        """
+        import whisperx
+
+        if self._align_model is None:
+            self.load()
+
+        result = {"segments": segments}
+        return whisperx.align(
+            result["segments"],
+            self._align_model,
+            self._align_metadata,
+            audio_path,
+            self.device
+        )
+
+    def transcribe_and_align(self, audio_path: str) -> Dict:
+        """Transcribe and align audio file.
+
+        Returns:
+            Result with word-level alignments.
+        """
+        result = self.transcribe(audio_path)
+        return self.align(result["segments"], audio_path)
+
+    def extract_features(self, audio: np.ndarray, sample_rate: int = 24000) -> torch.Tensor:
+        """Extract encoder features from audio.
+
+        Args:
+            audio: Audio samples as numpy array.
+            sample_rate: Audio sample rate.
+
+        Returns:
+            Encoder features tensor [seq_len, 1280].
+        """
+        import ctranslate2
+
+        if self._encoder is None:
+            self.load()
+
+        # Resample to 16kHz for Whisper
+        if sample_rate != 16000:
+            audio_tensor = torch.from_numpy(audio)
+            audio_16k = torchaudio.functional.resample(audio_tensor, sample_rate, 16000)
+            audio = audio_16k.numpy()
+
+        audio = audio.astype(np.float32)
+
+        # Extract mel features
+        mel_features = self._encoder.feature_extractor(audio)
+
+        # Encode
+        encoded = self._encoder.encode(mel_features)
+
+        # Convert ctranslate2 StorageView to numpy
+        cpu_view = encoded.to_device(ctranslate2.Device.cpu)
+        features_np = np.array(cpu_view, copy=True)  # [1, seq_len, 1280]
+        features = torch.from_numpy(features_np).squeeze(0).float()  # [seq_len, 1280]
+
+        return features
+
+    def batch_extract_features(
+        self,
+        audios: List[np.ndarray],
+        sample_rate: int = 24000
+    ) -> List[torch.Tensor]:
+        """Extract features from multiple audios.
+
+        Args:
+            audios: List of audio samples.
+            sample_rate: Audio sample rate.
+
+        Returns:
+            List of feature tensors.
+        """
+        return [self.extract_features(audio, sample_rate) for audio in audios]
+
+    def get_word_alignments(
+        self,
+        audio_path: str,
+        start_time: float,
+        end_time: float
+    ) -> List[Dict]:
+        """Get word alignments for a segment of audio.
+
+        Args:
+            audio_path: Path to audio file.
+            start_time: Segment start time in seconds.
+            end_time: Segment end time in seconds.
+
+        Returns:
+            List of word alignment dicts.
+        """
+        result = self.transcribe_and_align(audio_path)
+        words = result.get("word_segments", [])
+
+        segment_words = []
+        for w in words:
+            word_start = w.get('start', 0)
+            word_end = w.get('end', 0)
+            if word_start >= start_time - 0.1 and word_end <= end_time + 0.1:
+                segment_words.append({
+                    'word': w.get('word', ''),
+                    'start': word_start - start_time,
+                    'end': word_end - start_time,
+                    'start_frame': int((word_start - start_time) * 75),
+                    'end_frame': int((word_end - start_time) * 75)
+                })
+
+        return segment_words
+
+    @property
+    def model(self):
+        """Get the underlying WhisperX model."""
+        if self._model is None:
+            self.load()
+        return self._model
+
+    @property
+    def encoder(self):
+        """Get the internal encoder for feature extraction."""
+        if self._encoder is None:
+            self.load()
+        return self._encoder