datasets/analyze_dataset.py

#!/usr/bin/env python3
"""
Deep Dataset Analyzer - Comprehensive analysis of generated dataset.

Analyzes:
1. Structure and field presence
2. Whisper features (shape, statistics, quality)
3. SNAC tokens (offsets, distribution, frame integrity)
4. Text content (questions, answers, tokenization)
5. Word alignments (timing, coverage, token mapping)
6. Cross-field consistency
7. Training readiness checks
"""

import argparse
import sys
from pathlib import Path
from collections import defaultdict, Counter
import json

import torch
import numpy as np

# SNAC constants
SNAC_BASE = 128266
SNAC_LAYERS = 7
SNAC_VOCAB_PER_LAYER = 4096
WHISPER_DIM = 1280
EOS_TOKEN = 128009


def print_header(title):
    print(f"\n{'='*70}")
    print(f" {title}")
    print('='*70)


def print_subheader(title):
    print(f"\n{'-'*50}")
    print(f" {title}")
    print('-'*50)


def analyze_whisper_features(dataset):
    """Deep analysis of Whisper features."""
    print_header("WHISPER FEATURES ANALYSIS")

    lengths = []
    dims = []
    min_vals = []
    max_vals = []
    mean_vals = []
    std_vals = []
    has_nan = 0
    has_inf = 0
    dtypes = Counter()

    for i, sample in enumerate(dataset):
        if "whisper_features" not in sample:
            continue

        wf = sample["whisper_features"]

        if not isinstance(wf, torch.Tensor):
            print(f"  [WARN] Sample {i}: whisper_features is {type(wf).__name__}, not Tensor")
            continue

        dtypes[str(wf.dtype)] += 1
        lengths.append(wf.shape[0])

        if wf.dim() >= 2:
            dims.append(wf.shape[1])

        if torch.isnan(wf).any():
            has_nan += 1
        if torch.isinf(wf).any():
            has_inf += 1

        min_vals.append(wf.min().item())
        max_vals.append(wf.max().item())
        mean_vals.append(wf.mean().item())
        std_vals.append(wf.std().item())

    if not lengths:
        print("  No whisper_features found!")
        return

    lengths = np.array(lengths)

    print(f"\n  Samples with whisper_features: {len(lengths)}/{len(dataset)}")
    print(f"  Data types: {dict(dtypes)}")

    print_subheader("Sequence Length")
    print(f"  Min: {lengths.min()}, Max: {lengths.max()}")
    print(f"  Mean: {lengths.mean():.1f}, Std: {lengths.std():.1f}")
    print(f"  Median: {np.median(lengths):.1f}")

    # Length distribution
    percentiles = [10, 25, 50, 75, 90, 95, 99]
    print(f"  Percentiles: {dict(zip(percentiles, [int(np.percentile(lengths, p)) for p in percentiles]))}")

    if dims:
        dims = np.array(dims)
        unique_dims = np.unique(dims)
        print(f"\n  Feature dimension: {unique_dims}")
        if len(unique_dims) == 1 and unique_dims[0] == WHISPER_DIM:
            print(f"  ✅ Correct Whisper dimension ({WHISPER_DIM})")
        else:
            print(f"  ⚠️ Expected dimension {WHISPER_DIM}")

    print_subheader("Value Statistics")
    print(f"  Min value: {np.min(min_vals):.4f}")
    print(f"  Max value: {np.max(max_vals):.4f}")
    print(f"  Mean of means: {np.mean(mean_vals):.4f}")
    print(f"  Mean of stds: {np.mean(std_vals):.4f}")

    print_subheader("Quality Checks")
    print(f"  Samples with NaN: {has_nan} {'⚠️' if has_nan > 0 else '✅'}")
    print(f"  Samples with Inf: {has_inf} {'⚠️' if has_inf > 0 else '✅'}")

    # Check for constant/zero features
    zero_features = sum(1 for s in std_vals if s < 0.001)
    print(f"  Samples with near-zero std: {zero_features} {'⚠️' if zero_features > 0 else '✅'}")


