meeting-assistant/MeetingAssistant/Speakers/PyannoteSpeakerIdentityMatchValidator.cs

using MeetingAssistant.Transcription;
using Microsoft.Extensions.Options;
using NAudio.Wave;

namespace MeetingAssistant.Speakers;

public sealed class PyannoteSpeakerIdentityMatchValidator : ISpeakerIdentityMatchValidator
{
    private readonly PyannoteTranscriptFinalizer finalizer;
    private readonly SpeakerIdentityPyannoteValidationOptions options;
    private readonly ILogger<PyannoteSpeakerIdentityMatchValidator> logger;

    public PyannoteSpeakerIdentityMatchValidator(
        PyannoteTranscriptFinalizer finalizer,
        IOptions<MeetingAssistantOptions> options,
        ILogger<PyannoteSpeakerIdentityMatchValidator> logger)
    {
        this.finalizer = finalizer;
        this.options = options.Value.SpeakerIdentification.PyannoteValidation;
        this.logger = logger;
    }

    public async Task<bool> ValidateSampleAsync(
        byte[] wavBytes,
        CancellationToken cancellationToken)
    {
        if (!options.Enabled)
        {
            return true;
        }

        var segments = await DiarizeBytesAsync(wavBytes, cancellationToken);
        var result = AnalyzeSingleSpeakerCoverage(segments);
        if (!result.Accepted)
        {
            logger.LogInformation(
                "pyannote rejected speaker identity sample because it did not contain one dominant speaker: segments {SegmentCount}, speakers {SpeakerCount}, dominant speaker {DominantSpeaker}, coverage {Coverage:P2}, required coverage {RequiredCoverage:P2}, duration {Duration}",
                segments.Count,
                result.SpeakerCount,
                result.DominantSpeaker,
                result.Coverage,
                Math.Clamp(options.MinimumSingleSpeakerCoverage, 0, 1),
                result.TotalDuration);
            return false;
        }

        logger.LogInformation(
            "pyannote accepted speaker identity sample: segments {SegmentCount}, speakers {SpeakerCount}, dominant speaker {DominantSpeaker}, coverage {Coverage:P2}, duration {Duration}, sample bytes {SampleBytes}",
            segments.Count,
            result.SpeakerCount,
            result.DominantSpeaker,
            result.Coverage,
            result.TotalDuration,
            wavBytes.Length);
        return true;
    }

    public async Task<bool> ValidateMatchAsync(
        SpeakerIdentityMatchValidationRequest request,
        CancellationToken cancellationToken)
    {
        if (!options.Enabled)
        {
            return true;
        }

        var tempPath = Path.Combine(
            Path.GetTempPath(),
            "meeting-assistant-speaker-match-pyannote",
            $"{Guid.NewGuid():N}.wav");
        try
        {
            Directory.CreateDirectory(Path.GetDirectoryName(tempPath)!);
            var layout = WriteCompositeWav(
                tempPath,
                request.KnownSnippets,
                request.UnknownSnippet,
                TimeSpan.FromSeconds(1));
            if (layout.UnknownSegment is null || layout.KnownSegments.Count == 0)
            {
                logger.LogInformation(
                    "pyannote rejected speaker identity match for {DiarizedSpeaker} and identity {IdentityId} because no readable composite snippets were available",
                    request.DiarizedSpeaker,
                    request.IdentityId);
                return false;
            }

            var segments = await DiarizePathAsync(tempPath, cancellationToken);
            logger.LogInformation(
                "pyannote diarized speaker identity match for {DiarizedSpeaker} and identity {IdentityId}: {SegmentCount} segment(s), {KnownSnippetCount} known snippet(s)",
                request.DiarizedSpeaker,
                request.IdentityId,
                segments.Count,
                layout.KnownSegments.Count);
            var unknownSpeaker = FindBestSpeaker(layout.UnknownSegment.Value, segments);
            if (string.IsNullOrWhiteSpace(unknownSpeaker))
            {
                logger.LogInformation(
                    "pyannote rejected speaker identity match for {DiarizedSpeaker} because the unknown sample had no speaker",
                    request.DiarizedSpeaker);
                return false;
            }

            var compared = 0;
            var matching = 0;
            foreach (var knownSegment in layout.KnownSegments)
            {
                var knownSpeaker = FindBestSpeaker(knownSegment, segments);
                if (string.IsNullOrWhiteSpace(knownSpeaker))
                {
                    continue;
                }

                compared++;
                if (string.Equals(knownSpeaker, unknownSpeaker, StringComparison.Ordinal))
                {
                    matching++;
                }
            }

            var minimumRatio = Math.Clamp(options.MinimumMatchingKnownSnippetRatio, 0, 1);
            var accepted = compared > 0 && (double)matching / compared >= minimumRatio;
            if (!accepted)
            {
                logger.LogInformation(
                    "pyannote rejected speaker identity match for {DiarizedSpeaker} and identity {IdentityId}: {Matching}/{Compared} known sample(s) matched, required ratio {RequiredRatio:P2}",
                    request.DiarizedSpeaker,
                    request.IdentityId,
                    matching,
                    compared,
                    minimumRatio);
            }
            else
            {
                logger.LogInformation(
                    "pyannote accepted speaker identity match for {DiarizedSpeaker} and identity {IdentityId}: {Matching}/{Compared} known sample(s) matched, required ratio {RequiredRatio:P2}",
                    request.DiarizedSpeaker,
                    request.IdentityId,
                    matching,
                    compared,
                    minimumRatio);
            }

            return accepted;
        }
        finally
        {
            SpeakerCompositeWav.TryDelete(tempPath);
        }
    }

