voxr 0.7.0

Audio and speech: codecs, DSP, G2P, voices, and TTS/ASR pipelines over boostr models
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

//! Finds the checkpoint file in a training output directory.
//!
//! oxidizr names a checkpoint `{prefix}_final.safetensors` at the end of
//! training and `{prefix}_step_{N}.safetensors` for an intermediate save.
//! `prefix` encodes the model name, sequence length, parameter count and
//! dtype. This crate cannot rebuild that prefix, so it scans the directory
//! instead of guessing the name.

use std::path::{Path, PathBuf};

use crate::error::{Error, Result};

/// Picks the checkpoint to load from a directory holding one or more
/// `*.safetensors` files: the one named `*_final.safetensors` if present,
/// else the one with the highest `_step_<N>` suffix, else the sole
/// `*.safetensors` file. Errors, naming every candidate, when the directory
/// holds several files and none of those rules picks one.
pub(crate) fn locate(dir: &Path) -> Result<PathBuf> {
    let mut candidates = candidates_in(dir)?;
    candidates.sort();

    if candidates.is_empty() {
        return Err(Error::ModelError {
            reason: format!("no *.safetensors checkpoint in {}", dir.display()),
        });
    }
    if candidates.len() == 1 {
        return Ok(candidates.remove(0));
    }

    if let Some(final_ckpt) = candidates.iter().find(|p| stem(p).ends_with("_final")) {
        return Ok(final_ckpt.clone());
    }

    let by_step: Vec<(u64, &PathBuf)> = candidates
        .iter()
        .filter_map(|p| step_suffix(&stem(p)).map(|n| (n, p)))
        .collect();
    if by_step.len() == candidates.len()
        && let Some((_, path)) = by_step.iter().max_by_key(|(n, _)| *n)
    {
        return Ok((*path).clone());
    }

    let names: Vec<String> = candidates.iter().map(|p| p.display().to_string()).collect();
    Err(Error::ModelError {
        reason: format!(
            "{} holds {} safetensors files and none is named *_final or *_step_<N>: {}",
            dir.display(),
            candidates.len(),
            names.join(", ")
        ),
    })
}

/// Every `*.safetensors` path directly inside `dir`.
fn candidates_in(dir: &Path) -> Result<Vec<PathBuf>> {
    Ok(std::fs::read_dir(dir)
        .map_err(|e| Error::ModelError {
            reason: format!("reading {}: {e}", dir.display()),
        })?
        .filter_map(|entry| entry.ok())
        .map(|entry| entry.path())
        .filter(|path| path.extension().is_some_and(|ext| ext == "safetensors"))
        .collect())
}

/// The file stem (name without extension), or an empty string.
fn stem(p: &Path) -> String {
    p.file_stem()
        .and_then(|s| s.to_str())
        .unwrap_or("")
        .to_string()
}

/// Parses the trailing `_step_<N>` off a file stem, if present.
fn step_suffix(stem: &str) -> Option<u64> {
    let (_, n) = stem.rsplit_once("_step_")?;
    n.parse().ok()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;

    #[test]
    fn picks_the_sole_checkpoint() {
        let dir = tempfile::tempdir().unwrap_or_else(|e| panic!("{e}"));
        fs::write(
            dir.path().join("g2p-char_128_4m_bf16_step_491.safetensors"),
            b"",
        )
        .unwrap_or_else(|e| panic!("{e}"));
        let picked = locate(dir.path()).unwrap_or_else(|e| panic!("{e}"));
        assert!(picked.ends_with("g2p-char_128_4m_bf16_step_491.safetensors"));
    }

    #[test]
    fn prefers_final_over_numbered_steps() {
        let dir = tempfile::tempdir().unwrap_or_else(|e| panic!("{e}"));
        for name in [
            "g2p-char_128_4m_bf16_step_491.safetensors",
            "g2p-char_128_4m_bf16_final.safetensors",
        ] {
            fs::write(dir.path().join(name), b"").unwrap_or_else(|e| panic!("{e}"));
        }
        let picked = locate(dir.path()).unwrap_or_else(|e| panic!("{e}"));
        assert!(picked.ends_with("g2p-char_128_4m_bf16_final.safetensors"));
    }

    #[test]
    fn picks_the_highest_step_when_no_final_exists() {
        let dir = tempfile::tempdir().unwrap_or_else(|e| panic!("{e}"));
        for name in [
            "g2p-char_128_4m_bf16_step_491.safetensors",
            "g2p-char_128_4m_bf16_step_982.safetensors",
        ] {
            fs::write(dir.path().join(name), b"").unwrap_or_else(|e| panic!("{e}"));
        }
        let picked = locate(dir.path()).unwrap_or_else(|e| panic!("{e}"));
        assert!(picked.ends_with("g2p-char_128_4m_bf16_step_982.safetensors"));
    }

    #[test]
    fn ambiguous_directory_names_every_candidate() {
        let dir = tempfile::tempdir().unwrap_or_else(|e| panic!("{e}"));
        for name in ["a.safetensors", "b.safetensors"] {
            fs::write(dir.path().join(name), b"").unwrap_or_else(|e| panic!("{e}"));
        }
        let err = locate(dir.path()).unwrap_err();
        let msg = err.to_string();
        assert!(msg.contains("a.safetensors") && msg.contains("b.safetensors"));
    }
}