def analyze_snac_tokens(dataset):
    """Deep analysis of SNAC tokens."""
    print_header("SNAC TOKENS ANALYSIS")

    lengths = []
    frame_counts = []
    token_values = []
    offset_distribution = defaultdict(int)
    raw_token_distribution = defaultdict(int)
    incomplete_frames = 0
    has_offset = 0
    no_offset = 0
    dtypes = Counter()

    for i, sample in enumerate(dataset):
        if "snac_tokens" not in sample:
            continue

        st = sample["snac_tokens"]

        if not isinstance(st, torch.Tensor):
            print(f"  [WARN] Sample {i}: snac_tokens is {type(st).__name__}, not Tensor")
            continue

        dtypes[str(st.dtype)] += 1
        length = len(st)
        lengths.append(length)
        frame_counts.append(length // SNAC_LAYERS)

        if length % SNAC_LAYERS != 0:
            incomplete_frames += 1

        # Analyze token values
        for j, tok in enumerate(st.tolist()):
            token_values.append(tok)

            if tok >= SNAC_BASE:
                has_offset += 1
                # Determine which position offset
                offset_idx = (tok - SNAC_BASE) // SNAC_VOCAB_PER_LAYER
                offset_distribution[offset_idx] += 1
                # Get raw token value
                raw_tok = (tok - SNAC_BASE) % SNAC_VOCAB_PER_LAYER
                raw_token_distribution[raw_tok] += 1
            else:
                no_offset += 1

    if not lengths:
        print("  No snac_tokens found!")
        return

    lengths = np.array(lengths)
    frame_counts = np.array(frame_counts)
    token_values = np.array(token_values)

    print(f"\n  Samples with snac_tokens: {len(lengths)}/{len(dataset)}")
    print(f"  Data types: {dict(dtypes)}")

    print_subheader("Token Counts")
    print(f"  Total tokens: {len(token_values):,}")
    print(f"  Min per sample: {lengths.min()}, Max: {lengths.max()}")
    print(f"  Mean: {lengths.mean():.1f}, Std: {lengths.std():.1f}")

    print_subheader("Frame Analysis (7 tokens per frame)")
    print(f"  Min frames: {frame_counts.min()}, Max: {frame_counts.max()}")
    print(f"  Mean frames: {frame_counts.mean():.1f}")
    print(f"  Incomplete frames (not multiple of 7): {incomplete_frames} {'⚠️' if incomplete_frames > 0 else '✅'}")

    # Duration estimation (75 frames/second for SNAC 24kHz)
    durations = frame_counts / 75.0
    print(f"\n  Audio duration (estimated):")
    print(f"    Min: {durations.min():.2f}s, Max: {durations.max():.2f}s")
    print(f"    Mean: {durations.mean():.2f}s, Total: {durations.sum()/60:.1f} min")

    print_subheader("Token Offset Analysis")
    print(f"  Tokens with offset (>= {SNAC_BASE}): {has_offset:,} ({100*has_offset/len(token_values):.1f}%)")
    print(f"  Tokens without offset: {no_offset:,} ({100*no_offset/len(token_values):.1f}%)")

    if has_offset > 0:
        print(f"\n  Offset distribution by position (0-6):")
        for pos in range(SNAC_LAYERS):
            count = offset_distribution.get(pos, 0)
            expected_pct = 100 / SNAC_LAYERS
            actual_pct = 100 * count / has_offset if has_offset > 0 else 0
            status = "✅" if abs(actual_pct - expected_pct) < 5 else "⚠️"
            print(f"    Position {pos}: {count:,} ({actual_pct:.1f}%) {status}")

    print_subheader("Token Value Range")
    print(f"  Min token: {token_values.min()}")
    print(f"  Max token: {token_values.max()}")

    if has_offset > 0:
        expected_min = SNAC_BASE
        expected_max = SNAC_BASE + (SNAC_LAYERS * SNAC_VOCAB_PER_LAYER) - 1
        print(f"  Expected range: [{expected_min}, {expected_max}]")
        if token_values.min() >= expected_min and token_values.max() <= expected_max:
            print(f"  ✅ Token range valid")
        else:
            print(f"  ⚠️ Some tokens outside expected range")


def analyze_text_content(dataset):
    """Analyze text fields (question, answer, text_tokens)."""
    print_header("TEXT CONTENT ANALYSIS")

    questions = []
    answers = []
    text_token_lengths = []
    has_text = 0
    has_answer = 0
    has_text_tokens = 0

    for sample in dataset:
        if "text" in sample and sample["text"]:
            has_text += 1
            questions.append(sample["text"])

        if "answer" in sample and sample["answer"]:
            has_answer += 1
            answers.append(sample["answer"])

        if "text_tokens" in sample:
            has_text_tokens += 1
            tt = sample["text_tokens"]
            if isinstance(tt, torch.Tensor):
                text_token_lengths.append(len(tt))
            elif isinstance(tt, list):
                text_token_lengths.append(len(tt))

    print(f"\n  Field presence:")
    print(f"    text (question): {has_text}/{len(dataset)} ({100*has_text/len(dataset):.1f}%)")
    print(f"    answer: {has_answer}/{len(dataset)} ({100*has_answer/len(dataset):.1f}%)")
    print(f"    text_tokens: {has_text_tokens}/{len(dataset)} ({100*has_text_tokens/len(dataset):.1f}%)")

    if questions:
        print_subheader("Questions Analysis")
        q_lengths = [len(q) for q in questions]
        q_words = [len(q.split()) for q in questions]
        print(f"  Character length: min={min(q_lengths)}, max={max(q_lengths)}, mean={np.mean(q_lengths):.1f}")
        print(f"  Word count: min={min(q_words)}, max={max(q_words)}, mean={np.mean(q_words):.1f}")

        # Sample questions
        print(f"\n  Sample questions:")
        for q in questions[:3]:
            print(f"    - {q[:80]}{'...' if len(q) > 80 else ''}")

    if answers:
        print_subheader("Answers Analysis")
        a_lengths = [len(a) for a in answers]
        a_words = [len(a.split()) for a in answers]
        print(f"  Character length: min={min(a_lengths)}, max={max(a_lengths)}, mean={np.mean(a_lengths):.1f}")
        print(f"  Word count: min={min(a_words)}, max={max(a_words)}, mean={np.mean(a_words):.1f}")

        # Sample answers
        print(f"\n  Sample answers:")
        for a in answers[:3]:
            print(f"    - {a[:100]}{'...' if len(a) > 100 else ''}")

    if text_token_lengths:
        print_subheader("Pre-tokenized Answer Tokens")
        tl = np.array(text_token_lengths)
        print(f"  Token count: min={tl.min()}, max={tl.max()}, mean={tl.mean():.1f}")


def analyze_word_alignments(dataset):
    """Analyze word alignments for IST-LM interleaving."""
    print_header("WORD ALIGNMENTS ANALYSIS")

    has_alignments = 0
    word_counts = []
    frame_coverages = []
    has_tokens = 0
    token_counts = []
    timing_issues = 0

    for i, sample in enumerate(dataset):
        if "word_alignments" not in sample or not sample["word_alignments"]:
            continue

        has_alignments += 1
        wa = sample["word_alignments"]
        word_counts.append(len(wa))

        # Check alignment quality
        last_end = 0
        sample_has_tokens = True
        sample_token_count = 0

        for j, alignment in enumerate(wa):
            # Check required fields
            if 'word' not in alignment or 'start_frame' not in alignment or 'end_frame' not in alignment:
                continue

            start = alignment['start_frame']
            end = alignment['end_frame']

            # Check timing consistency
            if start < last_end - 1:  # Allow small overlap
                timing_issues += 1
            last_end = end

            # Check token presence
            if 'tokens' in alignment and alignment['tokens']:
                sample_token_count += len(alignment['tokens'])
            else:
                sample_has_tokens = False

        if sample_has_tokens and sample_token_count > 0:
            has_tokens += 1
            token_counts.append(sample_token_count)

        # Calculate frame coverage
        if "snac_tokens" in sample:
            snac_len = len(sample["snac_tokens"])
            total_frames = snac_len // SNAC_LAYERS
            if total_frames > 0 and wa:
                covered_frames = wa[-1]['end_frame'] if wa[-1]['end_frame'] else 0
                frame_coverages.append(min(100, 100 * covered_frames / total_frames))

    print(f"\n  Samples with word_alignments: {has_alignments}/{len(dataset)} ({100*has_alignments/len(dataset):.1f}%)")

    if not has_alignments:
        print("  No word alignments found!")
        return

    print(f"  Samples with pre-computed tokens: {has_tokens}/{has_alignments} ({100*has_tokens/has_alignments:.1f}%)")

    word_counts = np.array(word_counts)
    print_subheader("Word Count per Sample")
    print(f"  Min: {word_counts.min()}, Max: {word_counts.max()}, Mean: {word_counts.mean():.1f}")

    if token_counts:
        token_counts = np.array(token_counts)
        print_subheader("Token Count per Sample (from alignments)")
        print(f"  Min: {token_counts.min()}, Max: {token_counts.max()}, Mean: {token_counts.mean():.1f}")

    if frame_coverages:
        fc = np.array(frame_coverages)
        print_subheader("Frame Coverage")
        print(f"  Min: {fc.min():.1f}%, Max: {fc.max():.1f}%, Mean: {fc.mean():.1f}%")

    print_subheader("Quality Checks")
    print(f"  Timing issues (overlaps): {timing_issues} {'⚠️' if timing_issues > 0 else '✅'}")


def analyze_cross_field_consistency(dataset):
    """Check consistency between related fields."""
    print_header("CROSS-FIELD CONSISTENCY")

    issues = defaultdict(int)

    for i, sample in enumerate(dataset):
        # Check answer vs text_tokens consistency
        if "answer" in sample and "text_tokens" in sample:
            answer = sample["answer"]
            text_tokens = sample["text_tokens"]
            if isinstance(text_tokens, torch.Tensor):
                token_len = len(text_tokens)
            else:
                token_len = len(text_tokens) if text_tokens else 0

            # Rough check: tokens should be roughly proportional to text length
            if answer and token_len > 0:
                chars_per_token = len(answer) / token_len
                if chars_per_token < 1 or chars_per_token > 10:
                    issues["answer_token_mismatch"] += 1

        # Check word_alignments vs answer consistency
        if "word_alignments" in sample and "answer" in sample:
            wa = sample["word_alignments"]
            answer = sample["answer"]
            if wa and answer:
                wa_words = len(wa)
                answer_words = len(answer.split())
                if abs(wa_words - answer_words) > answer_words * 0.3:  # 30% tolerance
                    issues["alignment_word_mismatch"] += 1

        # Check snac_tokens frame count vs alignments
        if "snac_tokens" in sample and "word_alignments" in sample:
            snac = sample["snac_tokens"]
            wa = sample["word_alignments"]
            if isinstance(snac, torch.Tensor) and wa:
                snac_frames = len(snac) // SNAC_LAYERS
                last_frame = max((a.get('end_frame', 0) for a in wa), default=0)
                if last_frame > snac_frames * 1.1:  # 10% tolerance
                    issues["frame_overflow"] += 1

    print(f"\n  Consistency checks:")
    print(f"    Answer/token length mismatch: {issues['answer_token_mismatch']} {'⚠️' if issues['answer_token_mismatch'] > 0 else '✅'}")
    print(f"    Alignment/answer word mismatch: {issues['alignment_word_mismatch']} {'⚠️' if issues['alignment_word_mismatch'] > 0 else '✅'}")
    print(f"    Frame overflow in alignments: {issues['frame_overflow']} {'⚠️' if issues['frame_overflow'] > 0 else '✅'}")


def analyze_training_readiness(dataset):
    """Check if dataset is ready for training."""
    print_header("TRAINING READINESS CHECK")

    required_fields = ["whisper_features", "snac_tokens"]
    optional_fields = ["text", "answer", "text_tokens", "word_alignments"]

    field_presence = defaultdict(int)
    complete_samples = 0

    issues = []

    for sample in dataset:
        has_required = all(f in sample for f in required_fields)
        if has_required:
            complete_samples += 1

        for field in required_fields + optional_fields:
            if field in sample and sample[field] is not None:
                if isinstance(sample[field], (torch.Tensor, list, str)):
                    if len(sample[field]) > 0:
                        field_presence[field] += 1
                else:
                    field_presence[field] += 1

    print(f"\n  Total samples: {len(dataset)}")
    print(f"  Complete samples (with required fields): {complete_samples} ({100*complete_samples/len(dataset):.1f}%)")

    print(f"\n  Field presence:")
    for field in required_fields:
        count = field_presence[field]
        status = "✅" if count == len(dataset) else "❌"
        print(f"    {field}: {count}/{len(dataset)} {status} (required)")

    for field in optional_fields:
        count = field_presence[field]
        status = "✅" if count == len(dataset) else "⚠️"
        print(f"    {field}: {count}/{len(dataset)} {status} (optional)")

    # Overall assessment
    print_subheader("Overall Assessment")

    ready = True

    if complete_samples < len(dataset):
        print(f"  ❌ {len(dataset) - complete_samples} samples missing required fields")
        ready = False
    else:
        print(f"  ✅ All samples have required fields")

    if field_presence.get("text_tokens", 0) == len(dataset):
        print(f"  ✅ Pre-tokenized text available (faster training)")
    else:
        print(f"  ⚠️ text_tokens missing in some samples (will tokenize on-the-fly)")

    if field_presence.get("word_alignments", 0) == len(dataset):
        print(f"  ✅ Word alignments available (semantic interleaving)")
    else:
        print(f"  ⚠️ word_alignments missing (will use positional interleaving)")

    if ready:
        print(f"\n  🎉 DATASET READY FOR TRAINING!")
    else:
        print(f"\n  ⚠️ Dataset has issues that need to be fixed")

    return ready


def print_sample_inspection(dataset, num_samples=3):
    """Print detailed inspection of sample items."""
    print_header(f"SAMPLE INSPECTION (first {num_samples} samples)")

    for i in range(min(num_samples, len(dataset))):
        sample = dataset[i]
        print(f"\n  Sample {i}:")

        for key, value in sample.items():
            if isinstance(value, torch.Tensor):
                print(f"    {key}: Tensor {list(value.shape)} {value.dtype}")
                if key == "snac_tokens" and len(value) > 0:
                    print(f"      First 7 tokens: {value[:7].tolist()}")
                    print(f"      Last 7 tokens: {value[-7:].tolist()}")
            elif isinstance(value, str):
                preview = value[:60] + "..." if len(value) > 60 else value
                print(f"    {key}: '{preview}'")
            elif isinstance(value, list):
                print(f"    {key}: list[{len(value)}]")
                if key == "word_alignments" and len(value) > 0:
                    print(f"      First: {value[0]}")
                    if len(value) > 1:
                        print(f"      Last: {value[-1]}")
            else:
                print(f"    {key}: {type(value).__name__}")


def main():
    parser = argparse.ArgumentParser(description="Deep dataset analysis")
    parser.add_argument("--path", type=str, required=True, help="Path to dataset .pt file")
    parser.add_argument("--max-samples", type=int, default=None, help="Max samples to analyze")
    parser.add_argument("--quick", action="store_true", help="Quick analysis (skip detailed checks)")
    args = parser.parse_args()

    print("\n" + "="*70)
    print(" DEEP DATASET ANALYZER")
    print("="*70)
    print(f"\n  File: {args.path}")

    # Check file exists
    if not Path(args.path).exists():
        print(f"\n  ❌ File not found: {args.path}")
        sys.exit(1)

    # Load dataset
    print(f"  Loading dataset...")
    try:
        dataset = torch.load(args.path, map_location="cpu", weights_only=False)
    except Exception as e:
        print(f"\n  ❌ Failed to load: {e}")
        sys.exit(1)

    if not isinstance(dataset, list):
        print(f"\n  ❌ Dataset should be list, got {type(dataset).__name__}")
        sys.exit(1)

    print(f"  Loaded {len(dataset):,} samples")

    # Limit samples if requested
    if args.max_samples and args.max_samples < len(dataset):
        print(f"  Analyzing first {args.max_samples} samples")
        dataset = dataset[:args.max_samples]

    # File size
    file_size = Path(args.path).stat().st_size / (1024**3)
    print(f"  File size: {file_size:.2f} GB")
    print(f"  Avg per sample: {file_size*1024/len(dataset):.2f} MB")

    # Run analyses
    analyze_whisper_features(dataset)
    analyze_snac_tokens(dataset)
    analyze_text_content(dataset)

    if not args.quick:
        analyze_word_alignments(dataset)
        analyze_cross_field_consistency(dataset)

    ready = analyze_training_readiness(dataset)
    print_sample_inspection(dataset)

    print("\n" + "="*70)
    print(" ANALYSIS COMPLETE")
    print("="*70 + "\n")

    sys.exit(0 if ready else 1)


if __name__ == "__main__":
    main()