    private async Task<IReadOnlyList<TranscriptionSegment>> DiarizeBytesAsync(
        byte[] wavBytes,
        CancellationToken cancellationToken)
    {
        var tempPath = Path.Combine(
            Path.GetTempPath(),
            "meeting-assistant-speaker-sample-pyannote",
            $"{Guid.NewGuid():N}.wav");
        try
        {
            Directory.CreateDirectory(Path.GetDirectoryName(tempPath)!);
            await File.WriteAllBytesAsync(tempPath, wavBytes, cancellationToken);
            return await DiarizePathAsync(tempPath, cancellationToken);
        }
        finally
        {
            SpeakerCompositeWav.TryDelete(tempPath);
        }
    }

    private async Task<IReadOnlyList<TranscriptionSegment>> DiarizePathAsync(
        string wavPath,
        CancellationToken cancellationToken)
    {
        TimeSpan duration;
        try
        {
            using var reader = new WaveFileReader(wavPath);
            duration = reader.TotalTime;
        }
        catch (Exception exception) when (exception is IOException or InvalidDataException or FormatException)
        {
            logger.LogInformation(exception, "pyannote speaker identity validation could not read WAV input");
            return [];
        }

        if (duration <= TimeSpan.Zero)
        {
            return [];
        }

        return await finalizer.FinalizeAsync(
            wavPath,
            [new TranscriptionSegment(TimeSpan.Zero, duration, "Unknown", "speaker identity validation sample")],
            options.Diarization,
            SpeechRecognitionPipelineOptions.Default,
            cancellationToken);
    }

    private SingleSpeakerCoverage AnalyzeSingleSpeakerCoverage(IReadOnlyList<TranscriptionSegment> segments)
    {
        var durationsBySpeaker = segments
            .Where(segment => !string.IsNullOrWhiteSpace(segment.Speaker) && segment.End > segment.Start)
            .GroupBy(segment => segment.Speaker)
            .Select(group => new
            {
                Speaker = group.Key,
                Duration = group.Sum(segment => (segment.End - segment.Start).TotalSeconds)
            })
            .Where(group => group.Duration > 0)
            .ToList();
        if (durationsBySpeaker.Count == 0)
        {
            return new SingleSpeakerCoverage(false, 0, null, 0, TimeSpan.Zero);
        }

        var total = durationsBySpeaker.Sum(group => group.Duration);
        var dominant = durationsBySpeaker.OrderByDescending(group => group.Duration).First();
        var coverage = total > 0 ? dominant.Duration / total : 0;
        if (durationsBySpeaker.Count == 1)
        {
            return new SingleSpeakerCoverage(
                true,
                durationsBySpeaker.Count,
                dominant.Speaker,
                coverage,
                TimeSpan.FromSeconds(total));
        }

        var accepted = total > 0 && coverage >= Math.Clamp(options.MinimumSingleSpeakerCoverage, 0, 1);
        return new SingleSpeakerCoverage(
            accepted,
            durationsBySpeaker.Count,
            dominant.Speaker,
            coverage,
            TimeSpan.FromSeconds(total));
    }

    private CompositeLayout WriteCompositeWav(
        string path,
        IReadOnlyList<byte[]> knownSnippets,
        byte[] unknownSnippet,
        TimeSpan silence)
    {
        using var firstReader = SpeakerCompositeWav.OpenFirstReadableWave(knownSnippets.Append(unknownSnippet));
        if (firstReader is null)
        {
            return new CompositeLayout(null, []);
        }

        using var writer = new WaveFileWriter(path, firstReader.WaveFormat);
        var current = TimeSpan.Zero;
        var knownSegments = new List<TimeSegment>();

        foreach (var snippet in knownSnippets.Where(storedSnippet => storedSnippet.Length > 0))
        {
            var segment = SpeakerCompositeWav.AppendSnippet(writer, firstReader.WaveFormat, snippet, current);
            if (segment is null)
            {
                continue;
            }

            knownSegments.Add(segment.Value);
            current = segment.Value.End + silence;
            SpeakerCompositeWav.WriteSilence(writer, firstReader.WaveFormat, silence);
        }

        var unknownSegment = SpeakerCompositeWav.AppendSnippet(
            writer,
            firstReader.WaveFormat,
            unknownSnippet,
            current);
        return new CompositeLayout(unknownSegment, knownSegments);
    }

    private static string? FindBestSpeaker(
        TimeSegment segment,
        IReadOnlyList<TranscriptionSegment> segments)
    {
        return SpeakerCompositeWav.FindBestSpeaker(segment, segments);
    }

    private sealed record CompositeLayout(TimeSegment? UnknownSegment, IReadOnlyList<TimeSegment> KnownSegments);

    private sealed record SingleSpeakerCoverage(
        bool Accepted,
        int SpeakerCount,
        string? DominantSpeaker,
        double Coverage,
        TimeSpan TotalDuration);
}