93 lines
2.9 KiB
C++
93 lines
2.9 KiB
C++
// sherpa-onnx/csrc/context-graph.h
|
|
//
|
|
// Copyright (c) 2023 Xiaomi Corporation
|
|
|
|
#ifndef SHERPA_ONNX_CSRC_CONTEXT_GRAPH_H_
|
|
#define SHERPA_ONNX_CSRC_CONTEXT_GRAPH_H_
|
|
|
|
#include <memory>
|
|
#include <string>
|
|
#include <tuple>
|
|
#include <unordered_map>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
#include "sherpa-onnx/csrc/log.h"
|
|
|
|
namespace sherpa_onnx {
|
|
|
|
class ContextGraph;
|
|
using ContextGraphPtr = std::shared_ptr<ContextGraph>;
|
|
|
|
struct ContextState {
|
|
int32_t token;
|
|
float token_score;
|
|
float node_score;
|
|
float output_score;
|
|
int32_t level;
|
|
float ac_threshold;
|
|
bool is_end;
|
|
std::string phrase;
|
|
std::unordered_map<int32_t, std::unique_ptr<ContextState>> next;
|
|
const ContextState *fail = nullptr;
|
|
const ContextState *output = nullptr;
|
|
|
|
ContextState() = default;
|
|
ContextState(int32_t token, float token_score, float node_score,
|
|
float output_score, int32_t level = 0, float ac_threshold = 0.0f,
|
|
bool is_end = false, const std::string &phrase = {})
|
|
: token(token),
|
|
token_score(token_score),
|
|
node_score(node_score),
|
|
output_score(output_score),
|
|
level(level),
|
|
ac_threshold(ac_threshold),
|
|
is_end(is_end),
|
|
phrase(phrase) {}
|
|
};
|
|
|
|
class ContextGraph {
|
|
public:
|
|
ContextGraph() = default;
|
|
ContextGraph(const std::vector<std::vector<int32_t>> &token_ids,
|
|
float context_score, float ac_threshold,
|
|
const std::vector<float> &scores = {},
|
|
const std::vector<std::string> &phrases = {},
|
|
const std::vector<float> &ac_thresholds = {})
|
|
: context_score_(context_score), ac_threshold_(ac_threshold) {
|
|
root_ = std::make_unique<ContextState>(-1, 0, 0, 0);
|
|
root_->fail = root_.get();
|
|
Build(token_ids, scores, phrases, ac_thresholds);
|
|
}
|
|
|
|
ContextGraph(const std::vector<std::vector<int32_t>> &token_ids,
|
|
float context_score, const std::vector<float> &scores = {})
|
|
: ContextGraph(token_ids, context_score, 0.0f, scores,
|
|
std::vector<std::string>(), std::vector<float>()) {}
|
|
|
|
std::tuple<float, const ContextState *, const ContextState *> ForwardOneStep(
|
|
const ContextState *state, int32_t token_id,
|
|
bool strict_mode = true) const;
|
|
|
|
std::pair<bool, const ContextState *> IsMatched(
|
|
const ContextState *state) const;
|
|
|
|
std::pair<float, const ContextState *> Finalize(
|
|
const ContextState *state) const;
|
|
|
|
const ContextState *Root() const { return root_.get(); }
|
|
|
|
private:
|
|
float context_score_;
|
|
float ac_threshold_;
|
|
std::unique_ptr<ContextState> root_;
|
|
void Build(const std::vector<std::vector<int32_t>> &token_ids,
|
|
const std::vector<float> &scores,
|
|
const std::vector<std::string> &phrases,
|
|
const std::vector<float> &ac_thresholds) const;
|
|
void FillFailOutput() const;
|
|
};
|
|
|
|
} // namespace sherpa_onnx
|
|
#endif // SHERPA_ONNX_CSRC_CONTEXT_GRAPH_H_
|