piper_phoneme_streaming/g2p/
streaming.rs1use std::collections::HashMap;
2use std::sync::Arc;
3
4use crate::embedded_data::materialized_data_dir;
5use crate::error::Result;
6use crate::lang_detect::StreamingLanguageDetector;
7use crate::phoneme::PhonemeData;
8use crate::{
9 G2pToken, Language, SentenceUnit, StreamingSentencePhonemeUpgrade,
10 StreamingSentencePhonemeUpgradeSession, TextExpand, TextUnit, WordPhonemizer,
11};
12
13pub struct StreamingG2pSession {
18 expander: TextExpand,
19 session: StreamingSentencePhonemeUpgradeSession,
20}
21
22pub struct StreamingG2P {
28 phonemizers: HashMap<Language, WordPhonemizer>,
29 stream_sentence_upgrade: StreamingSentencePhonemeUpgrade,
30 default_language: Language,
31 languages: Vec<Language>,
32}
33
34impl StreamingG2P {
35 pub fn new(lang: Language) -> Result<Self> {
37 Self::with_languages(&[lang], lang)
38 }
39
40 pub fn with_languages(languages: &[Language], default_language: Language) -> Result<Self> {
46 let data_dir = materialized_data_dir()?;
47 let phdata = Arc::new(PhonemeData::load(data_dir)?);
48
49 let mut phonemizers = HashMap::new();
50 for &lang in languages {
51 phonemizers.insert(
52 lang,
53 WordPhonemizer::new_with_data(lang, Arc::clone(&phdata))?,
54 );
55 }
56 Ok(Self {
57 phonemizers,
58 stream_sentence_upgrade: StreamingSentencePhonemeUpgrade::new(default_language)?,
59 default_language,
60 languages: languages.to_vec(),
61 })
62 }
63
64 pub fn new_session(&self) -> StreamingG2pSession {
68 let expander = if self.languages.len() == 1 {
69 TextExpand::with_language(self.default_language)
70 } else {
71 let detector =
72 StreamingLanguageDetector::with_lingua(&self.languages, self.default_language);
73 TextExpand::with_detector(&self.languages, self.default_language, detector)
74 };
75
76 StreamingG2pSession {
77 session: self.stream_sentence_upgrade.new_session(),
78 expander,
79 }
80 }
81
82 pub fn push_text(
87 &self,
88 session: &mut StreamingG2pSession,
89 text: &str,
90 ) -> Result<Vec<G2pToken>> {
91 let mut outs = Vec::new();
92 for ch in text.chars() {
93 if let Some((unit, lang)) = session.expander.push(ch) {
94 let text_unit = TextUnit::from_expand_unit(unit, lang);
95 let phonemizer = self.phonemizer_for(lang);
96 let su = SentenceUnit::from_text_unit(text_unit, phonemizer)?;
97 outs.extend(session.session.push(su));
98 }
99 }
100 Ok(outs)
101 }
102
103 pub fn finish(&self, session: &mut StreamingG2pSession) -> Result<Vec<G2pToken>> {
108 let mut outs = Vec::new();
109 while let Some((unit, lang)) = session.expander.finish() {
110 let text_unit = TextUnit::from_expand_unit(unit, lang);
111 let phonemizer = self.phonemizer_for(lang);
112 let su = SentenceUnit::from_text_unit(text_unit, phonemizer)?;
113 outs.extend(session.session.push(su));
114 }
115 outs.extend(session.session.finish());
116 Ok(outs)
117 }
118
119 fn phonemizer_for(&self, lang: Language) -> &WordPhonemizer {
120 self.phonemizers
121 .get(&lang)
122 .or_else(|| self.phonemizers.get(&self.default_language))
123 .expect("at least the default language phonemizer must be present")
124 }
125}