From bc723ea200144bd6fa8a5dff4b9a780feda144fc Mon Sep 17 00:00:00 2001
From: 游雁 <zhifu.gzf@alibaba-inc.com>
Date: 星期四, 29 六月 2023 18:55:01 +0800
Subject: [PATCH] dcos

---
 funasr/runtime/onnxruntime/src/punc-model.cpp |   15 ++++++++++++---
 1 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/funasr/runtime/onnxruntime/src/punc-model.cpp b/funasr/runtime/onnxruntime/src/punc-model.cpp
index 1e619ab..54b8d6a 100644
--- a/funasr/runtime/onnxruntime/src/punc-model.cpp
+++ b/funasr/runtime/onnxruntime/src/punc-model.cpp
@@ -1,10 +1,17 @@
 #include "precomp.h"
 
-PuncModel *CreatePuncModel(std::map<std::string, std::string>& model_path, int thread_num)
+namespace funasr {
+PuncModel *CreatePuncModel(std::map<std::string, std::string>& model_path, int thread_num, PUNC_TYPE type)
 {
     PuncModel *mm;
-    mm = new CTTransformer();
-
+    if (type==PUNC_OFFLINE){
+        mm = new CTTransformer();
+    }else if(type==PUNC_ONLINE){
+        mm = new CTTransformerOnline();
+    }else{
+        LOG(ERROR) << "Wrong PUNC TYPE";
+        exit(-1);
+    }
     string punc_model_path;
     string punc_config_path;
 
@@ -17,3 +24,5 @@
     mm->InitPunc(punc_model_path, punc_config_path, thread_num);
     return mm;
 }
+
+} // namespace funasr

--
Gitblit v1.9.1