Skip to content

Commit a2b64a2

Browse files
pkufoolpkufool
committed
Add hotwords for ctc prefix beam search
1 parent 34aa323 commit a2b64a2

10 files changed

Lines changed: 166 additions & 20 deletions

‎sherpa-onnx/csrc/hypothesis.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,9 @@ struct Hypothesis {
8585
const ContextState *context_state = nullptr)
8686
: ys(ys), log_prob(log_prob), context_state(context_state) {}
8787

88+
explicit Hypothesis(const ContextState *context_state)
89+
: context_state(context_state) {}
90+
8891
double TotalLogProb(bool use_ctc = false) const {
8992
return LogProb(use_ctc) + lm_log_prob;
9093
}

‎sherpa-onnx/csrc/offline-ctc-decoder.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include <vector>
99

1010
#include "onnxruntime_cxx_api.h" // NOLINT
11+
#include "sherpa-onnx/csrc/offline-stream.h"
1112

1213
namespace sherpa_onnx {
1314

@@ -42,7 +43,8 @@ class OfflineCtcDecoder {
4243
* @return Return a vector of size `N` containing the decoded results.
4344
*/
4445
virtual std::vector<OfflineCtcDecoderResult> Decode(
45-
Ort::Value log_probs, Ort::Value log_probs_length) = 0;
46+
Ort::Value log_probs, Ort::Value log_probs_length,
47+
OfflineStream **ss = nullptr, int32_t n = 0) = 0;
4648
};
4749

4850
} // namespace sherpa_onnx

