add destination for model option

neurlang · Oct 13, 2024 · 581baae · 581baae
1 parent 12806c2
commit 581baae
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/cmd/train_phonemizer/main.go b/cmd/train_phonemizer/main.go
@@ -21,8 +21,11 @@ func error_abs(a, b uint32) uint32 {
 
 func main() {
 	cleantsv := flag.String("cleantsv", "", "clean tsv dataset for the language")
+	dstmodel := flag.String("dstmodel", "", "model destination .json.lzw file")
 	flag.Parse()
 
+	var improved_success_rate = 0
+
 	if cleantsv == nil || *cleantsv == "" {
 		println("clean tsv is mandatory")
 		return
@@ -133,6 +136,14 @@ func main() {
 			println(err.Error())
 		}
 
+		if dstmodel != nil && len(*dstmodel) > 0 && improved_success_rate < success {
+			improved_success_rate = success
+			err := net.WriteCompressedWeightsToFile(*dstmodel)
+			if err != nil {
+				println(err.Error())
+			}
+		}
+
 		if success == 100 {
 			println("Max accuracy or wrong data. Exiting")
 			os.Exit(0)