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
144
145
146
147
148

use std::path::PathBuf;

use crate::clustering::plda::PldaTransform;
use crate::inference::ExecutionMode;
use crate::inference::embedding::EmbeddingModel;
use crate::inference::segmentation::SegmentationModel;
use crate::models::ModelBundle;
use crate::powerset::PowersetMapping;

use super::OwnedDiarizationPipeline;
use super::config::{PipelineConfig, RuntimeConfig, segmentation_step_seconds};
use super::queued::{QueueConfig, QueueReceiver, QueueSender};
use super::types::PipelineError;

/// Builder for constructing diarization pipelines
///
/// # Examples
///
/// ```no_run
/// use speakrs::{ExecutionMode, PipelineBuilder};
///
/// // minimal
/// let mut pipeline = PipelineBuilder::from_pretrained(ExecutionMode::Cpu)?.build()?;
///
/// // from local directory
/// let mut pipeline = PipelineBuilder::from_dir("./models", ExecutionMode::Cpu)?
///     .build()?;
/// # Ok::<(), Box<dyn std::error::Error + Send + Sync>>(())
/// ```
pub struct PipelineBuilder {
    bundle: ModelBundle,
    mode: ExecutionMode,
    runtime: Option<RuntimeConfig>,
    pipeline: Option<PipelineConfig>,
    queue: Option<QueueConfig>,
}

impl PipelineBuilder {
    /// Start from a local models directory
    pub fn from_dir(
        models_dir: impl Into<PathBuf>,
        mode: ExecutionMode,
    ) -> Result<Self, PipelineError> {
        Ok(Self {
            bundle: ModelBundle::from_dir(models_dir)?,
            mode,
            runtime: None,
            pipeline: None,
            queue: None,
        })
    }

    /// Start from a pre-resolved [`ModelBundle`](ModelBundle)
    pub fn from_bundle(bundle: ModelBundle, mode: ExecutionMode) -> Self {
        Self {
            bundle,
            mode,
            runtime: None,
            pipeline: None,
            queue: None,
        }
    }

    /// Download models from HuggingFace and start building
    #[cfg(feature = "online")]
    #[cfg_attr(docsrs, doc(cfg(feature = "online")))]
    pub fn from_pretrained(mode: ExecutionMode) -> Result<Self, PipelineError> {
        mode.validate()?;
        let bundle = ModelBundle::from_pretrained(mode)?;
        Ok(Self::from_bundle(bundle, mode))
    }

    /// Override runtime config
    pub fn runtime(mut self, config: RuntimeConfig) -> Self {
        self.runtime = Some(config);
        self
    }

    /// Override pipeline config (thresholds, clustering)
    pub fn pipeline(mut self, config: PipelineConfig) -> Self {
        self.pipeline = Some(config);
        self
    }

    /// Override queue construction config
    ///
    /// Controls request-channel capacity. The local queue is not durable
    pub fn queue(mut self, config: QueueConfig) -> Self {
        self.queue = Some(config);
        self
    }

    /// Build the owned pipeline
    pub fn build(self) -> Result<OwnedDiarizationPipeline, PipelineError> {
        self.mode.validate()?;

        let pipeline = self
            .pipeline
            .unwrap_or_else(|| PipelineConfig::for_mode(self.mode));
        let runtime = self.runtime.unwrap_or_default();
        #[cfg(feature = "coreml")]
        let coreml_chunk_execution_policy = runtime.coreml_chunk_execution_policy();

        #[cfg(feature = "_metrics")]
        let step = match runtime.experiment {
            Some(experiment) => {
                experiment
                    .validate(self.mode)
                    .map_err(|error| PipelineError::Other(error.to_string()))?;
                experiment.segmentation_step_seconds()
            }
            None => segmentation_step_seconds(self.mode),
        };
        #[cfg(not(feature = "_metrics"))]
        let step = segmentation_step_seconds(self.mode);

        let seg_model = SegmentationModel::with_mode_and_config(
            self.bundle.segmentation_path(),
            step as f32,
            self.mode,
            &runtime,
        )?;
        let emb_model = EmbeddingModel::with_mode_and_config(
            self.bundle.embedding_path(),
            self.mode,
            &runtime,
        )?;
        let plda = PldaTransform::from_dir(self.bundle.plda_dir())?;

        Ok(OwnedDiarizationPipeline {
            seg_model,
            emb_model,
            plda,
            powerset: PowersetMapping::new(3, 2),
            default_config: pipeline,
            #[cfg(feature = "coreml")]
            coreml_chunk_execution_policy,
        })
    }

    /// Build and immediately convert to a background-processing queue
    pub fn build_queued(self) -> Result<(QueueSender, QueueReceiver), PipelineError> {
        let queue = self.queue.unwrap_or_default();
        queue.validate()?;
        let pipeline = self.build()?;
        Ok(pipeline.into_queued_with_queue_config(queue)?)
    }
}