248 lines
11 KiB
C++
248 lines
11 KiB
C++
#pragma once
|
|
#include <onnxruntime_cxx_api.h>
|
|
#include <tokenizers_cpp.h> // mlc-ai/tokenizers-cpp — reads tokenizer.json
|
|
|
|
#include <string>
|
|
#include <vector>
|
|
#include <memory>
|
|
#include <fstream>
|
|
#include <sstream>
|
|
#include <algorithm>
|
|
#include <cctype>
|
|
#include <unicode/unistr.h> // ICU — for NFKC + case fold
|
|
#include <unicode/normalizer2.h>
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// EOUDetector — End-of-Utterance probability via Qwen2.5-0.5B ONNX model
|
|
//
|
|
// ONNX interface:
|
|
// input "input_ids" INT64 [1, seq_len] (seq_len dynamic)
|
|
// output "eou_prob" FLOAT [1]
|
|
//
|
|
// Usage:
|
|
// EOUDetector eou("path/to/model.onnx", "path/to/tokenizer.json");
|
|
// float p = eou.score(history, current_text);
|
|
// if (p >= EOUDetector::THRESHOLD) { /* fire endpoint */ }
|
|
// ---------------------------------------------------------------------------
|
|
|
|
class EOUDetector {
|
|
public:
|
|
// Tuned on 200 labelled samples: FA=22.3%, FR=6.0%, F1=0.851
|
|
static constexpr float THRESHOLD = 0.0766f;
|
|
static constexpr int MAX_TOKENS = 512;
|
|
|
|
struct Turn {
|
|
enum class Role { User, Assistant } role;
|
|
std::string text; // raw text — preprocessing applied internally
|
|
};
|
|
|
|
// modelPath — directory containing model.onnx AND model.onnx.data
|
|
// (ORT finds the .data file automatically; both must coexist)
|
|
// tokenizerJson — path to tokenizer.json from the model directory
|
|
EOUDetector(const std::string& modelPath,
|
|
const std::string& tokenizerJson,
|
|
bool useGpu = false)
|
|
: env_(ORT_LOGGING_LEVEL_WARNING, "eou")
|
|
, memoryInfo_("Cpu", OrtDeviceAllocator, 0, OrtMemTypeDefault)
|
|
{
|
|
// ── session ──────────────────────────────────────────────────────────
|
|
Ort::SessionOptions opts;
|
|
opts.SetIntraOpNumThreads(1); // per-request; tune for your CPU core budget
|
|
opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);
|
|
if (useGpu) {
|
|
OrtCUDAProviderOptions cuda{};
|
|
cuda.device_id = 0;
|
|
opts.AppendExecutionProvider_CUDA(cuda);
|
|
}
|
|
|
|
std::string onnxPath = modelPath + "/model.onnx";
|
|
session_ = std::make_unique<Ort::Session>(env_, onnxPath.c_str(), opts);
|
|
|
|
// ── cache I/O names (avoids per-call allocation) ─────────────────────
|
|
inputName_ = session_->GetInputNameAllocated(0, allocator_).get();
|
|
outputName_ = session_->GetOutputNameAllocated(0, allocator_).get();
|
|
|
|
// ── tokenizer ────────────────────────────────────────────────────────
|
|
std::ifstream f(tokenizerJson);
|
|
std::string json((std::istreambuf_iterator<char>(f)),
|
|
std::istreambuf_iterator<char>());
|
|
tokenizer_ = tokenizers::Tokenizer::FromBlobJSON(json);
|
|
}
|
|
|
|
// Returns P(end-of-utterance) in [0, 1].
|
|
// history — prior turns in chronological order, alternating user/assistant
|
|
// (may be empty for first turn)
|
|
// currentText — the latest partial or complete utterance being evaluated
|
|
float score(const std::vector<Turn>& history,
|
|
const std::string& currentText)
|
|
{
|
|
std::vector<int64_t> ids = buildInputIds(history, currentText);
|
|
|
|
// ── dims ─────────────────────────────────────────────────────────────
|
|
int64_t seqLen = static_cast<int64_t>(ids.size());
|
|
std::vector<int64_t> inputDims {1, seqLen};
|
|
std::vector<int64_t> outputDims {1};
|
|
|
|
// ── I/O name pointers ─────────────────────────────────────────────────
|
|
std::vector<const char*> inputNames {inputName_.c_str()};
|
|
std::vector<const char*> outputNames {outputName_.c_str()};
|
|
|
|
// ── input tensor (INT64) ──────────────────────────────────────────────
|
|
std::vector<Ort::Value> inputTensors;
|
|
inputTensors.push_back(
|
|
Ort::Value::CreateTensor<int64_t>(
|
|
memoryInfo_,
|
|
ids.data(), ids.size(),
|
|
inputDims.data(), inputDims.size()));
|
|
|
|
// ── output tensor (FLOAT32) ───────────────────────────────────────────
|
|
std::vector<float> outputValues(1);
|
|
std::vector<Ort::Value> outputTensors;
|
|
outputTensors.push_back(
|
|
Ort::Value::CreateTensor<float>(
|
|
memoryInfo_,
|
|
outputValues.data(), outputValues.size(),
|
|
outputDims.data(), outputDims.size()));
|
|
|
|
// ── run ───────────────────────────────────────────────────────────────
|
|
session_->Run(
|
|
Ort::RunOptions{nullptr},
|
|
inputNames.data(), inputTensors.data(), inputTensors.size(),
|
|
outputNames.data(), outputTensors.data(), outputTensors.size());
|
|
|
|
return outputValues[0];
|
|
}
|
|
|
|
bool isEndOfUtterance(const std::vector<Turn>& history,
|
|
const std::string& currentText)
|
|
{
|
|
return score(history, currentText) >= THRESHOLD;
|
|
}
|
|
|
|
private:
|
|
// ── Qwen2.5 chat-template token IDs (hardcoded — never change) ───────────
|
|
//
|
|
// <|im_start|>system\n
|
|
// You are Qwen, created by Alibaba Cloud. You are a helpful assistant.
|
|
// <|im_end|>\n
|
|
static const std::vector<int64_t> SYSTEM_BLOCK;
|
|
static const std::vector<int64_t> USER_OPEN; // <|im_start|>user\n
|
|
static const std::vector<int64_t> USER_CLOSE; // <|im_end|>\n
|
|
static const std::vector<int64_t> ASST_OPEN; // <|im_start|>assistant\n
|
|
static const std::vector<int64_t> ASST_CLOSE; // <|im_end|>\n
|
|
|
|
// ── members ───────────────────────────────────────────────────────────────
|
|
Ort::Env env_;
|
|
Ort::AllocatorWithDefaultOptions allocator_;
|
|
Ort::MemoryInfo memoryInfo_;
|
|
std::unique_ptr<Ort::Session> session_;
|
|
std::unique_ptr<tokenizers::Tokenizer> tokenizer_;
|
|
std::string inputName_;
|
|
std::string outputName_;
|
|
|
|
// ── text preprocessing ────────────────────────────────────────────────────
|
|
// Matches training: NFKC → lowercase → strip punctuation (keep ' -)
|
|
// → collapse whitespace
|
|
std::string preprocess(const std::string& raw)
|
|
{
|
|
// NFKC normalise + fold case via ICU
|
|
icu::UnicodeString u = icu::UnicodeString::fromUTF8(raw);
|
|
UErrorCode err = U_ZERO_ERROR;
|
|
const icu::Normalizer2* nfkc =
|
|
icu::Normalizer2::getNFKCInstance(err);
|
|
u = nfkc->normalize(u, err);
|
|
u.foldCase(U_FOLD_CASE_DEFAULT);
|
|
|
|
std::string s;
|
|
u.toUTF8String(s);
|
|
|
|
// strip punctuation (keep apostrophe and hyphen), collapse spaces
|
|
std::string out;
|
|
out.reserve(s.size());
|
|
bool lastWasSpace = true;
|
|
for (unsigned char c : s) {
|
|
if (c == '\'' || c == '-') {
|
|
out += c; lastWasSpace = false;
|
|
} else if (std::ispunct(c)) {
|
|
// skip
|
|
} else if (std::isspace(c)) {
|
|
if (!lastWasSpace) { out += ' '; lastWasSpace = true; }
|
|
} else {
|
|
out += c; lastWasSpace = false;
|
|
}
|
|
}
|
|
while (!out.empty() && out.back() == ' ') out.pop_back();
|
|
return out;
|
|
}
|
|
|
|
// ── BPE-encode text, return int64 ids ────────────────────────────────────
|
|
std::vector<int64_t> encode(const std::string& text)
|
|
{
|
|
auto ids32 = tokenizer_->Encode(text, /*add_special_tokens=*/false);
|
|
std::vector<int64_t> ids64(ids32.begin(), ids32.end());
|
|
return ids64;
|
|
}
|
|
|
|
// ── append helper ─────────────────────────────────────────────────────────
|
|
static void append(std::vector<int64_t>& dst, const std::vector<int64_t>& src)
|
|
{
|
|
dst.insert(dst.end(), src.begin(), src.end());
|
|
}
|
|
|
|
// ── build full input_ids for ONNX ────────────────────────────────────────
|
|
std::vector<int64_t> buildInputIds(const std::vector<Turn>& history,
|
|
const std::string& currentText)
|
|
{
|
|
std::vector<int64_t> ids;
|
|
ids.reserve(MAX_TOKENS);
|
|
|
|
// system block (always present)
|
|
append(ids, SYSTEM_BLOCK);
|
|
|
|
// prior turns — merge consecutive same-role turns
|
|
std::vector<Turn> merged;
|
|
for (const auto& t : history) {
|
|
std::string clean = preprocess(t.text);
|
|
if (!merged.empty() && merged.back().role == t.role)
|
|
merged.back().text += " " + clean;
|
|
else
|
|
merged.push_back({t.role, clean});
|
|
}
|
|
|
|
for (const auto& t : merged) {
|
|
if (t.role == Turn::Role::User) {
|
|
append(ids, USER_OPEN);
|
|
append(ids, encode(t.text));
|
|
append(ids, USER_CLOSE);
|
|
} else {
|
|
append(ids, ASST_OPEN);
|
|
append(ids, encode(t.text));
|
|
append(ids, ASST_CLOSE);
|
|
}
|
|
}
|
|
|
|
// current (last) user turn — NO closing <|im_end|>
|
|
append(ids, USER_OPEN);
|
|
append(ids, encode(preprocess(currentText)));
|
|
|
|
// truncate to MAX_TOKENS from the right (keep most recent context)
|
|
if ((int)ids.size() > MAX_TOKENS)
|
|
ids = std::vector<int64_t>(ids.end() - MAX_TOKENS, ids.end());
|
|
|
|
return ids;
|
|
}
|
|
};
|
|
|
|
// ── static member definitions ─────────────────────────────────────────────────
|
|
|
|
const std::vector<int64_t> EOUDetector::SYSTEM_BLOCK = {
|
|
151644, 8948, 198, // <|im_start|>system\n
|
|
2610, 525, 1207, 16948, 11, 3465, 553, 54364, 14817, // You are Qwen, created by Alibaba Cloud.
|
|
13, 1446, 525, 264, 10950, 17847, 13, // You are a helpful assistant.
|
|
151645, 198 // <|im_end|>\n
|
|
};
|
|
const std::vector<int64_t> EOUDetector::USER_OPEN = {151644, 872, 198};
|
|
const std::vector<int64_t> EOUDetector::USER_CLOSE = {151645, 198};
|
|
const std::vector<int64_t> EOUDetector::ASST_OPEN = {151644, 77091, 198};
|
|
const std::vector<int64_t> EOUDetector::ASST_CLOSE = {151645, 198};
|