# Voice Designer Scripts

This directory contains a collection of Python scripts for processing and analyzing vocal stem data from v6 metadata files. These scripts form a pipeline for voice data analysis, extraction, sampling, and AI-powered captioning for the Voice Designer project.

## Pipeline Overview

The typical workflow follows this sequence:
1. **Analyze** → Understand the dataset structure and vocal stem distribution
2. **Extract** → Filter records containing vocal stems 
3. **Sample** → Create manageable subsets for processing
4. **Verify** → Validate the extracted data quality
5. **Caption** → Generate AI descriptions of vocal characteristics

## Scripts

### 1. `analyze_v6_data.py`
**Purpose**: Comprehensive analysis of v6 metadata files to understand field distributions and vocal data statistics.

**Key Functions**:
- Field frequency analysis and distribution plots
- Stem name counting and categorization 
- Statistical summaries with sampling support
- Generates analysis reports and visualizations

**Usage**:
```bash
python analyze_v6_data.py --data-path /path/to/metas_v6.jsonl --sample-size 10000 --output-dir ./analysis
```

**Outputs**: JSON reports, frequency distributions, matplotlib visualizations

---

### 2. `extract_vocal_stems.py`
**Purpose**: Filter and extract records containing vocal stems from large metadata JSONL files.

**Key Functions**:
- Identifies records with stems containing "vocal" or "vox" (case insensitive)
- Streams processing for memory efficiency on large files
- Progress tracking with tqdm

**Usage**:
```bash
python extract_vocal_stems.py --input-path /app2/suno/data/auk_v0/metas_v6_tr.jsonl --output-path ./vocal_stems.jsonl
```

**Filter Logic**: Records must have a "stems" dict with at least one key containing "vocal" or "vox"

---

### 3. `sample_vocal_stem_records.py`
**Purpose**: Sample and display individual vocal stem records with detailed formatting for manual inspection.

**Key Functions**:
- Supports forward/backward file traversal
- Configurable sampling with skip patterns
- Pretty-printed JSON output for human review
- Export to JSON files for further processing

**Usage**:
```bash
python sample_vocal_stem_records.py --max-samples 10 --output-file ./samples.json --reverse
```

**Sampling Options**: Random sampling, sequential processing, reverse traversal, skip patterns

---

### 4. `sample_vocal_stems.py`
**Purpose**: Create statistically representative samples from vocal stems files while preserving order.

**Key Functions**:
- Order-preserving uniform sampling across file length
- Ratio-based sampling (e.g., 0.1% of total records)
- Count-based sampling with uniform distribution
- Sequential sampling for consecutive records

**Usage**:
```bash
# Sample 100 records uniformly distributed
python sample_vocal_stems.py --count 100 --output-path ./sample_100.jsonl

# Sample 0.1% of records
python sample_vocal_stems.py --ratio 0.001 --output-path ./sample_0.1pct.jsonl
```

**Sampling Modes**: `--count`, `--ratio`, `--sequential`

---

### 5. `verify_vocal_stems.py`
**Purpose**: Quality assurance to verify that extracted vocal stem files contain valid vocal stems.

**Key Functions**:
- Validates presence of vocal-related stem keys
- Identifies and reports invalid records
- Shows examples of valid vocal stem structures
- Statistics on vocal stem types and distributions

**Usage**:
```bash
python verify_vocal_stems.py --input-path ./vocal_stems.jsonl --max-examples 5
```

**Verification**: Ensures all records have stems dict with vocal/vox-containing keys

---

### 6. `caption_vocal_stems.py`
**Purpose**: Advanced AI-powered captioning system for vocal stem analysis using Google Gemini models with intelligent retry logic, keyword accumulation, and probabilistic model selection.

**Key Capabilities**:
- **Parallel Processing**: Configurable multi-worker execution for high throughput
- **Dual Model Support**: Gemini 2.5 Flash (fast) and Pro (high-quality) models
- **Smart Model Selection**: Probabilistic selection (80% Flash, 20% Pro) or forced model use
- **Keyword Accumulation**: Automatically accumulates keywords across multiple API calls to meet minimum count requirements
- **Robust Validation**: Comma-separated format validation with punctuation filtering (allows trailing dots)
- **Incremental Saving**: Periodic batch saves with order preservation for long-running jobs
- **Audio Processing**: Automatic OPUS→MP3 conversion for Gemini compatibility
- **Intelligent Retry**: Exponential backoff with different strategies for various error types

**Core Arguments**:
```bash
python caption_vocal_stems.py \
  --input-jsonl /path/to/vocal_stems.jsonl \
  --output-jsonl /path/to/output.jsonl \
  --prompts-file /path/to/prompts.json \
  --num-workers 4 \
  --save-every 20
```

