feat: add utils to export the checkpoint into a traced torchscript checkpoint

1 files changed, 10 insertions, 0 deletions

diff --git a/utils/export.py b/utils/export.py
new file mode 100755
index 0000000..e7043ce
--- /dev/null
+++ b/utils/export.py
@@ -0,0 +1,10 @@
+#!/bin/env python3
+import torch
+from dp.phonemizer import Phonemizer
+from dp.preprocessing.text import Preprocessor, LanguageTokenizer, SequenceTokenizer
+torch.serialization.add_safe_globals([Preprocessor, LanguageTokenizer, SequenceTokenizer])
+
+phonemizer = Phonemizer.from_checkpoint('data/latin_ipa_forward.pt')
+model = phonemizer.predictor.model
+save = torch.jit.script(model)
+save.save('data/model.pt')
\ No newline at end of file