diff --git a/README.md b/README.md index 3b9924f..cef0445 100644 --- a/README.md +++ b/README.md @@ -148,6 +148,22 @@ cargo run -p floravox-cli -- g2p --phonetisaurus cmudict-20170708.o8.fst \ The Phonetisaurus decoder is a clean-room Rust implementation of the OpenFst file format plus a shortest-path search. Both layouts load (embedded symbol tables, or `model.fst` plus separate table files), and 16-byte and 20-byte arc encodings are detected automatically. +## P2G: turning phonemes back into words + +The G2P tiers are direction-symmetric, so the same tools spell pronunciations — including made-up ones ("moo", "eek"): + +- **Tier 1 — inverted lexicon** (`floravox-fst-compile --reverse LEXICON.tsv en-p2g`): exact homophone sets for in-lexicon pronunciations. `/ɹ ˈaɪ t/` → right, write, wright, rite (+ reit — whatever the corpus really contains). +- **Tier 2 — reversed WFST** (`floravox-train-phonetisaurus --reverse LEXICON.tsv en-p2g.fst`): generalizes to unseen phoneme strings through the same segment-alignment statistics. Measured on the published en bundle (124k entries): 39% exact / 41% with homophone credit / CER 40% / 99% coverage — at published Phonetisaurus-P2G baselines. Strong on medium/long words, weak on very short ones. +- **Tier 3 — ByT5 P2G** (optional, `onnx` feature): any byte-level phoneme→orthography checkpoint (e.g. a PolyIPA export) as the final fallback; it knows the whole IPA, so it can spell phonemes the WFST never saw. + +```console +# chain all three tiers (first non-empty candidate list wins): +cargo run -p floravox-g2p --bin floravox-p2g -- --lexicon en-lex-p2g --model en-p2g.fst --byt5 polyipa-onnx/ "ʃ ˈi k" +# sheek +``` + +In Rust, the tiers implement the `P2g` trait (`floravox_g2p::p2g`) and chain with `ChainedP2g`, mirroring the G2P fallback chain. Note the phoneme strings must match the training lexicon's conventions (the en bundle carries gruut-style stress marks: `k ˈæ t`, not `k æ t`). + ## How accurate is it? Two audit tools ship in `python/`: diff --git a/crates/floravox-g2p/src/bin/floravox-fst-compile.rs b/crates/floravox-g2p/src/bin/floravox-fst-compile.rs index 339b352..194542b 100644 --- a/crates/floravox-g2p/src/bin/floravox-fst-compile.rs +++ b/crates/floravox-g2p/src/bin/floravox-fst-compile.rs @@ -3,9 +3,15 @@ //! //! Usage: //! ```text -//! floravox-fst-compile [--format FMT] INPUT OUTPUT_STEM +//! floravox-fst-compile [--format FMT] [--reverse] INPUT OUTPUT_STEM //! ``` //! +//! `--reverse` inverts the mapping (P2G): the output stem holds +//! phoneme-string keys → candidate words. Homophones are merged under +//! one key (space-separated, first-seen order), so /ɹ aɪ t/ → +//! "right write wright rite". Keys keep their exact case (IPA is +//! case-significant). +//! //! FMT: auto (default) | cmudict | ipa-tsv | tsv //! //! * cmudict — `WORD P HH R AH1 N` (converted ARPABET → IPA) @@ -20,12 +26,33 @@ use std::io::Read; use floravox_g2p::{ingest, SourceFormat}; fn usage_exit() -> ! { - eprintln!("usage: floravox-fst-compile [--format auto|cmudict|ipa-tsv|tsv] INPUT OUTPUT_STEM"); + eprintln!( + "usage: floravox-fst-compile [--format auto|cmudict|ipa-tsv|tsv] [--reverse] INPUT OUTPUT_STEM" + ); std::process::exit(2); } +/// Invert (word, phonemes) rows into (phonemes, homophone words) rows, +/// merging duplicates in first-seen order. +fn reverse_rows(rows: Vec<(String, String)>) -> Vec<(String, String)> { + use std::collections::BTreeMap; + let mut merged: BTreeMap> = BTreeMap::new(); + for (word, phones) in rows { + let key = phones.split_whitespace().collect::>().join(" "); + let entry = merged.entry(key).or_default(); + if !entry.contains(&word) { + entry.push(word); + } + } + merged + .into_iter() + .map(|(k, words)| (k, words.join(" "))) + .collect() +} + fn main() { let mut format_arg: Option = None; + let mut reverse = false; let mut positional: Vec = Vec::new(); let mut args = std::env::args().skip(1); while let Some(arg) = args.next() { @@ -36,6 +63,8 @@ fn main() { }; } else if let Some(value) = arg.strip_prefix("--format=") { format_arg = Some(value.to_string()); + } else if arg == "--reverse" { + reverse = true; } else { positional.push(arg); } @@ -79,9 +108,17 @@ fn main() { ); } - match floravox_g2p::LexiconWriter::new(&stem).write(ingested.rows) { + let rows = if reverse { + reverse_rows(ingested.rows) + } else { + ingested.rows + }; + match floravox_g2p::LexiconWriter::new(&stem).write(rows) { Ok(count) => { - println!("wrote {stem}.fst + {stem}.pho ({count} entries, format {format:?})"); + let direction = if reverse { "p2g" } else { "g2p" }; + println!( + "wrote {stem}.fst + {stem}.pho ({count} entries, format {format:?}, {direction})" + ); } Err(e) => { eprintln!("compile failed: {e}"); diff --git a/crates/floravox-g2p/src/bin/floravox-p2g.rs b/crates/floravox-g2p/src/bin/floravox-p2g.rs new file mode 100644 index 0000000..ad8909f --- /dev/null +++ b/crates/floravox-g2p/src/bin/floravox-p2g.rs @@ -0,0 +1,131 @@ +//! Spell a pronunciation: phoneme-to-grapheme across the floravox tiers. +//! +//! Usage: +//! floravox-p2g [OPTIONS] PHONEMES... +//! +//! PHONEMES are whitespace-separated symbols (`floravox-p2g m uː`) or a +//! single slashed argument (`floravox-p2g /ɹ aɪ t/`). +//! +//! Tiers (first non-empty candidate list wins, exactly like the G2P +//! chain): +//! --lexicon STEM inverted FST lexicon (floravox-fst-compile --reverse) +//! --model STEM reversed WFST (floravox-train-phonetisaurus --reverse) +//! --byt5 DIR ByT5/PolyIPA ONNX export (encoder+decoder) +//! +//! Candidates print one per line (lexicon homophones keep corpus order; +//! model/byt5 tiers contribute their 1-best). Exit code 0 even when +//! nothing matched (empty output); 2 on usage/config errors. +//! +//! ```text +//! $ floravox-p2g --lexicon cmudict-p2g --model en-p2g "ɹ aɪ t" +//! right +//! write +//! wright +//! rite +//! ``` + +#[cfg(feature = "onnx")] +use floravox_g2p::p2g::ByT5P2g; +use floravox_g2p::p2g::{ChainedP2g, P2g, P2gLexicon}; +use floravox_g2p::phonetisaurus::PhonetisaurusP2g; + +fn usage_exit() -> ! { + eprintln!( + "usage: floravox-p2g [--lexicon STEM] [--model STEM] [--byt5 DIR] [--lang eng] PHONEMES..." + ); + std::process::exit(2); +} + +/// Chained tiers are boxed so any subset of tiers can be active. +type DynP2g<'a> = Box; + +fn main() { + let mut lexicon: Option = None; + let mut model: Option = None; + let mut byt5: Option = None; + let mut lang: Option = None; + let mut phones: Vec = Vec::new(); + let mut args = std::env::args().skip(1); + while let Some(a) = args.next() { + match a.as_str() { + "--lexicon" => lexicon = args.next(), + "--model" => model = args.next(), + "--byt5" => byt5 = args.next(), + "--lang" => lang = args.next(), + other if other.starts_with("--") => usage_exit(), + // Whitespace-separated symbols, however many args they came in. + other => phones.extend(other.split_whitespace().map(str::to_owned)), + } + } + if phones.is_empty() || (lexicon.is_none() && model.is_none() && byt5.is_none()) { + usage_exit(); + } + // Slashed single-argument form: /ɹ aɪ t/ → ["ɹ", "aɪ", "t"]. + if phones.len() == 1 { + let s = phones[0].trim(); + if s.len() >= 2 && s.starts_with('/') && s.ends_with('/') { + phones = s[1..s.len() - 1] + .split_whitespace() + .map(str::to_owned) + .collect(); + } + } + if phones.iter().any(String::is_empty) { + usage_exit(); + } + + // Chain: lexicon → WFST → ByT5. Boxes keep the tiers optional while + // the chain type stays concrete at each link. + let mut chain: Option = None; + + if let Some(stem) = &lexicon { + let lex = P2gLexicon::open(stem).unwrap_or_else(|e| { + eprintln!("cannot open lexicon {stem}: {e}"); + std::process::exit(2); + }); + chain = Some(match chain { + Some(inner) => Box::new(ChainedP2g(lex, inner)), + None => Box::new(lex), + }); + } + if let Some(stem) = &model { + let m = PhonetisaurusP2g::open(stem).unwrap_or_else(|e| { + eprintln!("cannot open model {stem}: {e}"); + std::process::exit(2); + }); + chain = Some(match chain { + Some(inner) => Box::new(ChainedP2g(m, inner)), + None => Box::new(m), + }); + } + #[cfg(feature = "onnx")] + if let Some(dir) = &byt5 { + let mut engine = ByT5P2g::load( + std::path::Path::new(dir).join("encoder_model.onnx"), + std::path::Path::new(dir).join("decoder_model.onnx"), + ) + .unwrap_or_else(|e| { + eprintln!("cannot load byt5 model from {dir}: {e}"); + std::process::exit(2); + }); + if let Some(tag) = &lang { + engine.lang.clone_from(tag); + } + chain = Some(match chain { + Some(inner) => Box::new(ChainedP2g(engine, inner)), + None => Box::new(engine), + }); + } + #[cfg(not(feature = "onnx"))] + if byt5.is_some() || lang.is_some() { + eprintln!("--byt5/--lang require a build with the onnx feature"); + std::process::exit(2); + } + + let Some(mut p2g) = chain else { + usage_exit(); + }; + for candidate in p2g.spell(&phones) { + println!("{candidate}"); + } +} diff --git a/crates/floravox-g2p/src/bin/floravox-train-phonetisaurus.rs b/crates/floravox-g2p/src/bin/floravox-train-phonetisaurus.rs index 2ba3a7a..7eee071 100644 --- a/crates/floravox-g2p/src/bin/floravox-train-phonetisaurus.rs +++ b/crates/floravox-g2p/src/bin/floravox-train-phonetisaurus.rs @@ -23,6 +23,15 @@ //! floravox-train-phonetisaurus LEXICON.tsv MODEL.fst //! [--order 7] [--iters 8] [--gmax 2] [--pmax 2] //! [--holdout 0.05] [--seed 7] [--metrics metrics.json] +//! [--reverse] +//! +//! `--reverse` trains a P2G model (phonemes on the input side, +//! graphemes on the output side) loadable by `PhonetisaurusP2g`. The +//! M2M alignment is direction-symmetric, so reversal is the ilabel/ +//! olabel swap; `--gmax`/`--pmax` keep their grapheme/phoneme meanings +//! whichever way the model points. The holdout evaluation reports +//! word-level exact match and character error rate (spelled vs +//! reference word). // This is a training CLI: the DP math uses single-letter bindings (g/p), // count/size casts, and a long orchestrating main by design. @@ -37,7 +46,7 @@ clippy::type_complexity )] -use floravox_g2p::phonetisaurus::write_model; +use floravox_g2p::phonetisaurus::{write_model, PhonetisaurusP2g}; use floravox_g2p::PhonetisaurusG2p; use std::collections::HashMap; @@ -59,8 +68,10 @@ fn main() { let mut holdout = 0.05f64; let mut seed: u64 = 7; let mut metrics_path: Option = None; + let mut reverse = false; while let Some(a) = args.next() { match a.as_str() { + "--reverse" => reverse = true, "--order" => order = parse(&mut args, &a), "--iters" => iters = parse(&mut args, &a), "--gmax" => gmax = parse(&mut args, &a), @@ -247,12 +258,20 @@ fn main() { // ---- symbol tables ---- let mut isyms: Vec<(String, i32)> = vec![("".into(), 0), ("|".into(), 1), ("_".into(), 2)]; let mut osyms = isyms.clone(); + // Reversal is exactly the input/output assignment: phoneme + // compounds on the input side, grapheme compounds on the output. let tok_labels: Vec<(i32, i32)> = tok_pairs .iter() .map(|&(g, p)| { - let il = label(g, &glist, &mut isyms, true); - let ol = label(p, &plist, &mut osyms, false); - (il, ol) + if reverse { + let il = label(p, &plist, &mut isyms, false); + let ol = label(g, &glist, &mut osyms, true); + (il, ol) + } else { + let il = label(g, &glist, &mut isyms, true); + let ol = label(p, &plist, &mut osyms, false); + (il, ol) + } }) .collect(); @@ -288,6 +307,41 @@ fn main() { ); // ---- evaluation through the shipped file ---- + if reverse { + let model = PhonetisaurusP2g::open(&output).expect("reload model for eval"); + // Homophone-aware gold set: a spelling is exact when it matches + // ANY word sharing the pronunciation (right/write/wright/rite). + let mut by_phones: std::collections::HashMap> = + std::collections::HashMap::new(); + for (w, ph) in &lex { + by_phones.entry(ph.join(" ")).or_default().push(w.clone()); + } + let (exact, cer, coverage) = evaluate_p2g(&model, &hold, &by_phones); + eprintln!( + "eval (p2g): exact {:.1}%, CER {:.1}%, coverage {:.1}%", + exact * 100.0, + cer * 100.0, + coverage * 100.0 + ); + if let Some(path) = metrics_path { + let json = format!( + "{{\n \"direction\": \"p2g\",\n \"train\": {},\n \"holdout\": {},\n \"exact_match\": {:.4},\n \"cer\": {:.4},\n \"coverage\": {:.4},\n \"order\": {},\n \"iters\": {},\n \"gmax\": {},\n \"pmax\": {},\n \"states\": {},\n \"arcs\": {}\n}}\n", + train.len(), + hold.len(), + exact, + cer, + coverage, + order, + iters, + gmax, + pmax, + n_states, + arc_count + ); + std::fs::write(path, json).expect("write metrics"); + } + return; + } let model = PhonetisaurusG2p::open(&output).expect("reload model for eval"); let (exact, per, coverage) = evaluate(&model, &hold); eprintln!( @@ -513,6 +567,47 @@ fn align( path } +/// P2G holdout metrics: homophone-aware exact match (any word sharing +/// the pronunciation counts), character error rate against the nearest +/// gold homophone, decode coverage. +fn evaluate_p2g( + model: &PhonetisaurusP2g, + hold: &[(String, Vec)], + by_phones: &std::collections::HashMap>, +) -> (f64, f64, f64) { + let mut exact = 0usize; + let mut cer_sum = 0.0f64; + let mut covered = 0usize; + let n = hold.len().max(1); + for (word, ref_phones) in hold { + let Some(spelled) = model.spell(ref_phones) else { + continue; + }; + covered += 1; + let gold = by_phones + .get(&ref_phones.join(" ")) + .map_or(std::slice::from_ref(word), |v| v.as_slice()); + if gold.contains(&spelled) { + exact += 1; + } + // CER against the nearest gold homophone. + let a: Vec = spelled.chars().map(|c| c.to_string()).collect(); + let best = gold + .iter() + .map(|g| { + let b: Vec = g.chars().map(|c| c.to_string()).collect(); + edit_distance(&a, &b) as f64 / g.chars().count().max(1) as f64 + }) + .fold(f64::INFINITY, f64::min); + cer_sum += best; + } + ( + exact as f64 / n as f64, + cer_sum / covered.max(1) as f64, + covered as f64 / n as f64, + ) +} + /// Exact-match rate, mean PER, decode coverage over holdout pairs. fn evaluate(model: &PhonetisaurusG2p, hold: &[(String, Vec)]) -> (f64, f64, f64) { let mut exact = 0usize; diff --git a/crates/floravox-g2p/src/lib.rs b/crates/floravox-g2p/src/lib.rs index b356a94..9fd525b 100644 --- a/crates/floravox-g2p/src/lib.rs +++ b/crates/floravox-g2p/src/lib.rs @@ -44,6 +44,7 @@ use std::sync::Arc; pub mod ingest; pub mod misaki; +pub mod p2g; #[cfg(feature = "uroman")] pub mod uroman; @@ -162,6 +163,43 @@ impl FstLexicon> { Self::from_rows(rows) } + /// Compile from (key, value) rows **without** case folding or + /// de-duplication (last row wins for duplicate keys). Used by the + /// P2G lexicons, whose keys are phoneme strings where case is + /// significant (IPA) and values may contain multiple words. + /// + /// # Errors + /// + /// [`G2pError::Compile`] when the FST cannot be built. + pub fn from_rows_raw(rows: Vec<(String, String)>) -> Result { + // Sort (fst keys must ascend) without case folding; exact keys, + // last row wins on duplicates. + let mut unique: std::collections::BTreeMap = + std::collections::BTreeMap::new(); + for (k, v) in rows { + unique.insert(k, v); + } + let mut blob = Vec::new(); + let mut entries: Vec<(String, u64)> = Vec::with_capacity(unique.len()); + for (key, value) in unique { + let offset = blob.len(); + blob.extend_from_slice(value.as_bytes()); + let len = blob.len() - offset; + entries.push((key, pack_value(offset, len))); + } + let mut builder = fst::MapBuilder::memory(); + for (key, value) in entries { + builder + .insert(key.as_bytes(), value) + .map_err(|e| G2pError::Compile(e.to_string()))?; + } + let fst_bytes = builder + .into_inner() + .map_err(|e| G2pError::Compile(e.to_string()))?; + let map = fst::Map::new(fst_bytes).map_err(|e| G2pError::Compile(e.to_string()))?; + Ok(Self { map, blob }) + } + /// Compile an in-memory lexicon from (word, phonemes) rows. /// # Errors /// @@ -220,6 +258,18 @@ impl> FstLexicon { self.map.is_empty() } + /// Look up a key **exactly** (no case folding); `None` when absent. + /// Returns the raw value string. Used by P2G lexicons whose keys are + /// phoneme strings. + #[must_use] + pub fn lookup_raw(&self, key: &str) -> Option { + let value = self.map.get(key.as_bytes())?; + let (off, len) = unpack_value(value); + let blob = self.blob.as_ref(); + let raw = blob.get(off..off + len)?; + std::str::from_utf8(raw).ok().map(str::to_owned) + } + /// Look up a word (case-insensitive); `None` when out of vocabulary. #[must_use] pub fn lookup(&self, word: &str) -> Option> { diff --git a/crates/floravox-g2p/src/p2g.rs b/crates/floravox-g2p/src/p2g.rs new file mode 100644 index 0000000..8040899 --- /dev/null +++ b/crates/floravox-g2p/src/p2g.rs @@ -0,0 +1,243 @@ +//! # Phoneme-to-grapheme (P2G) — the inverse of the G2P tiers. +//! +//! Same three-tier architecture, mirrored: +//! +//! * **Tier 1** — [`P2gLexicon`]: an *inverted* [`FstLexicon`] compiled +//! with `floravox-fst-compile --reverse` (phoneme string → candidate +//! words). Exact homophone sets for in-lexicon pronunciations +//! (/ɹ aɪ t/ → right, write, wright, rite), mmap'd, sub-100 µs. +//! * **Tier 2** — [`PhonetisaurusP2g`](crate::phonetisaurus::PhonetisaurusP2g): +//! a *reversed* WFST trained with `floravox-train-phonetisaurus +//! --reverse`. Generalizes to unseen phoneme strings (/m uː/ → "moo", +//! /iː k/ → "eek") — the alignment statistics are segment-level, so +//! novel pronunciations still spell plausibly. +//! * **Tier 3** — [`ByT5P2g`] (behind the `onnx` feature): a +//! byte-level seq2seq P2G model (e.g. a `PolyIPA` `ByT5` export) as the +//! final fallback. +//! +//! Chain tiers with [`ChainedP2g`] (first non-empty candidate list wins), +//! mirroring [`ChainedFallback`](crate::ChainedFallback). +//! +//! ``` +//! use floravox_g2p::p2g::{P2g, P2gLexicon}; +//! +//! let mut lex = P2gLexicon::from_rows(vec![ +//! ("ɹ aɪ t".into(), "right write wright rite".into()), +//! ("k æ t".into(), "cat".into()), +//! ]) +//! .unwrap(); +//! assert_eq!(lex.spell(&["ɹ".into(), "aɪ".into(), "t".into()]), +//! ["right".to_string(), "write".into(), "wright".into(), "rite".into()]); +//! // OOV pronunciation: no candidates (chain a PhonetisaurusP2g or +//! // ByT5P2g behind it for those). +//! assert!(lex.spell(&["m".into(), "uː".into()]).is_empty()); +//! ``` + +use crate::{FstLexicon, Phoneme}; + +#[cfg(feature = "onnx")] +use crate::byt5::Byt5G2p; + +/// A phoneme-to-grapheme strategy: given a pronunciation, return ranked +/// candidate spellings (best first, possibly empty). +pub trait P2g { + /// Spell `phonemes`, returning ranked candidate words. + fn spell(&mut self, phonemes: &[Phoneme]) -> Vec; +} + +/// **Tier 1**: exact-match inverted lexicon (phoneme string → words). +/// +/// Build with `floravox-fst-compile --reverse LEXICON.tsv OUT-p2g` and +/// [`P2gLexicon::open`], or [`P2gLexicon::from_rows`] in memory. Keys are +/// compared exactly (IPA case is significant). +pub struct P2gLexicon = Vec> { + inner: FstLexicon, +} + +impl P2gLexicon> { + /// Compile from (phonemes, words) rows. The phoneme side must be the + /// space-separated key; the value is one or more space-separated + /// candidate words. Duplicate keys keep the last row (use the CLI's + /// `--reverse` mode, which merges homophones, for corpus builds). + /// + /// # Errors + /// + /// [`crate::G2pError::Compile`] when the FST cannot be built. + pub fn from_rows(rows: Vec<(String, String)>) -> Result { + Ok(Self { + inner: FstLexicon::from_rows_raw(rows)?, + }) + } +} + +impl P2gLexicon { + /// Open a compiled inverted lexicon stem (`OUT-p2g.fst` + + /// `OUT-p2g.pho` produced by `floravox-fst-compile --reverse`). + /// Both halves are memory-mapped. + /// + /// # Errors + /// + /// [`crate::G2pError::Open`] when the files cannot be read or parsed. + pub fn open(stem: impl AsRef) -> Result { + Ok(Self { + inner: FstLexicon::open(stem)?, + }) + } +} + +impl> P2gLexicon { + /// Number of distinct pronunciations in the lexicon. + #[must_use] + pub fn len(&self) -> usize { + self.inner.len() + } + + /// True when the lexicon has no entries. + #[must_use] + pub fn is_empty(&self) -> bool { + self.inner.is_empty() + } +} + +impl> P2g for P2gLexicon { + fn spell(&mut self, phonemes: &[Phoneme]) -> Vec { + let key = phonemes.join(" "); + self.inner + .lookup_raw(&key) + .map(|v| v.split_whitespace().map(str::to_owned).collect()) + .unwrap_or_default() + } +} + +impl P2g for Box { + fn spell(&mut self, phonemes: &[Phoneme]) -> Vec { + (**self).spell(phonemes) + } +} + +/// Chain two P2G tiers: the first non-empty candidate list wins (lexicon +/// → model, model → neural, …), mirroring +/// [`ChainedFallback`](crate::ChainedFallback). +pub struct ChainedP2g(pub A, pub B); + +impl P2g for ChainedP2g { + fn spell(&mut self, phonemes: &[Phoneme]) -> Vec { + let first = self.0.spell(phonemes); + if first.is_empty() { + self.1.spell(phonemes) + } else { + first + } + } +} + +/// **Tier 3**: byte-level seq2seq P2G (a `ByT5` ONNX export trained for +/// phoneme→orthography, e.g. `PolyIPA`: input `: IPA`). +/// +/// Wrap an encoder/decoder pair exported with `optimum-cli export onnx` +/// (same layout [`Byt5G2p`] loads). The input convention — language tag +/// prefix, IPA joined by spaces — matches PolyIPA-style checkpoints; +/// adjust [`ByT5P2g::lang`] for other conventions. +#[cfg(feature = "onnx")] +pub struct ByT5P2g { + engine: Byt5G2p, + /// Language tag prepended to the input (`: …` for `PolyIPA`). + pub lang: String, +} + +#[cfg(feature = "onnx")] +impl ByT5P2g { + /// Load from an `optimum` ONNX export of a `ByT5` P2G model — the + /// encoder and decoder session files + /// (`encoder_model.onnx` / `decoder_model.onnx`). + /// + /// # Errors + /// + /// [`crate::G2pError::Inference`] when the sessions cannot be created. + pub fn load( + encoder_path: impl AsRef, + decoder_path: impl AsRef, + ) -> Result { + Ok(Self { + engine: Byt5G2p::load(encoder_path, decoder_path)?, + lang: "eng".into(), + }) + } +} + +#[cfg(feature = "onnx")] +impl P2g for ByT5P2g { + fn spell(&mut self, phonemes: &[Phoneme]) -> Vec { + let input = format!("<{}>: {}", self.lang, phonemes.join(" ")); + self.engine + .phonemize_word(&input) + .ok() + .map(|tokens| vec![tokens.join("")]) + .unwrap_or_default() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ph(list: &[&str]) -> Vec { + list.iter().map(|&s| s.to_owned()).collect() + } + + #[test] + fn lexicon_exact_match_and_homophones() { + let mut lex = P2gLexicon::from_rows(vec![ + ("ɹ aɪ t".into(), "right write wright rite".into()), + ("k æ t".into(), "cat".into()), + ]) + .unwrap(); + assert_eq!(lex.len(), 2); + assert_eq!( + lex.spell(&ph(&["ɹ", "aɪ", "t"])), + ["right", "write", "wright", "rite"] + ); + assert_eq!(lex.spell(&ph(&["k", "æ", "t"])), ["cat"]); + } + + #[test] + fn lexicon_oov_returns_empty() { + let mut lex = P2gLexicon::from_rows(vec![("k æ t".into(), "cat".into())]).unwrap(); + assert!(lex.spell(&ph(&["m", "uː"])).is_empty()); + } + + #[test] + fn lexicon_keys_are_case_sensitive() { + // IPA case is significant: "ʃ" ≠ "S". + let mut lex = P2gLexicon::from_rows(vec![("ʃ iː".into(), "she".into())]).unwrap(); + assert_eq!(lex.spell(&ph(&["ʃ", "iː"])), ["she"]); + assert!(lex.spell(&ph(&["S", "iː"])).is_empty()); + } + + #[test] + fn chained_first_non_empty_wins() { + struct Never; + impl P2g for Never { + fn spell(&mut self, _phonemes: &[Phoneme]) -> Vec { + panic!("must not be consulted when tier 1 hits"); + } + } + struct Always; + impl P2g for Always { + fn spell(&mut self, _phonemes: &[Phoneme]) -> Vec { + vec!["from-tier-2".into()] + } + } + let mut hit = ChainedP2g( + P2gLexicon::from_rows(vec![("k æ t".into(), "cat".into())]).unwrap(), + Never, + ); + assert_eq!(hit.spell(&ph(&["k", "æ", "t"])), ["cat"]); + + let mut miss = ChainedP2g( + P2gLexicon::from_rows(vec![("k æ t".into(), "cat".into())]).unwrap(), + Always, + ); + assert_eq!(miss.spell(&ph(&["m", "uː"])), ["from-tier-2"]); + } +} diff --git a/crates/floravox-g2p/src/phonetisaurus.rs b/crates/floravox-g2p/src/phonetisaurus.rs index ca76ba9..fe5286d 100644 --- a/crates/floravox-g2p/src/phonetisaurus.rs +++ b/crates/floravox-g2p/src/phonetisaurus.rs @@ -418,6 +418,37 @@ fn parse_symbols(text: &str) -> Vec<(String, i32)> { /// Phonetisaurus WFST G2P engine: lexicon-free grapheme→phoneme /// transcription via shortest-path over the model transducer. +/// Read `stem.fst` and resolve its symbol tables (embedded, or from +/// `stem.grapheme.table` / `stem.phoneme.table`). Shared by the G2P and +/// P2G readers; the table field names are G2P-centric ("graphemes" = +/// input side, "phonemes" = output side) whatever the model direction. +fn load_parsed(stem: &Path) -> Result { + let fst_path = with_ext(stem, "fst"); + let fst_bytes = std::fs::read(&fst_path).map_err(G2pError::Open)?; + let mut parsed = parse_model(&fst_bytes)?; + if !parsed.embedded_input { + let text = std::fs::read_to_string(with_ext(&stem_without_ext(stem), "grapheme.table")) + .map_err(|e| { + G2pError::Open(std::io::Error::other(format!( + "no embedded input table and no {}: {e}", + with_ext(&stem_without_ext(stem), "grapheme.table").display() + ))) + })?; + parsed.tables.graphemes = parse_symbols(&text); + } + if !parsed.embedded_output { + let text = std::fs::read_to_string(with_ext(&stem_without_ext(stem), "phoneme.table")) + .map_err(|e| { + G2pError::Open(std::io::Error::other(format!( + "no embedded output table and no {}: {e}", + with_ext(&stem_without_ext(stem), "phoneme.table").display() + ))) + })?; + parsed.tables.phonemes = parse_symbols(&text); + } + Ok(parsed) +} + pub struct PhonetisaurusG2p { fst: VectorFst, /// Grapheme string (compounds normalized: `a|c` → `ac`) → symbol id. @@ -442,30 +473,7 @@ impl PhonetisaurusG2p { /// [`G2pError::Open`] for filesystem failures; [`G2pError::Compile`] /// for malformed FST or symbol-table data. pub fn open(stem: impl AsRef) -> Result { - let stem = stem.as_ref(); - let fst_path = with_ext(stem, "fst"); - let fst_bytes = std::fs::read(&fst_path).map_err(G2pError::Open)?; - let mut parsed = parse_model(&fst_bytes)?; - if !parsed.embedded_input { - let text = std::fs::read_to_string(with_ext(&stem_without_ext(stem), "grapheme.table")) - .map_err(|e| { - G2pError::Open(std::io::Error::other(format!( - "no embedded input table and no {}: {e}", - with_ext(&stem_without_ext(stem), "grapheme.table").display() - ))) - })?; - parsed.tables.graphemes = parse_symbols(&text); - } - if !parsed.embedded_output { - let text = std::fs::read_to_string(with_ext(&stem_without_ext(stem), "phoneme.table")) - .map_err(|e| { - G2pError::Open(std::io::Error::other(format!( - "no embedded output table and no {}: {e}", - with_ext(&stem_without_ext(stem), "phoneme.table").display() - ))) - })?; - parsed.tables.phonemes = parse_symbols(&text); - } + let parsed = load_parsed(stem.as_ref())?; Self::build(parsed.fst, &parsed.tables) } @@ -748,6 +756,241 @@ fn with_ext(path: &Path, ext: &str) -> PathBuf { PathBuf::from(s) } +/// Phonetisaurus-style **phoneme-to-grapheme** engine: loads a model +/// trained with `floravox-train-phonetisaurus --reverse` (phoneme +/// compounds on the input side, grapheme compounds on the output side) +/// and spells pronunciations through the same shortest-path search. +/// +/// The same file opened by [`PhonetisaurusG2p`] performs G2P — the +/// direction is a property of how the model was trained, and each +/// reader interprets the tables for its direction. +/// +/// Input symbols follow the joint-segment convention on the phoneme +/// side: `T|AH0` covers two consecutive phonemes, `_` aligns to nothing. +/// Output symbols are grapheme compounds (`c|a|t`); `|` is stripped when +/// concatenating the spelled word. +pub struct PhonetisaurusP2g { + fst: VectorFst, + /// Phoneme compound (verbatim, `|`-joined) → input symbol id. + inputs: HashMap, + /// Output symbol id → grapheme compound string. + outputs: Vec, + /// Longest input compound in phoneme units (bounds segmentation). + max_input_units: usize, +} + +impl PhonetisaurusP2g { + /// Load a reversed model (embedded tables: bare `model.fst`; external + /// tables: `model.grapheme.table` / `model.phoneme.table` beside it). + /// + /// # Errors + /// + /// [`G2pError::Open`] for filesystem failures; [`G2pError::Compile`] + /// for malformed FST or symbol-table data. + pub fn open(stem: impl AsRef) -> Result { + let parsed = load_parsed(stem.as_ref())?; + Self::build(parsed.fst, &parsed.tables) + } + + /// Build from raw parts: FST bytes plus external text symbol tables + /// (input table = phoneme compounds, output table = grapheme + /// compounds). + /// + /// # Errors + /// + /// [`G2pError::Compile`] for malformed FST or symbol-table data. + pub fn from_parts( + fst_bytes: &[u8], + input_table: &str, + output_table: &str, + ) -> Result { + let mut parsed = parse_model(fst_bytes)?; + if parsed.embedded_input { + return Err(G2pError::Compile( + "FST embeds its input table; pass empty external tables".into(), + )); + } + if parsed.embedded_output { + return Err(G2pError::Compile( + "FST embeds its output table; pass empty external tables".into(), + )); + } + parsed.tables.graphemes = parse_symbols(input_table); + parsed.tables.phonemes = parse_symbols(output_table); + Self::build(parsed.fst, &parsed.tables) + } + + fn build(fst: VectorFst, tables: &SymbolTables) -> Result { + let mut inputs = HashMap::new(); + let mut max_input_units = 1; + for (sym, id) in &tables.graphemes { + if *id <= 2 { + continue; // epsilon, '|' separator, '_' null marker + } + // Phoneme compounds keep their '|' joins verbatim — the search + // builds candidate keys by joining units with '|'. + let units = sym.split('|').count(); + max_input_units = max_input_units.max(units); + inputs.insert(sym.clone(), *id); + } + let mut outputs: Vec = Vec::new(); + for (sym, id) in &tables.phonemes { + let Some(idx) = usize::try_from(*id).ok() else { + continue; + }; + if outputs.len() <= idx { + outputs.resize(idx + 1, String::new()); + } + outputs[idx].clone_from(sym); + } + if inputs.is_empty() { + return Err(G2pError::Compile("empty input table".into())); + } + Ok(Self { + fst, + inputs, + outputs, + max_input_units, + }) + } + + /// Number of model states (diagnostics). + #[must_use] + pub fn num_states(&self) -> usize { + self.fst.num_states() + } + + /// Number of model arcs (diagnostics). + #[must_use] + pub fn num_arcs(&self) -> usize { + self.fst.num_arcs() + } + + /// Spell `phonemes` (`None` when no path exists — unknown phonemes, + /// search-cap exhaustion). Multi-phoneme input symbols + /// (`T|AH0` compounds) are considered at every position; `_` and + /// epsilon input arcs emit without consuming. + #[must_use] + pub fn spell(&self, phonemes: &[Phoneme]) -> Option { + if phonemes.is_empty() { + return Some(String::new()); + } + let units: Vec = phonemes.to_vec(); + let (dist, best) = self.search(&units)?; + let mut node = best; + let mut labels: Vec = Vec::new(); + for _ in 0..MAX_BACKTRACK { + let data = dist[&node]; + let Some((pred, olabel)) = data.pred else { + break; + }; + if olabel > 2 { + labels.push(olabel); + } + node = pred; + } + labels.reverse(); + let mut word = String::new(); + for id in labels { + if let Some(sym) = usize::try_from(id).ok().and_then(|i| self.outputs.get(i)) { + // Strip the compound separator: `c|a|t` → `cat`, `_` + // already filtered by the olabel > 2 guard above. + word.push_str(&sym.replace('|', "")); + } + } + (!word.is_empty()).then_some(word) + } + + /// Same Dijkstra label-correcting search as the G2P decoder, over + /// phoneme units: candidate input segments are `|`-joins of + /// consecutive units matched against the input table. + fn search(&self, units: &[String]) -> Option<(HashMap, u64)> { + let start = self.fst.start?; + let end = units.len(); + let mut dist: HashMap = HashMap::new(); + let mut heap: BinaryHeap> = BinaryHeap::new(); + dist.insert( + pack(0, start), + NodeData { + cost: 0.0, + pred: None, + }, + ); + heap.push(Reverse((Cost(0.0), pack(0, start)))); + + let mut best_complete: Option<(f32, u64)> = None; + let mut pops = 0usize; + while let Some(Reverse((c, node))) = heap.pop() { + let data = dist[&node]; + if c.0 > data.cost { + continue; // stale entry + } + pops += 1; + if pops > MAX_POPS { + break; + } + let (in_pos, state) = unpack(node); + if in_pos == end { + if let Some(fw) = self.fst.final_weight(state) { + let total = data.cost + fw; + if best_complete.is_none_or(|(b, _)| total < b) { + best_complete = Some((total, node)); + } + } + } + + // Input segments available at this position: epsilon and the + // '_' null marker (both consume nothing), plus every + // '|'-joined table prefix of the remaining units. + let mut segs: Vec<(usize, i32)> = vec![(in_pos, 0), (in_pos, 2)]; + for k in 1..=self.max_input_units { + if in_pos + k > end { + break; + } + let s = units[in_pos..in_pos + k].join("|"); + if let Some(&id) = self.inputs.get(&s) { + segs.push((in_pos + k, id)); + } + } + + for arc in self.fst.arcs(state) { + for &(new_pos, id) in &segs { + if arc.ilabel != 0 && arc.ilabel != id { + continue; + } + let next = pack(new_pos, arc.nextstate); + let cost = data.cost + arc.weight; + let better = match dist.get(&next) { + Some(d) => cost < d.cost, + None => true, + }; + if better { + dist.insert( + next, + NodeData { + cost, + pred: Some((node, arc.olabel)), + }, + ); + heap.push(Reverse((Cost(cost), next))); + } + } + } + } + best_complete.map(|(_, node)| (dist, node)) + } +} + +impl crate::p2g::P2g for PhonetisaurusP2g { + fn spell(&mut self, phonemes: &[Phoneme]) -> Vec { + // Disambiguate from this same-named trait method: call the + // inherent 1-best decoder explicitly. + PhonetisaurusP2g::spell(self, phonemes) + .into_iter() + .collect() + } +} + #[cfg(test)] mod tests { use super::*; @@ -808,6 +1051,70 @@ mod tests { const GRAPHEMES: &str = " 0\n| 1\n_ 2\nH 3\nI 4\nHI 5\n"; const PHONEMES: &str = " 0\n| 1\n_ 2\nh 3\naɪ 4\nhaɪ 5\n"; + // ---- PhonetisaurusP2g ------------------------------------------------- + + // Reversed-model tables: input side = phonemes (with compounds like + // `m|uː`), output side = graphemes (`m|oo` compounds). + const P2G_INPUT: &str = " 0\n| 1\n_ 2\nk 3\næ 4\nt 5\nm 6\nuː 7\nm|uː 8\n"; + const P2G_OUTPUT: &str = " 0\n| 1\n_ 2\nc 3\na 4\nt 5\nm 6\nmoo 7\n"; + + /// /k æ t/ → cat (three single-phoneme arcs); /m uː/ → moo via the + /// compound input `m|uː` emitting the compound output `moo`. + fn p2g_model() -> PhonetisaurusP2g { + let bytes = write_fst( + 0, + &[ + (0, 3, 3, 0.1, 1), // k -> c + (1, 4, 4, 0.1, 2), // æ -> a + (2, 5, 5, 0.1, 3), // t -> t + (0, 8, 7, 0.05, 4), // m|uː -> moo (compound, cheaper) + (0, 6, 6, 0.5, 5), // m -> m (fallback path pieces) + (5, 7, 0, 0.5, 6), // uː -> _ (null output) + (6, 0, 0, 0.1, 4), // eps hop into the final state + ], + &[(3, 0.0), (4, 0.0)], + false, + ); + PhonetisaurusP2g::from_parts(&bytes, P2G_INPUT, P2G_OUTPUT).unwrap() + } + + fn ph(list: &[&str]) -> Vec { + list.iter().map(|&s| s.to_owned()).collect() + } + + #[test] + fn p2g_spells_known_pronunciation() { + let m = p2g_model(); + assert_eq!(m.spell(&ph(&["k", "æ", "t"])).as_deref(), Some("cat")); + } + + #[test] + fn p2g_prefers_compound_input_segment() { + // m|uː (0.05) beats m + uː/_ (0.5 + 0.5 + 0.1): "moo". + let m = p2g_model(); + assert_eq!(m.spell(&ph(&["m", "uː"])).as_deref(), Some("moo")); + } + + #[test] + fn p2g_unknown_phoneme_is_none() { + let m = p2g_model(); + assert!(m.spell(&ph(&["q", "x"])).is_none()); + assert!(m.spell(&ph(&["k", "æ"])).is_none()); // incomplete path + } + + #[test] + fn p2g_trait_impl_returns_candidates() { + use crate::p2g::P2g; + let mut m = p2g_model(); + // Call through the trait: the inherent 1-best `spell` shadows it + // for method-call syntax on the concrete type. + assert_eq!( + P2g::spell(&mut m, &ph(&["k", "æ", "t"])), + ["cat".to_string()] + ); + assert!(P2g::spell(&mut m, &ph(&["z"])).is_empty()); + } + // States: 0 --H/h--> 1 --eps--> 2 --I/aɪ--> 3(final) // 0 --HI/haɪ---------(weight varies)--> 3 // 1 --_/(compound out)--> ... tests null markers