**Advanced Configuration**:
```bash
# High-quality processing with forced Pro model
python caption_vocal_stems.py \
  --input-jsonl ./vocals.jsonl \
  --output-jsonl ./captions_pro.jsonl \
  --prompts-file ./prompts.json \
  --model gemini-2.5-pro \
  --force-model \
  --min-keywords 30 \
  --num-workers 2 \
  --save-every 10

# Fast processing with probabilistic model selection
python caption_vocal_stems.py \
  --input-jsonl ./vocals.jsonl \
  --output-jsonl ./captions_mixed.jsonl \
  --prompts-file ./prompts.json \
  --min-keywords 25 \
  --num-workers 6 \
  --save-every 50

# Process all stems (not just vocals)
python caption_vocal_stems.py \
  --input-jsonl ./all_stems.jsonl \
  --output-jsonl ./all_captions.jsonl \
  --prompts-file ./prompts.json \
  --all-stems \
  --limit 100
```

**Command-Line Options**:
- `--input-jsonl`: Input JSONL file with vocal stem records
- `--output-jsonl`: Output JSONL file for captioned results
- `--prompts-file`: JSON file containing prompt configurations
- `--num-workers`: Parallel workers (default: 4, optimal: 2-8)
- `--model`: Base model (`gemini-2.5-flash` or `gemini-2.5-pro`)
- `--force-model`: Disable probabilistic selection, use specified model only
- `--min-keywords`: Minimum keywords per caption (default: 25)
- `--save-every`: Save results every N records (default: 10)
- `--all-stems`: Process all stems instead of vocal-only (default: vocal-only)
- `--limit`: Process only first N records (for testing)
- `--tmpdir`: Custom temporary directory for audio processing

**Model Selection Strategy**:
- **Probabilistic Mode** (default): 80% Flash, 20% Pro selection per caption attempt
- **Forced Mode** (`--force-model`): Uses specified model exclusively
- **Per-Caption Selection**: Each retry can use a different model for better success rates

**Keyword Accumulation System**:
- **Target**: Configurable minimum keywords (default: 25)
- **Validation**: Comma-separated format with punctuation limits (<20% non-alphanumeric)
- **Tolerance**: Accepts trailing dots (e.g., "rock.", "R&B.", "pop.")
- **Accumulation**: Combines keywords from multiple API calls until minimum reached
- **Deduplication**: Prevents duplicate keywords across accumulation attempts

**Error Handling & Retry Logic**:
- **Validation Failures**: Retry with different models/prompts
- **Empty Responses**: Treat as retryable errors (common with Pro model)
- **Rate Limits**: Exponential backoff with longer delays
- **Partial Results**: Return accumulated keywords if minimum not reached
- **Maximum Attempts**: 6 total attempts (3 base + 3 accumulation)

**Performance Characteristics**:
- **Processing Rate**: ~0.1-0.2 records/minute (varies by stem count and model mix)
- **Success Rate**: 85-95% (improved with per-caption model selection)
- **Memory Usage**: Streaming JSONL processing for large datasets
- **Scalability**: Handles 100,000+ record datasets with incremental saving

**Output Structure**:
```json
{
  "id": "record_id",
  "stems": {...},
  "stems_captions": {
    "Vocals": [{
      "prompt_type": "voice_description_keywords",
      "caption": "Male, adult, American accent, deep, resonant, soulful, bluesy...",
      "error": null
    }],
    "Backing_Vocals": [...]
  }
}
```

**Monitoring & Logging**:
- **Progress Bars**: Real-time batch progress tracking
- **Structured Logging**: Timestamps, worker IDs, success/failure details
- **Error Categories**: Validation failures, API errors, empty responses
- **Model Usage**: Logs which model is selected for each attempt

**Best Practices**:
- **Worker Count**: Start with 2-4 workers, scale up if API limits allow
- **Save Frequency**: Use `--save-every 20-50` for large jobs to balance I/O
- **Testing**: Use `--limit 10` to validate configuration before full runs
- **Model Selection**: Use probabilistic mode for best cost/quality balance
- **Keyword Count**: 25-30 keywords provides good descriptive coverage

**Common Use Cases**:
```bash
# Quick test run
python caption_vocal_stems.py --input-jsonl sample.jsonl --output-jsonl test.jsonl --prompts-file prompts.json --limit 5

# Production run with balanced settings  
python caption_vocal_stems.py --input-jsonl vocals_10k.jsonl --output-jsonl captions.jsonl --prompts-file prompts.json --num-workers 4 --save-every 25

# High-quality run for critical data
python caption_vocal_stems.py --input-jsonl important_vocals.jsonl --output-jsonl quality_captions.jsonl --prompts-file prompts.json --model gemini-2.5-pro --force-model --min-keywords 30
```

