option to not use dictionary

Mobile-Artificial-Intelligence · Aug 13, 2024 · ba61245 · ba61245
1 parent 7bdc946
commit ba61245
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 25 deletions.
diff --git a/include/babylon.h b/include/babylon.h
@@ -5,13 +5,21 @@
 extern "C" {
 #endif
 
+#include <stdbool.h>
+
 #ifdef WIN32
    #define BABYLON_EXPORT __declspec(dllexport)
 #else
    #define BABYLON_EXPORT __attribute__((visibility("default"))) __attribute__((used))
 #endif
 
-BABYLON_EXPORT int babylon_g2p_init(const char* model_path, const char* language, int use_punctuation);
+typedef struct {
+   const char* language;
+   bool use_punctuation;
+   bool use_dictionary;
+} babylon_g2p_options;
+
+BABYLON_EXPORT int babylon_g2p_init(const char* model_path, babylon_g2p_options* options);
 
 BABYLON_EXPORT char* babylon_g2p(const char* text);
 

diff --git a/include/babylon.hpp b/include/babylon.hpp
@@ -32,7 +32,12 @@ namespace DeepPhonemizer {
 
   class Session {
     public:
-      Session(const std::string& model_path, const std::string language = "en_us", const bool use_punctuation = false);
+      Session(
+        const std::string& model_path, 
+        const std::string language = "en_us", 
+        const bool use_punctuation = false, 
+        const bool use_dictionary = false
+      );
       ~Session();
 
       std::vector<std::string> g2p(const std::string& text);

diff --git a/src/babylon.cpp b/src/babylon.cpp
@@ -6,9 +6,9 @@ static DeepPhonemizer::Session* dp;
 static Vits::Session* vits;
 
 extern "C" {
-    BABYLON_EXPORT int babylon_g2p_init(const char* model_path, const char* language, int use_punctuation) {
+    BABYLON_EXPORT int babylon_g2p_init(const char* model_path, babylon_g2p_options* options) {
         try {
-            dp = new DeepPhonemizer::Session(model_path, language, use_punctuation);
+            dp = new DeepPhonemizer::Session(model_path, options->language, options->use_punctuation);
             return 0;
         } 
         catch (const std::exception& e) {

diff --git a/src/cleaners.cpp b/src/cleaners.cpp
@@ -182,10 +182,6 @@ namespace DeepPhonemizer {
             }
         }
 
-        if (!word.empty()) {
-            words.push_back(word);
-        }
-
         return words;
     }
 }
diff --git a/src/phonemizer.cpp b/src/phonemizer.cpp
@@ -129,7 +129,29 @@ namespace DeepPhonemizer {
         return probabilities;
     }
 
-    Session::Session(const std::string& model_path, const std::string language, const bool use_punctuation) {
+    std::unordered_map<std::string, std::vector<std::string>> parse_dictionary(const std::string& dictionary_str) {
+        std::unordered_map<std::string, std::vector<std::string>> dictionary;
+
+        std::istringstream dictionary_stream(dictionary_str);
+        std::string line;
+        while (std::getline(dictionary_stream, line)) {
+            std::stringstream line_stream(line);
+            std::string word;
+            line_stream >> word;
+
+            std::vector<std::string> phonemes;
+            std::string phoneme;
+            while (line_stream >> phoneme) {
+                phonemes.push_back(phoneme);
+            }
+
+            dictionary[word] = phonemes;
+        }
+
+        return dictionary;
+    }
+
+    Session::Session(const std::string& model_path, const std::string language, const bool use_punctuation, const bool use_dictionary) {
         Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "DeepPhonemizer");
         env.DisableTelemetryEvents();
 
@@ -170,22 +192,9 @@ namespace DeepPhonemizer {
             phoneme_symbols.push_back(phoneme_symbol_buffer);
         }
 
-        std::string dictonary_str = model_metadata.LookupCustomMetadataMapAllocated("dictionary", allocator).get();
-
-        std::istringstream dictionary_stream(dictonary_str);
-        std::string line;
-        while (std::getline(dictionary_stream, line)) {
-            std::stringstream line_stream(line);
-            std::string word;
-            line_stream >> word;
-
-            std::vector<std::string> phonemes;
-            std::string phoneme;
-            while (line_stream >> phoneme) {
-                phonemes.push_back(phoneme);
-            }
-
-            dictionary[word] = phonemes;
+        if (use_dictionary) {
+            std::string dictonary_str = model_metadata.LookupCustomMetadataMapAllocated("dictionary", allocator).get();
+            dictionary = parse_dictionary(dictonary_str);
         }
 
         int char_repeats = model_metadata.LookupCustomMetadataMapAllocated("char_repeats", allocator).get()[0] - '0';