hypothesis.h
2.9 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
/**
* Copyright (c) 2023 Xiaomi Corporation
*
*/
#ifndef SHERPA_ONNX_CSRC_HYPOTHESIS_H_
#define SHERPA_ONNX_CSRC_HYPOTHESIS_H_
#include <sstream>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "sherpa-onnx/csrc/math.h"
namespace sherpa_onnx {
struct Hypothesis {
// The predicted tokens so far. Newly predicated tokens are appended.
std::vector<int64_t> ys;
// timestamps[i] contains the frame number after subsampling
// on which ys[i] is decoded.
std::vector<int32_t> timestamps;
// The total score of ys in log space.
double log_prob = 0;
int32_t num_trailing_blanks = 0;
Hypothesis() = default;
Hypothesis(const std::vector<int64_t> &ys, double log_prob)
: ys(ys), log_prob(log_prob) {}
// If two Hypotheses have the same `Key`, then they contain
// the same token sequence.
std::string Key() const {
// TODO(fangjun): Use a hash function?
std::ostringstream os;
std::string sep = "-";
for (auto i : ys) {
os << i << sep;
sep = "-";
}
return os.str();
}
// For debugging
std::string ToString() const {
std::ostringstream os;
os << "(" << Key() << ", " << log_prob << ")";
return os.str();
}
};
class Hypotheses {
public:
Hypotheses() = default;
explicit Hypotheses(std::vector<Hypothesis> hyps) {
for (auto &h : hyps) {
hyps_dict_[h.Key()] = std::move(h);
}
}
explicit Hypotheses(std::unordered_map<std::string, Hypothesis> hyps_dict)
: hyps_dict_(std::move(hyps_dict)) {}
// Add hyp to this object. If it already exists, its log_prob
// is updated with the given hyp using log-sum-exp.
void Add(Hypothesis hyp);
// Get the hyp that has the largest log_prob.
// If length_norm is true, hyp's log_prob is divided by
// len(hyp.ys) before comparison.
Hypothesis GetMostProbable(bool length_norm) const;
// Get the k hyps that have the largest log_prob.
// If length_norm is true, hyp's log_prob is divided by
// len(hyp.ys) before comparison.
std::vector<Hypothesis> GetTopK(int32_t k, bool length_norm) const;
int32_t Size() const { return hyps_dict_.size(); }
std::string ToString() const {
std::ostringstream os;
for (const auto &p : hyps_dict_) {
os << p.second.ToString() << "\n";
}
return os.str();
}
const auto begin() const { return hyps_dict_.begin(); }
const auto end() const { return hyps_dict_.end(); }
void Clear() { hyps_dict_.clear(); }
private:
// Return a list of hyps contained in this object.
std::vector<Hypothesis> Vec() const {
std::vector<Hypothesis> ans;
ans.reserve(hyps_dict_.size());
for (const auto &p : hyps_dict_) {
ans.push_back(p.second);
}
return ans;
}
private:
using Map = std ::unordered_map<std::string, Hypothesis>;
Map hyps_dict_;
};
} // namespace sherpa_onnx
#endif // SHERPA_ONNX_CSRC_HYPOTHESIS_H_