‎sherpa-onnx/csrc/offline-ctc-fst-decoder.cc‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,8 @@ OfflineCtcFstDecoder::OfflineCtcFstDecoder(
8484
: config_(config), fst_(ReadGraph(config_.graph)) {}
8585

8686
std::vector<OfflineCtcDecoderResult> OfflineCtcFstDecoder::Decode(
87-
Ort::Value log_probs, Ort::Value log_probs_length) {
87+
Ort::Value log_probs, Ort::Value log_probs_length,
88+
OfflineStream **ss /*= nullptr*/, int32_t n /*= 0*/) {
8889
std::vector<int64_t> shape = log_probs.GetTensorTypeAndShapeInfo().GetShape();
8990

9091
assert(static_cast<int32_t>(shape.size()) == 3);

‎sherpa-onnx/csrc/offline-ctc-fst-decoder.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,10 @@ class OfflineCtcFstDecoder : public OfflineCtcDecoder {
1919
public:
2020
explicit OfflineCtcFstDecoder(const OfflineCtcFstDecoderConfig &config);
2121

22-
std::vector<OfflineCtcDecoderResult> Decode(
23-
Ort::Value log_probs, Ort::Value log_probs_length) override;
22+
std::vector<OfflineCtcDecoderResult> Decode(Ort::Value log_probs,
23+
Ort::Value log_probs_length,
24+
OfflineStream **ss = nullptr,
25+
int32_t n = 0) override;
2426

2527
private:
2628
OfflineCtcFstDecoderConfig config_;

‎sherpa-onnx/csrc/offline-ctc-greedy-search-decoder.cc‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,8 @@
1313
namespace sherpa_onnx {
1414

1515
std::vector<OfflineCtcDecoderResult> OfflineCtcGreedySearchDecoder::Decode(
16-
Ort::Value log_probs, Ort::Value log_probs_length) {
16+
Ort::Value log_probs, Ort::Value log_probs_length,
17+
OfflineStream **ss /*= nullptr*/, int32_t n /*= 0*/) {
1718
std::vector<int64_t> shape = log_probs.GetTensorTypeAndShapeInfo().GetShape();
1819
int32_t batch_size = static_cast<int32_t>(shape[0]);
1920
int32_t num_frames = static_cast<int32_t>(shape[1]);

‎sherpa-onnx/csrc/offline-ctc-greedy-search-decoder.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,10 @@ class OfflineCtcGreedySearchDecoder : public OfflineCtcDecoder {
1616
explicit OfflineCtcGreedySearchDecoder(int32_t blank_id)
1717
: blank_id_(blank_id) {}
1818

19-
std::vector<OfflineCtcDecoderResult> Decode(
20-
Ort::Value log_probs, Ort::Value log_probs_length) override;
19+
std::vector<OfflineCtcDecoderResult> Decode(Ort::Value log_probs,
20+
Ort::Value log_probs_length,
21+
OfflineStream **ss = nullptr,
22+
int32_t n = 0) override;
2123

2224
private:
2325
int32_t blank_id_;

‎sherpa-onnx/csrc/offline-ctc-prefix-beam-search-decoder.cc‎

Lines changed: 36 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
#include <utility>
99
#include <vector>
1010

11+
#include "sherpa-onnx/csrc/context-graph.h"
1112
#include "sherpa-onnx/csrc/hypothesis.h"
1213
#include "sherpa-onnx/csrc/macros.h"
1314

@@ -16,14 +17,16 @@ namespace sherpa_onnx {
1617
static std::vector<Hypothesis> StepWorker(const float *p_log_probs,
1718
std::vector<Hypothesis> &hyps,
1819
int32_t blank_id, int32_t vocab_size,
19-
int32_t max_active_paths) {
20+
int32_t max_active_paths,
21+
const ContextGraph *context_graph) {
2022
auto topk = TopkIndex(p_log_probs, vocab_size, max_active_paths);
2123
Hypotheses next_hyps;
2224
for (auto &hyp : hyps) {
23-
Hypothesis new_hyp = hyp;
2425
for (auto k : topk) {
26+
Hypothesis new_hyp = hyp;
2527
int32_t new_token = k;
2628
float log_prob = p_log_probs[k];
29+
bool update_prefix = false;
2730
if (new_token == blank_id) {
2831
// Case 0: *a + ε => *a
2932
// *aε + ε => *a
@@ -36,7 +39,6 @@ static std::vector<Hypothesis> StepWorker(const float *p_log_probs,
3639
// Prefix does not change, update log_prob of non_blank
3740
new_hyp.log_prob_nb = hyp.log_prob_nb + log_prob;
3841
new_hyp.log_prob_b = -std::numeric_limits<float>::infinity();
39-
4042
next_hyps.Add(std::move(new_hyp));
4143

4244
// Case 2: *aε + a => *aa
@@ -45,14 +47,26 @@ static std::vector<Hypothesis> StepWorker(const float *p_log_probs,
4547
new_hyp.ys.push_back(new_token);
4648
new_hyp.log_prob_nb = hyp.log_prob_b + log_prob;
4749
new_hyp.log_prob_b = -std::numeric_limits<float>::infinity();
48-
next_hyps.Add(std::move(new_hyp));
50+
update_prefix = true;
4951
} else {
5052
// Case 3: *a + b => *ab, *aε + b => *ab
5153
// Prefix changes, update log_prob of non_blank
5254
// Caution: DO NOT use append, as clone is shallow copy
5355
new_hyp.ys.push_back(new_token);
5456
new_hyp.log_prob_nb = hyp.LogProb(true) + log_prob;
5557
new_hyp.log_prob_b = -std::numeric_limits<float>::infinity();
58+
update_prefix = true;
59+
}
60+
61+
if (update_prefix) {
62+
float lm_log_prob = hyp.lm_log_prob;
63+
if (context_graph != nullptr && hyp.context_state != nullptr) {
64+
auto context_res =
65+
context_graph->ForwardOneStep(hyp.context_state, new_token);
66+
lm_log_prob = lm_log_prob + std::get<0>(context_res);
67+
new_hyp.context_state = std::get<1>(context_res);
68+
}
69+
new_hyp.lm_log_prob = lm_log_prob;
5670
next_hyps.Add(std::move(new_hyp));
5771
}
5872
}
@@ -61,7 +75,8 @@ static std::vector<Hypothesis> StepWorker(const float *p_log_probs,
6175
}
6276

6377
std::vector<OfflineCtcDecoderResult> OfflineCtcPrefixBeamSearchDecoder::Decode(
64-
Ort::Value log_probs, Ort::Value log_probs_length) {
78+
Ort::Value log_probs, Ort::Value log_probs_length,
79+
OfflineStream **ss /*= nullptr*/, int32_t n /*= 0*/) {
6580
std::vector<int64_t> shape = log_probs.GetTensorTypeAndShapeInfo().GetShape();
6681
int32_t batch_size = static_cast<int32_t>(shape[0]);
6782
int32_t num_frames = static_cast<int32_t>(shape[1]);
@@ -75,8 +90,17 @@ std::vector<OfflineCtcDecoderResult> OfflineCtcPrefixBeamSearchDecoder::Decode(
7590
std::vector<std::vector<Hypothesis>> cur;
7691
cur.reserve(batch_size);
7792

93+
std::vector<ContextGraphPtr> context_graphs(batch_size, nullptr);
94+
7895
for (int32_t i = 0; i < batch_size; ++i) {
79-
cur.emplace_back(std::vector<Hypothesis>({Hypothesis()}));
96+
const ContextState *context_state = nullptr;
97+
if (ss != nullptr) {
98+
context_graphs[i] = ss[i]->GetContextGraph();
99+
if (context_graphs[i] != nullptr)
100+
context_state = context_graphs[i]->Root();
101+
}
102+
Hypothesis hyp(context_state);
103+
cur.emplace_back(std::vector<Hypothesis>({hyp}));
80104
}
81105

82106
for (int32_t t = 0; t < num_frames; ++t) {
@@ -85,7 +109,12 @@ std::vector<OfflineCtcDecoderResult> OfflineCtcPrefixBeamSearchDecoder::Decode(
85109
const float *p_log_probs = log_probs.GetTensorData<float>() +
86110
b * num_frames * vocab_size + t * vocab_size;
87111
cur[b] = StepWorker(p_log_probs, cur[b], blank_id_, vocab_size,
88-
max_active_paths_);
112+
max_active_paths_, context_graphs[b].get());
113+
// for (auto &x : cur[b]) {
114+
// SHERPA_ONNX_LOGE("step : %d, key : %s, ac : %f, lm : %f", t,
115+
// x.Key().c_str(), x.LogProb(true), x.lm_log_prob);
116+
// }
117+
// SHERPA_ONNX_LOGE("\n");
89118
}
90119
}
91120
}

‎sherpa-onnx/csrc/offline-ctc-prefix-beam-search-decoder.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,10 @@ class OfflineCtcPrefixBeamSearchDecoder : public OfflineCtcDecoder {
1616
OfflineCtcPrefixBeamSearchDecoder(int32_t max_active_paths, int32_t blank_id)
1717
: max_active_paths_(max_active_paths), blank_id_(blank_id) {}
1818

19-
std::vector<OfflineCtcDecoderResult> Decode(
20-
Ort::Value log_probs, Ort::Value log_probs_length) override;
19+
std::vector<OfflineCtcDecoderResult> Decode(Ort::Value log_probs,
20+
Ort::Value log_probs_length,
21+
OfflineStream **ss = nullptr,
22+
int32_t n = 0) override;
2123

2224
private:
2325
int32_t max_active_paths_;

‎sherpa-onnx/csrc/offline-recognizer-ctc-impl.h‎

Lines changed: 106 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77

88
#include <ios>
99
#include <memory>
10+
#include <regex> // NOLINT
1011
#include <sstream>
1112
#include <string>
1213
#include <utility>
@@ -21,6 +22,8 @@
2122
#include "sherpa-onnx/csrc/offline-recognizer-impl.h"
2223
#include "sherpa-onnx/csrc/pad-sequence.h"
2324
#include "sherpa-onnx/csrc/symbol-table.h"
25+
#include "sherpa-onnx/csrc/utils.h"
26+
#include "ssentencepiece/csrc/ssentencepiece.h"
2427

2528
namespace sherpa_onnx {
2629

@@ -207,6 +210,15 @@ class OfflineRecognizerCtcImpl : public OfflineRecognizerImpl {
207210
if (config_.decoding_method == "greedy_search") {
208211
decoder_ = std::make_unique<OfflineCtcGreedySearchDecoder>(blank_id);
209212
} else {
213+
if (!config_.model_config.bpe_vocab.empty()) {
214+
bpe_encoder_ = std::make_unique<ssentencepiece::Ssentencepiece>(
215+
config_.model_config.bpe_vocab);
216+
}
217+
218+
if (!config_.hotwords_file.empty()) {
219+
InitHotwords();
220+
}
221+
210222
decoder_ = std::make_unique<OfflineCtcPrefixBeamSearchDecoder>(
211223
config_.max_active_paths, blank_id);
212224
}
@@ -219,9 +231,47 @@ class OfflineRecognizerCtcImpl : public OfflineRecognizerImpl {
219231
}
220232
}
221233

234+
std::unique_ptr<OfflineStream> CreateStream(
235+
const std::string &hotwords) const override {
236+
auto hws = std::regex_replace(hotwords, std::regex("/"), "\n");
237+
std::istringstream is(hws);
238+
std::vector<std::vector<int32_t>> current;
239+
std::vector<float> current_scores;
240+
if (!EncodeHotwords(is, config_.model_config.modeling_unit, symbol_table_,
241+
bpe_encoder_.get(), &current, &current_scores)) {
242+
SHERPA_ONNX_LOGE("Encode hotwords failed, skipping, hotwords are : %s",
243+
hotwords.c_str());
244+
}
245+
246+
int32_t num_default_hws = hotwords_.size();
247+
int32_t num_hws = current.size();
248+
249+
current.insert(current.end(), hotwords_.begin(), hotwords_.end());
250+
251+
if (!current_scores.empty() && !boost_scores_.empty()) {
252+
current_scores.insert(current_scores.end(), boost_scores_.begin(),
253+
boost_scores_.end());
254+
} else if (!current_scores.empty() && boost_scores_.empty()) {
255+
current_scores.insert(current_scores.end(), num_default_hws,
256+
config_.hotwords_score);
257+
} else if (current_scores.empty() && !boost_scores_.empty()) {
258+
current_scores.insert(current_scores.end(), num_hws,
259+
config_.hotwords_score);
260+
current_scores.insert(current_scores.end(), boost_scores_.begin(),
261+
boost_scores_.end());
262+
} else {
263+
// Do nothing.
264+
}
265+
266+
auto context_graph = std::make_shared<ContextGraph>(
267+
current, config_.hotwords_score, current_scores);
268+
return std::make_unique<OfflineStream>(config_.feat_config, context_graph);
269+
}
270+
222271
std::unique_ptr<OfflineStream> CreateStream() const override {
223272
if (config_.model_config.omnilingual.model.empty()) {
224-
return std::make_unique<OfflineStream>(config_.feat_config);
273+
return std::make_unique<OfflineStream>(config_.feat_config,
274+
hotwords_graph_);
225275
} else {
226276
return std::make_unique<OfflineStream>(OmnilingualAsrTag{});
227277
}
@@ -287,7 +337,7 @@ class OfflineRecognizerCtcImpl : public OfflineRecognizerImpl {
287337
-23.025850929940457f);
288338
auto t = model_->Forward(std::move(x), std::move(x_length));
289339

290-
auto results = decoder_->Decode(std::move(t[0]), std::move(t[1]));
340+
auto results = decoder_->Decode(std::move(t[0]), std::move(t[1]), ss, n);
291341

292342
int32_t frame_shift_ms = 10;
293343
for (int32_t i = 0; i != n; ++i) {
@@ -330,7 +380,9 @@ class OfflineRecognizerCtcImpl : public OfflineRecognizerImpl {
330380
x_length_shape.data(), x_length_shape.size());
331381

332382
auto t = model_->Forward(std::move(x), std::move(x_length));
333-
auto results = decoder_->Decode(std::move(t[0]), std::move(t[1]));
383+
384+
OfflineStream *ss[1] = {s};
385+
auto results = decoder_->Decode(std::move(t[0]), std::move(t[1]), ss, 1);
334386
int32_t frame_shift_ms = 10;
335387

336388
if (!config_.model_config.omnilingual.model.empty()) {
@@ -344,9 +396,60 @@ class OfflineRecognizerCtcImpl : public OfflineRecognizerImpl {
344396
s->SetResult(r);
345397
}
346398

399+
void InitHotwords() {
400+
// each line in hotwords_file contains space-separated words
401+
402+
std::ifstream is(config_.hotwords_file);
403+
if (!is) {
404+
SHERPA_ONNX_LOGE("Open hotwords file failed: %s",
405+
config_.hotwords_file.c_str());
406+
exit(-1);
407+
}
408+
409+
if (!EncodeHotwords(is, config_.model_config.modeling_unit, symbol_table_,
410+
bpe_encoder_.get(), &hotwords_, &boost_scores_)) {
411+
SHERPA_ONNX_LOGE(
412+
"Failed to encode some hotwords, skip them already, see logs above "
413+
"for details.");
414+
}
415+
hotwords_graph_ = std::make_shared<ContextGraph>(
416+
hotwords_, config_.hotwords_score, boost_scores_);
417+
}
418+
419+
#if __ANDROID_API__ >= 9
420+
void InitHotwords(AAssetManager *mgr) {
421+
// each line in hotwords_file contains space-separated words
422+
423+
auto buf = ReadFile(mgr, config_.hotwords_file);
424+
425+
std::istringstream is(std::string(buf.begin(), buf.end()));
426+
427+
if (!is) {
428+
SHERPA_ONNX_LOGE("Open hotwords file failed: %s",
429+
config_.hotwords_file.c_str());
430+
exit(-1);
431+
}
432+
433+
if (!EncodeHotwords(is, config_.model_config.modeling_unit, symbol_table_,
434+
bpe_encoder_.get(), &hotwords_, &boost_scores_)) {
435+
SHERPA_ONNX_LOGE(
436+
"Failed to encode some hotwords, skip them already, see logs above "
437+
"for details.");
438+
}
439+
hotwords_graph_ = std::make_shared<ContextGraph>(
440+
hotwords_, config_.hotwords_score, boost_scores_);
441+
}
442+
#endif
443+
347444
private:
348445
OfflineRecognizerConfig config_;
349446
SymbolTable symbol_table_;
447+
448+
std::vector<std::vector<int32_t>> hotwords_;
449+
std::vector<float> boost_scores_;
450+
ContextGraphPtr hotwords_graph_;
451+
std::unique_ptr<ssentencepiece::Ssentencepiece> bpe_encoder_;
452+
350453
std::unique_ptr<OfflineCtcModel> model_;
351454
std::unique_ptr<OfflineCtcDecoder> decoder_;
352455
};

‎sherpa-onnx/csrc/offline-recognizer.cc‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,8 @@ bool OfflineRecognizerConfig::Validate() const {
8181
}
8282
}
8383

84-
if (!hotwords_file.empty() && decoding_method != "modified_beam_search") {
84+
if (!hotwords_file.empty() && (decoding_method != "modified_beam_search" &&
85+
decoding_method != "prefix_beam_search")) {
8586
SHERPA_ONNX_LOGE(
8687
"Please use --decoding-method=modified_beam_search if you"
8788
" provide --hotwords-file. Given --decoding-method='%s'",

0 commit comments

Comments
 (0)