admin
b5c083e06f
Some checks failed
CodeQL / Analyze (go) (push) Successful in 6m28s
Docker Image / build-docker (push) Failing after 13m26s
Lint and Testing / lint (push) Successful in 11m17s
Lint and Testing / test (push) Successful in 11m17s
Lint and Testing / golangci (push) Successful in 2m40s
73 lines
1.9 KiB
Go
73 lines
1.9 KiB
Go
package config
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
)
|
|
|
|
type Diarization struct {
|
|
Enabled bool `yaml:"enabled"`
|
|
ModelDir string `yaml:"model_dir"`
|
|
SegmentationModel string `yaml:"segmentation_model"`
|
|
EmbeddingModel string `yaml:"embedding_model"`
|
|
NumThreads int `yaml:"num_threads"`
|
|
NumClusters int `yaml:"num_clusters"`
|
|
ClusteringThreshold float32 `yaml:"clustering_threshold"`
|
|
MinDurationOn float32 `yaml:"min_duration_on"`
|
|
MinDurationOff float32 `yaml:"min_duration_off"`
|
|
}
|
|
|
|
func (d Diarization) WithDefaults() Diarization {
|
|
if d.ModelDir == "" {
|
|
d.ModelDir = "./models/diarization"
|
|
}
|
|
if d.SegmentationModel == "" {
|
|
d.SegmentationModel = "pyannote-segmentation-3-0/model.onnx"
|
|
}
|
|
if d.EmbeddingModel == "" {
|
|
d.EmbeddingModel = "3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx"
|
|
}
|
|
if d.NumThreads <= 0 {
|
|
d.NumThreads = 2
|
|
}
|
|
if d.ClusteringThreshold <= 0 {
|
|
d.ClusteringThreshold = 0.5
|
|
}
|
|
if d.MinDurationOn <= 0 {
|
|
d.MinDurationOn = 0.3
|
|
}
|
|
if d.MinDurationOff <= 0 {
|
|
d.MinDurationOff = 0.5
|
|
}
|
|
return d
|
|
}
|
|
|
|
func (d Diarization) SegmentationPath() string {
|
|
return resolveModelPath(d.ModelDir, d.SegmentationModel)
|
|
}
|
|
|
|
func (d Diarization) EmbeddingPath() string {
|
|
return resolveModelPath(d.ModelDir, d.EmbeddingModel)
|
|
}
|
|
|
|
func (d Diarization) ModelsPresent() bool {
|
|
if _, err := os.Stat(d.SegmentationPath()); err != nil {
|
|
return false
|
|
}
|
|
if _, err := os.Stat(d.EmbeddingPath()); err != nil {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
func resolveModelPath(dir, name string) string {
|
|
if filepath.IsAbs(name) {
|
|
return name
|
|
}
|
|
return filepath.Join(dir, name)
|
|
}
|
|
|
|
func (d Diarization) Active() bool {
|
|
return d.Enabled
|
|
}
|