Add CTC HLG decoding for JNI (#810)
This commit is contained in:
@@ -14,6 +14,7 @@ java_files += OnlineParaformerModelConfig.java
|
||||
java_files += OnlineZipformer2CtcModelConfig.java
|
||||
java_files += OnlineTransducerModelConfig.java
|
||||
java_files += OnlineModelConfig.java
|
||||
java_files += OnlineCtcFstDecoderConfig.java
|
||||
java_files += OnlineStream.java
|
||||
java_files += OnlineRecognizerConfig.java
|
||||
java_files += OnlineRecognizerResult.java
|
||||
@@ -22,6 +23,7 @@ java_files += OnlineRecognizer.java
|
||||
java_files += OfflineTransducerModelConfig.java
|
||||
java_files += OfflineParaformerModelConfig.java
|
||||
java_files += OfflineWhisperModelConfig.java
|
||||
java_files += OfflineNemoEncDecCtcModelConfig.java
|
||||
java_files += OfflineModelConfig.java
|
||||
java_files += OfflineRecognizerConfig.java
|
||||
java_files += OfflineRecognizerResult.java
|
||||
@@ -42,10 +44,12 @@ $(info -- class files $(class_files))
|
||||
all: $(out_jar)
|
||||
|
||||
$(out_jar): $(class_files)
|
||||
jar --create --verbose --file $(out_jar) -C $(out_dir) .
|
||||
# jar --create --verbose --file $(out_jar) -C $(out_dir) ./
|
||||
jar cvf $(out_jar) -C $(out_dir) ./
|
||||
|
||||
clean:
|
||||
$(RM) -rfv $(out_dir)
|
||||
|
||||
$(class_files): $(out_dir)/$(package_dir)/%.class: src/$(package_dir)/%.java
|
||||
javac -d $(out_dir) --class-path $(out_dir) $<
|
||||
mkdir -p build
|
||||
javac -d $(out_dir) -cp $(out_dir) $<
|
||||
|
||||
@@ -5,6 +5,7 @@ public class OfflineModelConfig {
|
||||
private final OfflineTransducerModelConfig transducer;
|
||||
private final OfflineParaformerModelConfig paraformer;
|
||||
private final OfflineWhisperModelConfig whisper;
|
||||
private final OfflineNemoEncDecCtcModelConfig nemo;
|
||||
private final String tokens;
|
||||
private final int numThreads;
|
||||
private final boolean debug;
|
||||
@@ -16,6 +17,7 @@ public class OfflineModelConfig {
|
||||
this.transducer = builder.transducer;
|
||||
this.paraformer = builder.paraformer;
|
||||
this.whisper = builder.whisper;
|
||||
this.nemo = builder.nemo;
|
||||
this.tokens = builder.tokens;
|
||||
this.numThreads = builder.numThreads;
|
||||
this.debug = builder.debug;
|
||||
@@ -64,6 +66,7 @@ public class OfflineModelConfig {
|
||||
private OfflineParaformerModelConfig paraformer = OfflineParaformerModelConfig.builder().build();
|
||||
private OfflineTransducerModelConfig transducer = OfflineTransducerModelConfig.builder().build();
|
||||
private OfflineWhisperModelConfig whisper = OfflineWhisperModelConfig.builder().build();
|
||||
private OfflineNemoEncDecCtcModelConfig nemo = OfflineNemoEncDecCtcModelConfig.builder().build();
|
||||
private String tokens = "";
|
||||
private int numThreads = 1;
|
||||
private boolean debug = true;
|
||||
@@ -84,6 +87,11 @@ public class OfflineModelConfig {
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder setNemo(OfflineNemoEncDecCtcModelConfig nemo) {
|
||||
this.nemo = nemo;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder setWhisper(OfflineWhisperModelConfig whisper) {
|
||||
this.whisper = whisper;
|
||||
return this;
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
// Copyright 2024 Xiaomi Corporation
|
||||
package com.k2fsa.sherpa.onnx;
|
||||
|
||||
public class OfflineNemoEncDecCtcModelConfig {
|
||||
private final String model;
|
||||
|
||||
private OfflineNemoEncDecCtcModelConfig(Builder builder) {
|
||||
this.model = builder.model;
|
||||
}
|
||||
|
||||
public static Builder builder() {
|
||||
return new Builder();
|
||||
}
|
||||
|
||||
public String getModel() {
|
||||
return model;
|
||||
}
|
||||
|
||||
public static class Builder {
|
||||
private String model = "";
|
||||
|
||||
public OfflineNemoEncDecCtcModelConfig build() {
|
||||
return new OfflineNemoEncDecCtcModelConfig(this);
|
||||
}
|
||||
|
||||
public Builder setModel(String model) {
|
||||
this.model = model;
|
||||
return this;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,43 @@
|
||||
// Copyright 2024 Xiaomi Corporation
|
||||
package com.k2fsa.sherpa.onnx;
|
||||
|
||||
public class OnlineCtcFstDecoderConfig {
|
||||
private final String graph;
|
||||
private final int maxActive;
|
||||
|
||||
private OnlineCtcFstDecoderConfig(Builder builder) {
|
||||
this.graph = builder.graph;
|
||||
this.maxActive = builder.maxActive;
|
||||
}
|
||||
|
||||
public static Builder builder() {
|
||||
return new Builder();
|
||||
}
|
||||
|
||||
public String getGraph() {
|
||||
return graph;
|
||||
}
|
||||
|
||||
public float getMaxActive() {
|
||||
return maxActive;
|
||||
}
|
||||
|
||||
public static class Builder {
|
||||
private String graph = "";
|
||||
private int maxActive = 3000;
|
||||
|
||||
public OnlineCtcFstDecoderConfig build() {
|
||||
return new OnlineCtcFstDecoderConfig(this);
|
||||
}
|
||||
|
||||
public Builder setGraph(String model) {
|
||||
this.graph = graph;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder setMaxActive(int maxActive) {
|
||||
this.maxActive = maxActive;
|
||||
return this;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,8 @@ public class OnlineRecognizerConfig {
|
||||
private final FeatureConfig featConfig;
|
||||
private final OnlineModelConfig modelConfig;
|
||||
private final OnlineLMConfig lmConfig;
|
||||
|
||||
private final OnlineCtcFstDecoderConfig ctcFstDecoderConfig;
|
||||
private final EndpointConfig endpointConfig;
|
||||
private final boolean enableEndpoint;
|
||||
private final String decodingMethod;
|
||||
@@ -17,6 +19,7 @@ public class OnlineRecognizerConfig {
|
||||
this.featConfig = builder.featConfig;
|
||||
this.modelConfig = builder.modelConfig;
|
||||
this.lmConfig = builder.lmConfig;
|
||||
this.ctcFstDecoderConfig = builder.ctcFstDecoderConfig;
|
||||
this.endpointConfig = builder.endpointConfig;
|
||||
this.enableEndpoint = builder.enableEndpoint;
|
||||
this.decodingMethod = builder.decodingMethod;
|
||||
@@ -37,6 +40,7 @@ public class OnlineRecognizerConfig {
|
||||
private FeatureConfig featConfig = FeatureConfig.builder().build();
|
||||
private OnlineModelConfig modelConfig = OnlineModelConfig.builder().build();
|
||||
private OnlineLMConfig lmConfig = OnlineLMConfig.builder().build();
|
||||
private OnlineCtcFstDecoderConfig ctcFstDecoderConfig = OnlineCtcFstDecoderConfig.builder().build();
|
||||
private EndpointConfig endpointConfig = EndpointConfig.builder().build();
|
||||
private boolean enableEndpoint = true;
|
||||
private String decodingMethod = "greedy_search";
|
||||
@@ -63,6 +67,11 @@ public class OnlineRecognizerConfig {
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder setCtcFstDecoderConfig(OnlineCtcFstDecoderConfig ctcFstDecoderConfig) {
|
||||
this.ctcFstDecoderConfig = ctcFstDecoderConfig;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder setEndpointConfig(EndpointConfig endpointConfig) {
|
||||
this.endpointConfig = endpointConfig;
|
||||
return this;
|
||||
|
||||
@@ -147,6 +147,19 @@ static OfflineRecognizerConfig GetOfflineConfig(JNIEnv *env, jobject config) {
|
||||
ans.model_config.whisper.tail_paddings =
|
||||
env->GetIntField(whisper_config, fid);
|
||||
|
||||
fid = env->GetFieldID(
|
||||
model_config_cls, "nemo",
|
||||
"Lcom/k2fsa/sherpa/onnx/OfflineNemoEncDecCtcModelConfig;");
|
||||
jobject nemo_config = env->GetObjectField(model_config, fid);
|
||||
jclass nemo_config_cls = env->GetObjectClass(nemo_config);
|
||||
|
||||
fid = env->GetFieldID(paraformer_config_cls, "model", "Ljava/lang/String;");
|
||||
|
||||
s = (jstring)env->GetObjectField(nemo_config, fid);
|
||||
p = env->GetStringUTFChars(s, nullptr);
|
||||
ans.model_config.nemo_ctc.model = p;
|
||||
env->ReleaseStringUTFChars(s, p);
|
||||
|
||||
return ans;
|
||||
}
|
||||
|
||||
|
||||
@@ -198,6 +198,22 @@ static OnlineRecognizerConfig GetConfig(JNIEnv *env, jobject config) {
|
||||
fid = env->GetFieldID(lm_model_config_cls, "scale", "F");
|
||||
ans.lm_config.scale = env->GetFloatField(lm_model_config, fid);
|
||||
|
||||
fid = env->GetFieldID(cls, "ctcFstDecoderConfig",
|
||||
"Lcom/k2fsa/sherpa/onnx/OnlineCtcFstDecoderConfig;");
|
||||
|
||||
jobject fst_decoder_config = env->GetObjectField(config, fid);
|
||||
jclass fst_decoder_config_cls = env->GetObjectClass(fst_decoder_config);
|
||||
|
||||
fid = env->GetFieldID(fst_decoder_config_cls, "graph", "Ljava/lang/String;");
|
||||
s = (jstring)env->GetObjectField(fst_decoder_config, fid);
|
||||
p = env->GetStringUTFChars(s, nullptr);
|
||||
ans.ctc_fst_decoder_config.graph = p;
|
||||
env->ReleaseStringUTFChars(s, p);
|
||||
|
||||
fid = env->GetFieldID(fst_decoder_config_cls, "maxActive", "I");
|
||||
ans.ctc_fst_decoder_config.max_active =
|
||||
env->GetIntField(fst_decoder_config, fid);
|
||||
|
||||
return ans;
|
||||
}
|
||||
} // namespace sherpa_onnx
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
#include <fstream>
|
||||
|
||||
#include "sherpa-onnx/csrc/macros.h"
|
||||
#include "sherpa-onnx/csrc/onnx-utils.h"
|
||||
#include "sherpa-onnx/jni/common.h"
|
||||
|
||||
static jobjectArray ReadWaveImpl(JNIEnv *env, std::istream &is,
|
||||
|
||||
@@ -18,6 +18,10 @@ data class OfflineParaformerModelConfig(
|
||||
var model: String = "",
|
||||
)
|
||||
|
||||
data class OfflineNemoEncDecCtcModelConfig(
|
||||
var model: String = "",
|
||||
)
|
||||
|
||||
data class OfflineWhisperModelConfig(
|
||||
var encoder: String = "",
|
||||
var decoder: String = "",
|
||||
@@ -30,6 +34,7 @@ data class OfflineModelConfig(
|
||||
var transducer: OfflineTransducerModelConfig = OfflineTransducerModelConfig(),
|
||||
var paraformer: OfflineParaformerModelConfig = OfflineParaformerModelConfig(),
|
||||
var whisper: OfflineWhisperModelConfig = OfflineWhisperModelConfig(),
|
||||
var nemo: OfflineNemoEncDecCtcModelConfig = OfflineNemoEncDecCtcModelConfig(),
|
||||
var numThreads: Int = 1,
|
||||
var debug: Boolean = false,
|
||||
var provider: String = "cpu",
|
||||
@@ -216,6 +221,16 @@ fun getOfflineModelConfig(type: Int): OfflineModelConfig? {
|
||||
)
|
||||
}
|
||||
|
||||
6 -> {
|
||||
val modelDir = "sherpa-onnx-nemo-ctc-en-citrinet-512"
|
||||
return OfflineModelConfig(
|
||||
nemo = OfflineNemoEncDecCtcModelConfig(
|
||||
model = "$modelDir/model.int8.onnx",
|
||||
),
|
||||
tokens = "$modelDir/tokens.txt",
|
||||
)
|
||||
}
|
||||
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
@@ -45,11 +45,17 @@ data class OnlineLMConfig(
|
||||
var scale: Float = 0.5f,
|
||||
)
|
||||
|
||||
data class OnlineCtcFstDecoderConfig(
|
||||
var graph: String = "",
|
||||
var maxActive: Int = 3000,
|
||||
)
|
||||
|
||||
|
||||
data class OnlineRecognizerConfig(
|
||||
var featConfig: FeatureConfig = FeatureConfig(),
|
||||
var modelConfig: OnlineModelConfig,
|
||||
var lmConfig: OnlineLMConfig = OnlineLMConfig(),
|
||||
var ctcFstDecoderConfig : OnlineCtcFstDecoderConfig = OnlineCtcFstDecoderConfig(),
|
||||
var endpointConfig: EndpointConfig = EndpointConfig(),
|
||||
var enableEndpoint: Boolean = true,
|
||||
var decodingMethod: String = "greedy_search",
|
||||
|
||||
Reference in New Issue
Block a user