---

### 7. `find_accent_records.py`
**Purpose**: Intelligent detection and extraction of records containing specific accents for accent recognition testing and evaluation.

**Key Capabilities**:
- **Multi-Accent Detection**: Searches for 16 different accent types (Scottish, Australian, German, Chinese, Japanese, Korean, etc.)
- **Smart Keyword Matching**: Uses region-specific terms and cultural markers for accurate detection
- **Multi-Field Search**: Analyzes text content, tags, artist info, and metadata for accent indicators
- **Configurable Filtering**: Vocal-stem-only option, per-accent limits, and target accent selection
- **Analysis Metadata**: Adds detailed accent analysis to records for evaluation purposes

**Core Usage**:
```bash
# Find all accent types in validation set
python find_accent_records.py \
  --input /app2/suno/data/auk_v0/metas_v6_val.jsonl \
  --output accent_test_records.jsonl \
  --max-per-accent 5

# Find specific accents for targeted testing
python find_accent_records.py \
  --input metas_v6_val.jsonl \
  --output scottish_aussie.jsonl \
  --accents scottish australian german \
  --max-per-accent 10
```

**Available Accents**: scottish, australian, german, bavarian, chinese, japanese, korean, irish, french, italian, spanish, russian, indian, british, southern_us, new_york

**Detection Strategy**:
- **Scottish**: "scottish", "glasgow", "edinburgh", "highland"
- **Australian**: "aussie", "sydney", "melbourne", "adelaide"
- **German**: "german", "deutsch", "berlin", "munich"
- **Chinese**: "chinese", "mandarin", "beijing", "shanghai"
- **Japanese**: "japanese", "tokyo", "osaka", "kyoto"
- **And many more region-specific keywords...**

**Output Structure**:
```json
{
  "id": "record_id",
  "stems": {...},
  "_accent_analysis": {
    "found_accents": {
      "scottish": [{
        "field": "tags",
        "matched_keywords": ["glasgow", "scottish"],
        "context": "Scottish folk song from Glasgow..."
      }]
    }
  }
}
```

---

### 8. `find_famous_artist_records.py` 
**Purpose**: Detection and extraction of records from famous artists for voice recognition testing and artist identification evaluation.

**Key Capabilities**:
- **70+ Famous Artists**: From Taylor Swift and Drake to The Beatles and Queen
- **Multi-Genre Coverage**: Pop, Hip-Hop, Rock, R&B spanning multiple decades
- **Comprehensive Keyword Matching**: Artist names, aliases, real names, and associated terms
- **Cross-Era Detection**: Modern artists (Billie Eilish, Post Malone) to legends (Elvis, Michael Jackson)
- **Collaboration Detection**: Identifies records mentioning multiple artists

**Core Usage**:
```bash
# Find all famous artists in dataset
python find_famous_artist_records.py \
  --input /app2/suno/data/auk_v0/metas_v6_val.jsonl \
  --output famous_artist_records.jsonl \
  --max-per-artist 3

# Find specific artists for targeted testing  
python find_famous_artist_records.py \
  --input metas_v6_val.jsonl \
  --output pop_stars.jsonl \
  --artists taylor_swift ariana_grande billie_eilish drake \
  --max-per-artist 5
```

**Artist Categories**:
- **Modern Pop**: taylor_swift, ariana_grande, billie_eilish, dua_lipa, olivia_rodrigo
- **Hip-Hop**: drake, kendrick_lamar, kanye_west, eminem, jay_z, nas, future, travis_scott
- **R&B/Soul**: beyonce, rihanna, bruno_mars, frank_ocean, sza, the_weeknd
- **Rock Legends**: the_beatles, queen, led_zeppelin, pink_floyd, nirvana, radiohead
- **Pop Icons**: michael_jackson, madonna, prince, elvis_presley, adele

**Detection Examples**:
- **Taylor Swift**: "taylor swift", "taylor", "swift", "tswift"  
- **Drake**: "drake", "drizzy", "champagne papi", "aubrey graham"
- **The Beatles**: "beatles", "john lennon", "paul mccartney", "george harrison"
- **Queen**: "queen", "freddie mercury", "brian may"

**Output Structure**:
```json
{
  "id": "record_id", 
  "stems": {...},
  "_artist_analysis": {
    "found_artists": {
      "taylor_swift": [{
        "field": "text",
        "matched_keywords": ["taylor swift", "swift"],
        "context": "Taylor Swift's storytelling approach..."
      }]
    }
  }
}
```

