speakrs 0.6.0

Fast Rust speaker diarization with pyannote-level accuracy and native CPU/CoreML/CUDA inference
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

//! Host-only evidence that every selectable host kernel is shipped in its PTX area

use std::collections::BTreeSet;

use super::{ComputeCapability, KernelModule, PtxTier, candidate, embedding, fbank, segmentation};

fn entries(ptx: &str) -> BTreeSet<&str> {
    ptx.lines()
        .filter_map(|line| line.trim().strip_prefix(".visible .entry "))
        .map(|entry| entry.split(['(', ' ', '\t']).next().unwrap())
        .collect()
}

fn missing<'a>(ptx: &str, required: &'a [&'a str]) -> Vec<&'a str> {
    let entries = entries(ptx);
    required
        .iter()
        .copied()
        .filter(|name| !entries.contains(name))
        .collect()
}

fn plans() -> Vec<(KernelModule, Vec<&'static str>)> {
    let plans = vec![
        (KernelModule::Fbank, fbank::REQUIRED_KERNELS.to_vec()),
        (
            KernelModule::Embedding,
            embedding::REQUIRED_KERNELS.to_vec(),
        ),
        (
            KernelModule::Segmentation,
            segmentation::REQUIRED_KERNELS.to_vec(),
        ),
        (KernelModule::Resnet, candidate::conv_kernel_inventory()),
        (KernelModule::Lstm, candidate::LSTM_KERNELS.to_vec()),
        (KernelModule::Sincnet, candidate::SINC_KERNELS.to_vec()),
        (KernelModule::Segdense, candidate::SEGDENSE_KERNELS.to_vec()),
        (
            KernelModule::FbankDft,
            candidate::FBANK_DFT_KERNELS.to_vec(),
        ),
        (KernelModule::LstmProj, candidate::LSTMPROJ_KERNELS.to_vec()),
        (KernelModule::Wideconv, candidate::WIDECONV_KERNELS.to_vec()),
    ];
    #[cfg(feature = "_cuda-libraries")]
    let plans = {
        let mut plans = plans;
        plans.push((KernelModule::Probe, super::probe::REQUIRED_KERNELS.to_vec()));
        plans
    };
    plans
}

#[test]
fn every_host_configuration_has_entries_in_every_paired_ptx_tier() {
    let mut pairings = 0;
    for (area, required) in plans() {
        assert!(
            area.variants().iter().next().is_some(),
            "{} has no embedded PTX",
            area.name()
        );
        assert!(
            !required.is_empty(),
            "{} has no host requirements",
            area.name()
        );
        for device in [
            ComputeCapability::new(7, 5),
            ComputeCapability::new(8, 0),
            ComputeCapability::new(8, 6),
            ComputeCapability::new(8, 9),
            ComputeCapability::new(9, 0),
            ComputeCapability::new(12, 0),
            ComputeCapability::new(13, 0),
        ] {
            for requested in std::iter::once(None).chain(PtxTier::ALL.map(Some)) {
                let Ok(limit) = PtxTier::select(device, requested) else {
                    continue;
                };
                // a lower shipped variant can pair with this target through fallback
                for (tier, ptx) in area.variants().iter() {
                    if tier > limit || tier.min_capability() > device {
                        continue;
                    }
                    pairings += 1;
                    let absent = missing(ptx, &required);
                    assert!(
                        absent.is_empty(),
                        "{} {tier} on {device}: missing {absent:?}",
                        area.name()
                    );
                }
            }
        }
    }

    assert!(pairings > 0, "enable a CUDA target feature for this test");
}

#[test]
fn missing_entry_is_detected_in_an_in_memory_ptx_fixture() {
    for (area, required) in plans() {
        for (tier, ptx) in area.variants().iter() {
            assert!(missing(ptx, &required).is_empty(), "{} {tier}", area.name());
            for name in &required {
                let fixture = ptx.replace(
                    &format!(".visible .entry {name}("),
                    &format!(".visible .entry removed_{name}("),
                );
                assert_ne!(fixture, ptx, "entry parser fixture did not remove {name}");
                assert!(
                    missing(&fixture, &required).contains(name),
                    "missing {name} was not detected"
                );
            }
        }
    }
}

#[cfg(feature = "cuda")]
#[test]
fn full_cuda_embeds_all_targets_and_every_shipped_areas_kernels() {
    assert_eq!(
        PtxTier::compiled_in().collect::<Vec<_>>(),
        [PtxTier::Sm75, PtxTier::Sm80, PtxTier::Sm90, PtxTier::Sm120]
    );
    for (area, required) in plans() {
        let variants = area.variants();
        for (_, ptx) in variants.iter() {
            assert!(
                missing(ptx, &required).is_empty(),
                "{} is missing shipped kernels",
                area.name()
            );
        }
        assert!(
            variants.iter().next().is_some(),
            "{} has no shipped module",
            area.name()
        );
    }
}