offline-transducer-modified-beam-search-decoder.cc
4.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
// sherpa-onnx/csrc/offline-transducer-modified-beam-search-decoder.cc
//
// Copyright (c) 2023 Xiaomi Corporation
#include "sherpa-onnx/csrc/offline-transducer-modified-beam-search-decoder.h"
#include <deque>
#include <utility>
#include <vector>
#include "sherpa-onnx/csrc/hypothesis.h"
#include "sherpa-onnx/csrc/onnx-utils.h"
#include "sherpa-onnx/csrc/packed-sequence.h"
#include "sherpa-onnx/csrc/slice.h"
namespace sherpa_onnx {
std::vector<OfflineTransducerDecoderResult>
OfflineTransducerModifiedBeamSearchDecoder::Decode(
Ort::Value encoder_out, Ort::Value encoder_out_length) {
PackedSequence packed_encoder_out = PackPaddedSequence(
model_->Allocator(), &encoder_out, &encoder_out_length);
int32_t batch_size =
static_cast<int32_t>(packed_encoder_out.sorted_indexes.size());
int32_t vocab_size = model_->VocabSize();
int32_t context_size = model_->ContextSize();
std::vector<int64_t> blanks(context_size, 0);
Hypotheses blank_hyp({{blanks, 0}});
std::deque<Hypotheses> finalized;
std::vector<Hypotheses> cur(batch_size, blank_hyp);
std::vector<Hypothesis> prev;
int32_t start = 0;
int32_t t = 0;
for (auto n : packed_encoder_out.batch_sizes) {
Ort::Value cur_encoder_out = packed_encoder_out.Get(start, n);
start += n;
if (n < static_cast<int32_t>(cur.size())) {
for (int32_t k = static_cast<int32_t>(cur.size()) - 1; k >= n; --k) {
finalized.push_front(std::move(cur[k]));
}
cur.erase(cur.begin() + n, cur.end());
} // if (n < static_cast<int32_t>(cur.size()))
// Due to merging paths with identical token sequences,
// not all utterances have "max_active_paths" paths.
auto hyps_row_splits = GetHypsRowSplits(cur);
int32_t num_hyps = hyps_row_splits.back();
prev.clear();
prev.reserve(num_hyps);
for (auto &hyps : cur) {
for (auto &h : hyps) {
prev.push_back(std::move(h.second));
}
}
cur.clear();
cur.reserve(n);
auto decoder_input = model_->BuildDecoderInput(prev, num_hyps);
// decoder_input shape: (num_hyps, context_size)
auto decoder_out = model_->RunDecoder(std::move(decoder_input));
// decoder_out is (num_hyps, joiner_dim)
cur_encoder_out =
Repeat(model_->Allocator(), &cur_encoder_out, hyps_row_splits);
// now cur_encoder_out is of shape (num_hyps, joiner_dim)
Ort::Value logit = model_->RunJoiner(
std::move(cur_encoder_out), Clone(model_->Allocator(), &decoder_out));
float *p_logit = logit.GetTensorMutableData<float>();
LogSoftmax(p_logit, vocab_size, num_hyps);
// now p_logit contains log_softmax output, we rename it to p_logprob
// to match what it actually contains
float *p_logprob = p_logit;
// add log_prob of each hypothesis to p_logprob before taking top_k
for (int32_t i = 0; i != num_hyps; ++i) {
float log_prob = prev[i].log_prob;
for (int32_t k = 0; k != vocab_size; ++k, ++p_logprob) {
*p_logprob += log_prob;
}
}
p_logprob = p_logit; // we changed p_logprob in the above for loop
// Now compute top_k for each utterance
for (int32_t i = 0; i != n; ++i) {
int32_t start = hyps_row_splits[i];
int32_t end = hyps_row_splits[i + 1];
auto topk =
TopkIndex(p_logprob, vocab_size * (end - start), max_active_paths_);
Hypotheses hyps;
for (auto k : topk) {
int32_t hyp_index = k / vocab_size + start;
int32_t new_token = k % vocab_size;
Hypothesis new_hyp = prev[hyp_index];
if (new_token != 0) {
// blank id is fixed to 0
new_hyp.ys.push_back(new_token);
new_hyp.timestamps.push_back(t);
}
new_hyp.log_prob = p_logprob[k];
hyps.Add(std::move(new_hyp));
} // for (auto k : topk)
p_logprob += (end - start) * vocab_size;
cur.push_back(std::move(hyps));
} // for (int32_t i = 0; i != n; ++i)
++t;
} // for (auto n : packed_encoder_out.batch_sizes)
for (auto &h : finalized) {
cur.push_back(std::move(h));
}
if (lm_) {
// use LM for rescoring
lm_->ComputeLMScore(lm_scale_, context_size, &cur);
}
std::vector<OfflineTransducerDecoderResult> unsorted_ans(batch_size);
for (int32_t i = 0; i != batch_size; ++i) {
Hypothesis hyp = cur[i].GetMostProbable(true);
auto &r = unsorted_ans[packed_encoder_out.sorted_indexes[i]];
// strip leading blanks
r.tokens = {hyp.ys.begin() + context_size, hyp.ys.end()};
r.timestamps = std::move(hyp.timestamps);
}
return unsorted_ans;
}
} // namespace sherpa_onnx