---

### 9. `jsonl_to_individual_files.py`
**Purpose**: Utility script to convert JSONL files into individual JSON files for easier inspection and analysis.

**Key Features**:
- **Flexible Naming**: Sequential numbering, custom prefixes, filesystem-safe names
- **Pretty Formatting**: Configurable JSON indentation for readability  
- **Progress Tracking**: Real-time progress bars for large datasets
- **Safe Processing**: Handles malformed JSON gracefully
- **Overwrite Protection**: Prevents accidental file replacement

**Usage**:
```bash
# Basic conversion with sequential numbering
python jsonl_to_individual_files.py \
  --input records.jsonl \
  --output individual_files/

# Custom formatting and naming
python jsonl_to_individual_files.py \
  --input large_dataset.jsonl \
  --output pretty_files/ \
  --prefix experiment_1 \
  --indent 4 \
  --limit 100
```

---

### 10. `add_accents_to_filenames.py` & `add_artists_to_filenames.py`
**Purpose**: Filename enhancement utilities that add detected accent/artist information to individual JSON filenames for easy identification.

**Key Features**:
- **Smart Labeling**: Extracts accent/artist info from analysis metadata
- **Multi-Label Support**: Handles records with multiple accents/artists (e.g., "Scottish+Irish")
- **Filesystem Safe**: Converts names to safe filename formats
- **Copy/Move Options**: Flexible file handling with dry-run support

**Usage**:
```bash
# Add accent labels to filenames
python add_accents_to_filenames.py \
  --input-dir /path/to/accent_files/ \
  --output-dir /path/to/labeled_files/ \
  --copy

# Add artist labels to filenames  
python add_artists_to_filenames.py \
  --input-dir /path/to/artist_files/ \
  --output-dir /path/to/labeled_files/ \
  --copy
```

**Output Examples**:
```
accent_001_Scottish_record_id.json
accent_002_German+Bavarian_record_id.json  
artist_001_Taylor-Swift_record_id.json
artist_002_Kanye-West+Eminem_record_id.json
```

## Data Structures

### Input Format (v6 Metadata)
```json
{
  "id": "record_id",
  "stems": {
    "Vocals": "/path/to/vocals.opus",
    "Backing_Vocals": "/path/to/backing.opus",
    "Drums": "/path/to/drums.opus"
  },
  "text": "lyrics...",
  "tags": ["genre", "style"]
}
```

### Output Format (with Captions)
```json
{
  "id": "record_id",
  "stems": {...},
  "stems_captions": {
    "Vocals": [{
      "prompt_type": "voice_description_keywords",
      "caption": "Male, adult, smooth, clear, soulful...",
      "error": null
    }]
  }
}
```

## Environment Requirements

- **Python**: 3.10+
- **Conda Environment**: `suno_3`
- **Key Dependencies**: 
  - `google-genai` (for captioning)
  - `tqdm` (progress bars)
  - `matplotlib` (visualizations)
  - `file-read-backwards` (reverse file traversal)

## Common Usage Patterns

### Full Pipeline Example
```bash
# 1. Analyze the dataset
python analyze_v6_data.py --data-path /app2/suno/data/auk_v0/metas_v6_tr.jsonl

# 2. Extract vocal records
python extract_vocal_stems.py --input-path /app2/suno/data/auk_v0/metas_v6_tr.jsonl

# 3. Sample for testing
python sample_vocal_stems.py --count 100 --output-path ./test_sample.jsonl

# 4. Verify quality
python verify_vocal_stems.py --input-path ./test_sample.jsonl

# 5. Generate captions
python caption_vocal_stems.py --input-jsonl ./test_sample.jsonl --output-jsonl ./captioned.jsonl
```

### Performance Considerations

- **Large Files**: Use streaming/sampling for files >1M records
- **Parallel Processing**: `caption_vocal_stems.py` supports 2-8 workers optimally
- **Memory Usage**: Most scripts use streaming to handle large files efficiently
- **API Limits**: Gemini has rate limits; adjust workers accordingly

## Error Handling

All scripts include:
- Progress bars for long-running operations
- Graceful handling of malformed JSON records
- Retry logic for API calls (captioning)
- Detailed error reporting and validation

## Development Notes

- Scripts are designed for the Suno AI audio processing pipeline
- All scripts support `--help` for detailed usage information  
- Designed to work with OPUS audio format prevalent in Suno datasets
- Integration with Weights & Biases for experiment tracking (where applicable)