Deep Learning for Speech Entrainment Detection - A Practical Guide

import torch
import torch.nn as nn
import librosa
import numpy as np
from typing import Tuple, List

class SpeechEntrainmentDetector(nn.Module):
    def __init__(self,
                 acoustic_dim: int = 128,
                 prosodic_dim: int = 64,
                 temporal_dim: int = 32,
                 hidden_dim: int = 256):
        super().__init__()

        # Acoustic feature encoder
        self.acoustic_encoder = nn.Sequential(
            nn.Linear(acoustic_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim, hidden_dim // 2)
        )

        # Prosodic feature encoder
        self.prosodic_encoder = nn.Sequential(
            nn.Linear(prosodic_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(hidden_dim // 2, hidden_dim // 4)
        )

        # Temporal dynamics module
        self.temporal_lstm = nn.LSTM(
            input_size=temporal_dim,
            hidden_size=hidden_dim // 4,
            num_layers=2,
            dropout=0.3,
            bidirectional=True
        )

        # Cross-modal attention
        self.attention = nn.MultiheadAttention(
            embed_dim=hidden_dim,
            num_heads=8,
            dropout=0.1
        )

        # Entrainment classifier
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_dim // 2, 1),
            nn.Sigmoid()
        )

    def forward(self, acoustic_features, prosodic_features, temporal_features):
        # Encode different feature types
        acoustic_encoded = self.acoustic_encoder(acoustic_features)
        prosodic_encoded = self.prosodic_encoder(prosodic_features)

        # Process temporal dynamics
        temporal_encoded, _ = self.temporal_lstm(temporal_features)
        temporal_encoded = temporal_encoded.mean(dim=1)  # Global temporal representation

        # Combine features
        combined_features = torch.cat([
            acoustic_encoded,
            prosodic_encoded,
            temporal_encoded
        ], dim=-1)

        # Apply attention mechanism
        attended_features, _ = self.attention(
            combined_features.unsqueeze(0),
            combined_features.unsqueeze(0),
            combined_features.unsqueeze(0)
        )

        # Classify entrainment
        entrainment_score = self.classifier(attended_features.squeeze(0))

        return entrainment_score

def extract_acoustic_features(audio_path: str, sr: int = 16000) -> np.ndarray:
    """
    Extract comprehensive acoustic features from audio file.
    """
    # Load audio
    y, _ = librosa.load(audio_path, sr=sr)

    # Extract features
    features = []

    # MFCCs (13 coefficients)
    mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13)
    features.append(mfccs.mean(axis=1))
    features.append(mfccs.std(axis=1))

    # Spectral features
    spectral_centroids = librosa.feature.spectral_centroid(y=y, sr=sr)
    spectral_rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr)
    spectral_bandwidth = librosa.feature.spectral_bandwidth(y=y, sr=sr)

    features.extend([
        spectral_centroids.mean(),
        spectral_rolloff.mean(),
        spectral_bandwidth.mean()
    ])

    # Zero crossing rate
    zcr = librosa.feature.zero_crossing_rate(y)
    features.append(zcr.mean())

    # Chromagram
    chroma = librosa.feature.chroma_stft(y=y, sr=sr)
    features.extend(chroma.mean(axis=1))

    return np.array(features)

def extract_prosodic_features(audio_path: str) -> np.ndarray:
    """
    Extract prosodic features using Praat integration.
    """
    try:
        import parselmouth
        from parselmouth.praat import call
    except ImportError:
        raise ImportError("Please install parselmouth: pip install praat-parselmouth")

    sound = parselmouth.Sound(audio_path)

    # Fundamental frequency analysis
    f0 = sound.to_pitch_ac(time_step=0.01, pitch_floor=75, pitch_ceiling=300)
    f0_values = f0.selected_array['frequency']
    f0_values = f0_values[f0_values > 0]  # Remove unvoiced frames

    # Intensity analysis
    intensity = sound.to_intensity(time_step=0.01)
    intensity_values = intensity.values[0]

    # Calculate prosodic statistics
    features = []

    if len(f0_values) > 0:
        features.extend([
            np.mean(f0_values),
            np.std(f0_values),
            np.max(f0_values) - np.min(f0_values),  # F0 range
            np.percentile(f0_values, 75) - np.percentile(f0_values, 25)  # IQR
        ])
    else:
        features.extend([0, 0, 0, 0])

    features.extend([
        np.mean(intensity_values),
        np.std(intensity_values),
        np.max(intensity_values) - np.min(intensity_values)
    ])

    return np.array(features)

def train_model(model, train_loader, val_loader, epochs=50):
    """
    Training loop for the speech entrainment detector.
    """
    criterion = nn.BCELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, patience=5, factor=0.5
    )

    best_val_loss = float('inf')

    for epoch in range(epochs):
        # Training phase
        model.train()
        train_loss = 0.0

        for batch_idx, (acoustic, prosodic, temporal, labels) in enumerate(train_loader):
            optimizer.zero_grad()

            predictions = model(acoustic, prosodic, temporal)
            loss = criterion(predictions.squeeze(), labels.float())

            loss.backward()
            optimizer.step()

            train_loss += loss.item()

        # Validation phase
        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():
            for acoustic, prosodic, temporal, labels in val_loader:
                predictions = model(acoustic, prosodic, temporal)
                loss = criterion(predictions.squeeze(), labels.float())
                val_loss += loss.item()

                # Calculate accuracy
                predicted = (predictions.squeeze() > 0.5).long()
                total += labels.size(0)
                correct += (predicted == labels).sum().item()

        # Learning rate scheduling
        scheduler.step(val_loss)

        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), 'best_entrainment_model.pth')

        print(f'Epoch {epoch+1}/{epochs}:')
        print(f'  Train Loss: {train_loss/len(train_loader):.4f}')
        print(f'  Val Loss: {val_loss/len(val_loader):.4f}')
        print(f'  Val Accuracy: {100*correct/total:.2f}%')

def enhance_conversational_ai(user_speech, ai_response_candidates):
    """
    Select AI responses that promote natural entrainment patterns.
    """
    entrainment_scores = []

    for candidate in ai_response_candidates:
        # Extract features from user speech and candidate response
        user_features = extract_all_features(user_speech)
        candidate_features = extract_all_features(candidate)

        # Predict entrainment compatibility
        score = model.predict_entrainment(user_features, candidate_features)
        entrainment_scores.append(score)

    # Select response with optimal entrainment score
    best_response_idx = np.argmax(entrainment_scores)
    return ai_response_candidates[best_response_idx]

def assess_communication_therapy_progress(patient_sessions):
    """
    Track patient progress in developing natural entrainment patterns.
    """
    progress_scores = []

    for session in patient_sessions:
        therapist_speech = session['therapist']
        patient_speech = session['patient']

        entrainment_score = detect_entrainment(therapist_speech, patient_speech)
        progress_scores.append(entrainment_score)

    return {
        'overall_progress': np.mean(progress_scores),
        'trend': np.polyfit(range(len(progress_scores)), progress_scores, 1)[0],
        'individual_scores': progress_scores
    }