diff --git a/.gitignore b/.gitignore index 672e3f88..9cd27b2b 100644 --- a/.gitignore +++ b/.gitignore @@ -22,4 +22,8 @@ log.txt /language-model/cache/ /language-model/sc.lm /language-model/sc.lm.incoming + +# 神经整句模型：lock.json 和 NOTICE.md 进仓库，下载下来的 safetensors 不进 —— +# 由 scripts/fetch-neural-model.ps1 按 lock.json 从 release 重新取到并校验 sha256。 +/neural-model/*.safetensors __pycache__/ diff --git a/engine/CMakeLists.txt b/engine/CMakeLists.txt index e89f07b8..ced881c5 100644 --- a/engine/CMakeLists.txt +++ b/engine/CMakeLists.txt @@ -50,6 +50,9 @@ set(METASEQUOIA_IME_ENGINE_SOURCES local_modes/kaomoji_query.cpp local_modes/quick_phrase_query.cpp local_modes/unicode_query.cpp + neural/sentence_model.cpp + neural/neural_decoder.cpp + neural/rescore_worker.cpp providers/japanese_candidate_provider.cpp providers/pinyin_candidate_provider.cpp providers/provider_registry.cpp @@ -58,6 +61,7 @@ set(METASEQUOIA_IME_ENGINE_SOURCES quanpin/quanpin_dictionary.cpp quanpin/quanpin_query.cpp quanpin/quanpin_utils.cpp + quanpin/lattice_rerank.cpp quanpin/word_lattice.cpp schemes/japanese_romaji_scheme.cpp schemes/quanpin_scheme.cpp @@ -207,6 +211,11 @@ if(BUILD_TESTING) target_link_libraries(test_engine_smoke PRIVATE MetasequoiaIme::Engine) add_test(NAME engine_smoke COMMAND test_engine_smoke) + # 神经整句排序集成测试：需数据目录里备好 sentence-model*.safetensors，缺文件时优雅跳过。 + add_executable(test_neural_sentence tests/src/test_neural_sentence.cpp) + target_link_libraries(test_neural_sentence PRIVATE MetasequoiaIme::Engine) + add_test(NAME neural_sentence COMMAND test_neural_sentence) + add_executable(test_input_session tests/src/test_input_session.cpp) target_compile_definitions(test_input_session PRIVATE METASEQUOIA_SKIP_FREQUENCY_TESTS) target_link_libraries(test_input_session PRIVATE MetasequoiaIme::Engine) diff --git a/engine/contracts/assets/assets.h b/engine/contracts/assets/assets.h index d5e5ddc4..88116b35 100644 --- a/engine/contracts/assets/assets.h +++ b/engine/contracts/assets/assets.h @@ -11,6 +11,9 @@ inline constexpr char japanese_model[] = "dict_japanese.dat"; inline constexpr char japanese_notice[] = "mozc_dictionary_oss_README.txt"; inline constexpr char dictionary_manifest[] = "dictionary-manifest.json"; inline constexpr char language_model[] = "sc.lm"; +inline constexpr char neural_model_desktop[] = "sentence-model-desktop.safetensors"; +inline constexpr char neural_model_keyboard[] = "sentence-model.safetensors"; +inline constexpr char notice_19[] = "licenses/chinese-ime-lm-NOTICE.md"; inline constexpr char pinyin_model[] = "dict_pinyin.dat"; inline constexpr char translations[] = "custom_translations.txt"; inline constexpr char helpcode_lantian[] = "helpcodes/helpcode.txt"; diff --git a/engine/contracts/assets/assets.json b/engine/contracts/assets/assets.json index 3a869666..f0d5ba42 100644 --- a/engine/contracts/assets/assets.json +++ b/engine/contracts/assets/assets.json @@ -66,6 +66,33 @@ "desktop" ] }, + { + "id": "neural_model_desktop", + "path": "sentence-model-desktop.safetensors", + "role": "resource", + "source": "neural-model/sentence-model-desktop.safetensors", + "profiles": [ + "desktop" + ] + }, + { + "id": "neural_model_keyboard", + "path": "sentence-model.safetensors", + "role": "resource", + "source": "neural-model/sentence-model.safetensors", + "profiles": [ + "desktop" + ] + }, + { + "id": "notice_19", + "path": "licenses/chinese-ime-lm-NOTICE.md", + "role": "notice", + "source": "neural-model/NOTICE.md", + "profiles": [ + "desktop" + ] + }, { "id": "pinyin_model", "path": "dict_pinyin.dat", diff --git a/engine/core/ime_session.cpp b/engine/core/ime_session.cpp index 6b2a5884..b5c469f0 100644 --- a/engine/core/ime_session.cpp +++ b/engine/core/ime_session.cpp @@ -288,6 +288,8 @@ void ImeSession::apply_request_options(QueryRequest &request) const (quanpin_autocorrect_types_ & quanpin::kAutocorrectTransposition) != 0; request.enable_quanpin_autocorrect_neighbor = (quanpin_autocorrect_types_ & quanpin::kAutocorrectNeighbor) != 0; request.fuzzy_pinyin = fuzzy_pinyin_; + request.sentence_association = sentence_association_; + request.rescoring_context = rescoring_context_; } void ImeSession::refresh_candidates() diff --git a/engine/core/ime_session.h b/engine/core/ime_session.h index 0269d492..b6361d3c 100644 --- a/engine/core/ime_session.h +++ b/engine/core/ime_session.h @@ -3,6 +3,7 @@ #include "composition_state.h" #include "input_session_types.h" #include "scheme_type.h" +#include "sentence_association_options.h" #include "../providers/provider_registry.h" #include "../schemes/input_scheme.h" #include "../schemes/wubi_scheme.h" @@ -25,6 +26,15 @@ class ImeSession { fuzzy_pinyin_ = options; } + void set_sentence_association(const SentenceAssociationOptions &options) + { + sentence_association_ = options; + } + // 本会话最近上屏的文本，随查询下发给神经重排当前文。 + void set_rescoring_context(std::string context) + { + rescoring_context_ = std::move(context); + } void set_wubi_input_options(metasequoia::WubiInputOptions options); void replace_shuangpin_raw_input(const std::string &raw_input, const std::string &raw_input_with_cases); void replace_quanpin_raw_input(const std::string &raw_input, const std::string &raw_input_with_cases); @@ -90,6 +100,8 @@ class ImeSession bool enable_quanpin_helpcode_ = false; unsigned quanpin_autocorrect_types_ = 0; metasequoia::FuzzyPinyinOptions fuzzy_pinyin_; + SentenceAssociationOptions sentence_association_; + std::string rescoring_context_; metasequoia::WubiInputOptions wubi_options_; // Resolved when the scheme changes rather than on every keystroke. WubiScheme *wubi_scheme_ = nullptr; diff --git a/engine/core/input_session.cpp b/engine/core/input_session.cpp index 4edc82f7..7249f8bd 100644 --- a/engine/core/input_session.cpp +++ b/engine/core/input_session.cpp @@ -978,7 +978,8 @@ std::optional InputSession::learn_candidate(std::size_t index) // 整句候选走造词，不走调频：它在词库里没有行可改。不受 frequency_adjustment_ // 约束（那是候选词频策略），也不跳过 index == 0——首位整句同样没有落库， // 下次仍要靠猜，存下来才稳定。 - if (selected.source == CandidateSource::Generated || selected.source == CandidateSource::Fallback) + if (selected.source == CandidateSource::Generated || selected.source == CandidateSource::Fallback || + selected.source == CandidateSource::NeuralDesktop || selected.source == CandidateSource::NeuralKeyboard) { return learn_sentence_candidate(selected); } diff --git a/engine/core/input_session.h b/engine/core/input_session.h index 8ccede6d..36bf882a 100644 --- a/engine/core/input_session.h +++ b/engine/core/input_session.h @@ -209,6 +209,14 @@ class InputSession { engine_.set_fuzzy_pinyin_options(options); } + void set_sentence_association(const SentenceAssociationOptions &options) + { + engine_.set_sentence_association(options); + } + void set_rescoring_context(std::string context) + { + engine_.set_rescoring_context(std::move(context)); + } void set_chinese_punctuation_enabled(bool enabled) { chinese_punctuation_enabled_ = enabled; diff --git a/engine/core/input_session_composition.cpp b/engine/core/input_session_composition.cpp index 7f01fa11..6cda3856 100644 --- a/engine/core/input_session_composition.cpp +++ b/engine/core/input_session_composition.cpp @@ -526,7 +526,8 @@ int InputSession::store_user_phrase_from_canonical_pinyin(std::string pinyin, st return canonical_phrase_engine_->create_word_from_canonical_pinyin(std::move(pinyin), std::move(word)); } -// 整句候选（词格 CandidateSource::Generated、Google 解码器 CandidateSource::Fallback） +// 整句候选（词格 CandidateSource::Generated、Google 解码器 CandidateSource::Fallback、 +// 神经整句 NeuralDesktop / NeuralKeyboard） // 是猜出来的，词库里没有它那一行，所以调频对它无效：update_weight_by_pinyin_and_word // 改的是 SQLite 里已存在的行，找不到行就什么也不做，用户于是发现自己选多少次都提不上来。 // 选中即落库才是对的处理：存成用户词组之后，下次同样的输入它以 UserDatabase 候选出现， diff --git a/engine/core/query_request.h b/engine/core/query_request.h index 2df72ea4..7de7e6ee 100644 --- a/engine/core/query_request.h +++ b/engine/core/query_request.h @@ -3,6 +3,7 @@ #include "key_event.h" #include "fuzzy_pinyin_options.h" #include "scheme_type.h" +#include "sentence_association_options.h" #include #include @@ -30,5 +31,11 @@ struct QueryRequest bool enable_quanpin_autocorrect_neighbor = false; std::vector key_strokes; metasequoia::FuzzyPinyinOptions fuzzy_pinyin; + // 整句候选来源与去重补位选项，默认全关。 + SentenceAssociationOptions sentence_association; + // 给神经重排看的前文：本会话最近上屏的文本。空着也能重排，只是模型看不到语境，「shanghai」 + // 到底是上海还是伤害就只能靠词格自己的静态分。词典层只取末尾若干字（RerankOptions:: + // context_chars），所以这里给多了也无妨。 + std::string rescoring_context; bool valid = false; }; diff --git a/engine/core/sentence_association_options.h b/engine/core/sentence_association_options.h new file mode 100644 index 00000000..24819fd9 --- /dev/null +++ b/engine/core/sentence_association_options.h @@ -0,0 +1,30 @@ +#pragma once + +// 整句联想的开关。全部默认关闭，由用户在设置里手动开启。 +// +// 整句候选最多四行：词格首选、Google 解码器，以及两个神经模型的首选各一条。 +// +// 前两个各是一条独立来源的整句。 +// +// 后两个不一样。神经模型不再自己造句，它只在词格内部解出的最多 12 条 n-best 里挑一条（见 +// engine/neural/neural_decoder.h）。它们不依赖 word_lattice 的显示开关：只开神经时词格仍会在内部 +// 生成备选，但只显示神经选中的句子；同时开 word_lattice 时才额外保留词格首选作对照。两个神经 +// 开关都开时各自独立打分，各出一条首选；两边选中相同文本时合并成一条。 +struct SentenceAssociationOptions +{ + bool word_lattice = false; // kenlm 词格整句，CandidateSource::Generated + bool google = false; // Google 解码器整句，CandidateSource::Fallback + bool neural_desktop = false; // 用高精度档的神经模型重排词格整句 + bool neural_keyboard = false; // 用快速档的神经模型重排词格整句 + bool show_next_on_duplicate = false; // 某来源首选被去重时，显示该来源下一条不同结果 + + bool operator==(const SentenceAssociationOptions &other) const + { + return word_lattice == other.word_lattice && google == other.google && neural_desktop == other.neural_desktop && + neural_keyboard == other.neural_keyboard && show_next_on_duplicate == other.show_next_on_duplicate; + } + bool operator!=(const SentenceAssociationOptions &other) const + { + return !(*this == other); + } +}; diff --git a/engine/core/word_item.h b/engine/core/word_item.h index 9a1d64dd..29c1b14b 100644 --- a/engine/core/word_item.h +++ b/engine/core/word_item.h @@ -16,6 +16,12 @@ enum class CandidateSource Kaomoji, Generated, Fallback, + // 神经整句重排挑中的那一条：句子本身还是词格解出来的 n-best，chinese-ime-lm 的字符级 + // Transformer 只负责在里面挑最准的一句，挑中的那行记在这里，词格自己的第一名仍记 + // Generated。重排器弃权或后台还没算完时不会出现这两个来源。desktop 档更准更慢， + // keyboard 档更快。两者都属于「猜出来的整句」，落库/学习判定与 Generated/Fallback 同类。 + NeuralDesktop, + NeuralKeyboard, }; struct WordItem diff --git a/engine/neural/json.h b/engine/neural/json.h new file mode 100644 index 00000000..42e84f9d --- /dev/null +++ b/engine/neural/json.h @@ -0,0 +1,464 @@ +#pragma once + +// Minimal, self-contained JSON parser for the safetensors header. +// +// The core engine target does not depend on a JSON library (nlohmann is only pulled in by the +// optional `voice` feature), and the safetensors header is small and regular. Rather than add a +// build-wide dependency, this parses the handful of shapes the loader needs: the flat object of +// tensor descriptors, plus the `__metadata__` block whose `config` and `vocab` values are +// themselves JSON strings parsed a second time. +// +// It is a straightforward recursive-descent parser over UTF-8 text. serde_json (which writes the +// files) emits raw UTF-8 for non-ASCII, but `\uXXXX` escapes are handled anyway so a re-encoded +// file still loads. + +#include +#include +#include +#include +#include + +namespace neural::json +{ + +struct Value +{ + enum class Type + { + Null, + Bool, + Number, + String, + Array, + Object, + }; + + Type type = Type::Null; + bool boolean = false; + double number = 0.0; + std::string string; + std::vector array; + std::vector> object; + + bool is_object() const + { + return type == Type::Object; + } + bool is_array() const + { + return type == Type::Array; + } + bool is_string() const + { + return type == Type::String; + } + bool is_number() const + { + return type == Type::Number; + } + + // Linear lookup: objects here have at most a few dozen keys, and the tensor table is read once. + const Value *find(const std::string &key) const + { + for (const auto &entry : object) + { + if (entry.first == key) + { + return &entry.second; + } + } + return nullptr; + } +}; + +class ParseError : public std::runtime_error +{ + public: + explicit ParseError(const std::string &message) : std::runtime_error(message) + { + } +}; + +namespace detail +{ + +class Parser +{ + public: + explicit Parser(const std::string &text) : text_(text) + { + } + + Value parse() + { + skip_whitespace(); + Value value = parse_value(); + skip_whitespace(); + if (position_ != text_.size()) + { + fail("trailing characters after JSON value"); + } + return value; + } + + private: + const std::string &text_; + size_t position_ = 0; + + [[noreturn]] void fail(const std::string &message) const + { + throw ParseError("json: " + message + " at offset " + std::to_string(position_)); + } + + char peek() const + { + if (position_ >= text_.size()) + { + fail("unexpected end of input"); + } + return text_[position_]; + } + + char next() + { + char c = peek(); + ++position_; + return c; + } + + void expect(char c) + { + if (next() != c) + { + fail(std::string("expected '") + c + "'"); + } + } + + void skip_whitespace() + { + while (position_ < text_.size()) + { + char c = text_[position_]; + if (c == ' ' || c == '\t' || c == '\n' || c == '\r') + { + ++position_; + } + else + { + break; + } + } + } + + Value parse_value() + { + skip_whitespace(); + char c = peek(); + switch (c) + { + case '{': + return parse_object(); + case '[': + return parse_array(); + case '"': { + Value value; + value.type = Value::Type::String; + value.string = parse_string(); + return value; + } + case 't': + case 'f': + return parse_bool(); + case 'n': + return parse_null(); + default: + return parse_number(); + } + } + + Value parse_object() + { + Value value; + value.type = Value::Type::Object; + expect('{'); + skip_whitespace(); + if (peek() == '}') + { + ++position_; + return value; + } + while (true) + { + skip_whitespace(); + std::string key = parse_string(); + skip_whitespace(); + expect(':'); + Value child = parse_value(); + value.object.emplace_back(std::move(key), std::move(child)); + skip_whitespace(); + char c = next(); + if (c == ',') + { + continue; + } + if (c == '}') + { + break; + } + fail("expected ',' or '}' in object"); + } + return value; + } + + Value parse_array() + { + Value value; + value.type = Value::Type::Array; + expect('['); + skip_whitespace(); + if (peek() == ']') + { + ++position_; + return value; + } + while (true) + { + value.array.push_back(parse_value()); + skip_whitespace(); + char c = next(); + if (c == ',') + { + continue; + } + if (c == ']') + { + break; + } + fail("expected ',' or ']' in array"); + } + return value; + } + + Value parse_bool() + { + Value value; + value.type = Value::Type::Bool; + if (text_.compare(position_, 4, "true") == 0) + { + position_ += 4; + value.boolean = true; + } + else if (text_.compare(position_, 5, "false") == 0) + { + position_ += 5; + value.boolean = false; + } + else + { + fail("invalid literal"); + } + return value; + } + + Value parse_null() + { + Value value; + value.type = Value::Type::Null; + if (text_.compare(position_, 4, "null") == 0) + { + position_ += 4; + } + else + { + fail("invalid literal"); + } + return value; + } + + Value parse_number() + { + size_t start = position_; + if (position_ < text_.size() && (text_[position_] == '-' || text_[position_] == '+')) + { + ++position_; + } + while (position_ < text_.size()) + { + char c = text_[position_]; + if ((c >= '0' && c <= '9') || c == '.' || c == 'e' || c == 'E' || c == '+' || c == '-') + { + ++position_; + } + else + { + break; + } + } + if (position_ == start) + { + fail("invalid number"); + } + Value value; + value.type = Value::Type::Number; + try + { + value.number = std::stod(text_.substr(start, position_ - start)); + } + catch (const std::exception &) + { + fail("invalid number"); + } + return value; + } + + std::string parse_string() + { + expect('"'); + std::string out; + while (true) + { + if (position_ >= text_.size()) + { + fail("unterminated string"); + } + char c = text_[position_++]; + if (c == '"') + { + break; + } + if (c == '\\') + { + if (position_ >= text_.size()) + { + fail("unterminated escape"); + } + char escape = text_[position_++]; + switch (escape) + { + case '"': + out.push_back('"'); + break; + case '\\': + out.push_back('\\'); + break; + case '/': + out.push_back('/'); + break; + case 'b': + out.push_back('\b'); + break; + case 'f': + out.push_back('\f'); + break; + case 'n': + out.push_back('\n'); + break; + case 'r': + out.push_back('\r'); + break; + case 't': + out.push_back('\t'); + break; + case 'u': + append_unicode(out); + break; + default: + fail("invalid escape"); + } + } + else + { + out.push_back(c); + } + } + return out; + } + + unsigned read_hex4() + { + if (position_ + 4 > text_.size()) + { + fail("truncated \\u escape"); + } + unsigned code = 0; + for (int i = 0; i < 4; ++i) + { + char c = text_[position_++]; + code <<= 4; + if (c >= '0' && c <= '9') + { + code |= static_cast(c - '0'); + } + else if (c >= 'a' && c <= 'f') + { + code |= static_cast(c - 'a' + 10); + } + else if (c >= 'A' && c <= 'F') + { + code |= static_cast(c - 'A' + 10); + } + else + { + fail("invalid hex digit in \\u escape"); + } + } + return code; + } + + void append_unicode(std::string &out) + { + unsigned code = read_hex4(); + if (code >= 0xD800 && code <= 0xDBFF) + { + // High surrogate: a low surrogate must follow to form a code point above the BMP. + if (position_ + 2 > text_.size() || text_[position_] != '\\' || text_[position_ + 1] != 'u') + { + fail("lone high surrogate in \\u escape"); + } + position_ += 2; + unsigned low = read_hex4(); + if (low < 0xDC00 || low > 0xDFFF) + { + fail("invalid low surrogate in \\u escape"); + } + code = 0x10000 + ((code - 0xD800) << 10) + (low - 0xDC00); + } + append_utf8(out, code); + } + + static void append_utf8(std::string &out, unsigned code) + { + if (code <= 0x7F) + { + out.push_back(static_cast(code)); + } + else if (code <= 0x7FF) + { + out.push_back(static_cast(0xC0 | (code >> 6))); + out.push_back(static_cast(0x80 | (code & 0x3F))); + } + else if (code <= 0xFFFF) + { + out.push_back(static_cast(0xE0 | (code >> 12))); + out.push_back(static_cast(0x80 | ((code >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (code & 0x3F))); + } + else + { + out.push_back(static_cast(0xF0 | (code >> 18))); + out.push_back(static_cast(0x80 | ((code >> 12) & 0x3F))); + out.push_back(static_cast(0x80 | ((code >> 6) & 0x3F))); + out.push_back(static_cast(0x80 | (code & 0x3F))); + } + } +}; + +} // namespace detail + +// Parses `text` as JSON. Throws ParseError on malformed input. +inline Value parse(const std::string &text) +{ + return detail::Parser(text).parse(); +} + +} // namespace neural::json diff --git a/engine/neural/neural_decoder.cpp b/engine/neural/neural_decoder.cpp new file mode 100644 index 00000000..13b7e3d7 --- /dev/null +++ b/engine/neural/neural_decoder.cpp @@ -0,0 +1,103 @@ +#include "neural_decoder.h" + +#include "sentence_model.h" + +#include +#include +#include +#include +#include +#include + +namespace neural +{ +namespace +{ + +// The lattice reports log10 (kenlm's unit); the model reports natural log. Interpolating the two +// without this would silently scale lambda by 2.3. +constexpr double kLog10ToNats = 2.302585092994046; + +} // namespace + +std::string last_characters(const std::string &text, std::size_t count) +{ + if (count == 0) + { + return {}; + } + std::size_t seen = 0; + std::size_t start = text.size(); + for (std::size_t i = text.size(); i-- > 0;) + { + if ((static_cast(text[i]) & 0xC0) == 0x80) + { + continue; // continuation byte + } + start = i; + if (++seen == count) + { + break; + } + } + return text.substr(start); +} + +const SentenceModel *shared_sentence_model(const std::string &path) +{ + static std::mutex mutex; + // Value may be null: an entry records that loading `path` was attempted, so a missing file is + // not re-read on every keystroke. + static std::map> cache; + + std::lock_guard lock(mutex); + auto it = cache.find(path); + if (it == cache.end()) + { + it = cache.emplace(path, SentenceModel::load_file(path)).first; + } + return it->second.get(); +} + +std::vector rerank_order(const SentenceModel &model, const std::string &context, + const std::vector &sentences, + const std::vector &static_scores, const RerankOptions &options) +{ + if (sentences.size() < 2 || sentences.size() != static_scores.size() || options.lambda == 0.0) + { + return {}; + } + + const std::size_t scored_count = std::min(sentences.size(), options.max_paths); + if (scored_count < 2) + { + return {}; + } + + const std::vector texts(sentences.begin(), + sentences.begin() + static_cast(scored_count)); + const std::vector neural = model.score_sentences(last_characters(context, options.context_chars), texts); + if (neural.size() != scored_count) + { + // The model failed or was inconsistent. Half a reranking is worse than none, so the caller + // keeps its own order. + return {}; + } + + std::vector combined(scored_count, 0.0); + for (std::size_t i = 0; i < scored_count; ++i) + { + const double statik = static_scores[i] * kLog10ToNats; + combined[i] = statik + options.lambda * (neural[i] - statik); + } + + std::vector order(sentences.size()); + std::iota(order.begin(), order.end(), std::size_t{0}); + // Stable so that sentences the two models score identically keep the lattice's ordering, and so + // that the unscored tail past max_paths stays as it was. + std::stable_sort(order.begin(), order.begin() + static_cast(scored_count), + [&](std::size_t a, std::size_t b) { return combined[a] > combined[b]; }); + return order; +} + +} // namespace neural diff --git a/engine/neural/neural_decoder.h b/engine/neural/neural_decoder.h new file mode 100644 index 00000000..5f5822ca --- /dev/null +++ b/engine/neural/neural_decoder.h @@ -0,0 +1,77 @@ +#pragma once + +// Neural reranking of whole-sentence candidates. +// +// The model does not produce sentences. The word lattice (engine/quanpin/word_lattice.*) already +// decodes the typed pinyin into its n-best sentences with a kenlm trigram; this layer asks the +// character-level Transformer how likely each of those sentences is given what the user has already +// committed, and reorders them by +// +// score = static + lambda * (neural - static) +// +// both terms in nats. That is a log-linear interpolation of the two language models, with lambda +// controlling how much the neural model is allowed to override the lattice. It cannot invent a +// sentence the lattice did not propose, which is the point: the lattice is good at coverage and bad +// at long-range choices, and the Transformer is the reverse. +// +// This replaces an earlier design that ran a per-syllable beam search over the model to generate +// sentences directly. That needed one full forward pass per beam per syllable with no K/V reuse -- +// on the order of 10^10 multiply-accumulates for a dozen syllables, seconds of wall clock on the +// typing path. Scoring a small batch of finished sentences against a cached context is far less work, because +// the sentences come from the lattice for free and the batch streams the weights once. +// +// The word lattice always supplies the paths when neural reranking is enabled. Its own candidate +// does not have to be enabled for display: in that mode the lattice is only an internal candidate +// generator and the user sees the neural model's pick. + +#include +#include +#include + +namespace neural +{ + +class SentenceModel; + +struct RerankOptions +{ + // How much of the model's opinion to apply. 0 leaves the lattice order untouched; 1 replaces + // the lattice score outright. 0.5 keeps the lattice's own evidence (user dictionary weights, + // reading priors) in play, which matters because the neural model has never seen them. + double lambda = 0.5; + // Paths to rescore. The lattice's own ordering past the first handful is already noise, and + // every extra path is another row in the batch. + std::size_t max_paths = 12; + // Characters of committed text shown to the model as context. Bounded by the model's own + // window; 64 matches what the model was trained to use. + std::size_t context_chars = 64; +}; + +// Last `count` Unicode characters of `text`, not bytes: the context bound is in characters because +// the model's window is in characters. Exposed so a caller that caches results keyed by context can +// trim to exactly what will be scored -- two histories with the same tail are the same context, and +// keying on the untrimmed text would miss the cache on every commit. +std::string last_characters(const std::string &text, std::size_t count); + +// Loads (once, process-wide) and returns the model at `path`, or nullptr if it is missing or +// invalid. The returned pointer is owned by the cache and lives for the process; a failed load is +// remembered so it is not retried on every keystroke. +const SentenceModel *shared_sentence_model(const std::string &path); + +// The order `sentences` should be shown in, as indices into it. +// +// `static_scores` are the lattice's own path scores in log10, as LatticePath::log_prob reports +// them; they are converted to nats here. `context` is the text already committed before the +// composition, and is truncated to the last `options.context_chars` characters. +// +// Returns an empty vector when the model could not be consulted at all (fewer than two sentences, +// lambda of zero, mismatched inputs). The caller must then leave its own order alone: a partial +// reorder is worse than none, because it mixes two incomparable rankings. +// +// Only the first `options.max_paths` sentences are rescored; any beyond that keep their relative +// order at the end. +std::vector rerank_order(const SentenceModel &model, const std::string &context, + const std::vector &sentences, + const std::vector &static_scores, const RerankOptions &options = {}); + +} // namespace neural diff --git a/engine/neural/rescore_worker.cpp b/engine/neural/rescore_worker.cpp new file mode 100644 index 00000000..e8a21d77 --- /dev/null +++ b/engine/neural/rescore_worker.cpp @@ -0,0 +1,166 @@ +#include "rescore_worker.h" + +#include +#include +#include +#include + +namespace neural +{ +namespace +{ +// 结果表只是给「同一次输入被重查一遍」兜底的，不是长期缓存：装满就整张丢掉，比维护 LRU 简单， +// 代价也只是重算一次。 +constexpr std::size_t kMaxResults = 128; +constexpr std::size_t kWorkerCount = 2; + +std::string make_key(const SentenceModel &model, const std::string &context, const std::vector &sentences) +{ + // \x01 不会出现在 UTF-8 文本里，拿它当分隔符就不会把两批不同的候选拼成同一个键。 + std::string key = std::to_string(reinterpret_cast(&model)); + key.push_back('\x01'); + key += context; + for (const std::string &sentence : sentences) + { + key.push_back('\x01'); + key += sentence; + } + return key; +} +} // namespace + +RescoreWorker &RescoreWorker::instance() +{ + // 故意泄漏：后台线程只碰这个对象的成员，而静态对象的析构发生在退出时线程可能已被系统干掉之 + // 后，join 会挂、detach 会悬空。永不析构把这道选择题去掉，进程退出时操作系统收走就是了。 + static RescoreWorker *worker = new RescoreWorker(); + return *worker; +} + +void RescoreWorker::set_ready_callback(Ready ready) +{ + std::lock_guard lock(mutex_); + ready_ = std::move(ready); +} + +void RescoreWorker::clear() +{ + std::lock_guard lock(mutex_); + done_.clear(); + pending_.clear(); +} + +std::optional> RescoreWorker::order_for(const SentenceModel &model, const std::string &context, + const std::vector &sentences, + const std::vector &static_scores, + const RerankOptions &options) +{ + std::string key = make_key(model, context, sentences); + { + std::lock_guard lock(mutex_); + if (stopping_) + { + return std::nullopt; + } + const auto found = done_.find(key); + if (found != done_.end()) + { + return found->second; + } + Job job; + job.model = &model; + job.key = std::move(key); + job.context = context; + job.sentences = sentences; + job.static_scores = static_scores; + job.options = options; + pending_[&model] = std::move(job); + ensure_threads(); + } + wake_.notify_one(); + return std::nullopt; +} + +void RescoreWorker::shutdown() +{ + std::unique_lock lock(mutex_); + stopping_ = true; + ready_ = nullptr; // 目标很可能已经析构了，这之后一律不再回调 + pending_.clear(); // 排着队还没开始的，直接不要了 + wake_.notify_all(); + idle_.wait(lock, [this] { return busy_workers_ == 0; }); +} + +void RescoreWorker::ensure_threads() +{ + // 调用方持有 mutex_。两个执行槽让速度优先模型不必等效果优先模型打完分。 + // 线程只在第一次真有活干时才起，没开神经重排的用户不会平白多线程。 + if (threads_started_) + { + return; + } + threads_started_ = true; + for (std::size_t i = 0; i < kWorkerCount; ++i) + std::thread([this] { run(); }).detach(); + std::atexit([] { instance().shutdown(); }); +} + +void RescoreWorker::run() +{ + for (;;) + { + Job job; + { + std::unique_lock lock(mutex_); + wake_.wait(lock, [this] { + return stopping_ || std::any_of(pending_.begin(), pending_.end(), [this](const auto &item) { + return active_models_.find(item.first) == active_models_.end(); + }); + }); + if (stopping_) + { + return; + } + auto next = std::find_if(pending_.begin(), pending_.end(), [this](const auto &item) { + return active_models_.find(item.first) == active_models_.end(); + }); + job = std::move(next->second); + pending_.erase(next); + active_models_.insert(job.model); + ++busy_workers_; + } + + // 打分在锁外做：它是这里最慢的一步，占着锁会让按键路径上的 order_for 跟着等。 + std::vector order = + rerank_order(*job.model, job.context, job.sentences, job.static_scores, job.options); + + Ready ready; + { + std::lock_guard lock(mutex_); + // 期间可能又来了新请求，说明用户还在打字，这一份存下来留给「他退回来」的情况。 + if (done_.size() >= kMaxResults) + { + done_.clear(); + } + done_.emplace(std::move(job.key), std::move(order)); + ready = ready_; // shutdown 已经清过就是空的，那就不叫了 + } + + // 回调在锁外发：它会一路走到重查，重查又会调 order_for，在锁内发就是自锁。busy_ 要压到回 + // 调之后才放，否则 shutdown 会在回调正往已析构的目标里走的时候就返回。 + if (ready) + { + ready(); + } + + { + std::lock_guard lock(mutex_); + active_models_.erase(job.model); + --busy_workers_; + } + wake_.notify_all(); + idle_.notify_all(); + } +} + +} // namespace neural diff --git a/engine/neural/rescore_worker.h b/engine/neural/rescore_worker.h new file mode 100644 index 00000000..28ffffb0 --- /dev/null +++ b/engine/neural/rescore_worker.h @@ -0,0 +1,91 @@ +#pragma once +// 神经整句重排的后台线程。 +// +// 给十二条候选打一次分要几十毫秒——比一次按键的预算大一个数量级，所以它不能长在按键路径上。这里 +// 把它挪开：查询照旧同步返回词格的静态顺序，同时把这一批候选交给后台线程；线程算完把顺序存进结 +// 果表，再叫醒调用方，调用方丢掉候选缓存重查一次，这次就能在表里查到顺序了。第一次查询按原顺序 +// 出，几十毫秒后原地重排——用户看到的是候选先出来、再自己理顺，而不是按键卡住。 +// +// 「每个模型最新的赢」：用户还在打字时，同一模型前几次请求的结果已经没人要，所以新请求直接盖掉 +// 该模型的旧待办。不同模型各有一个待办位和独立执行槽，速度优先不会排在效果优先后面等待。 +// +// 结果表按 (模型, 前文, 候选列表) 做键，不含静态分：同一批候选的静态分来自同一次词格解码， +// 必然相同。模型必须入键，否则两个模型会误读彼此的排序结果。 + +#include "neural_decoder.h" +#include "sentence_model.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace neural +{ + +class RescoreWorker +{ + public: + using Ready = std::function; + + // 进程唯一，且故意不析构（见 .cpp）。 + static RescoreWorker &instance(); + + // 把「算完了」的通知接出去，调用方据此重查。不设也行：结果照样算、照样进表，只是要等下一次 + // 按键才被取走。 + void set_ready_callback(Ready ready); + + // 有现成结果就返回；没有就把这一批排上队并返回空，调用方这一次照静态顺序走。 + std::optional> order_for(const SentenceModel &model, const std::string &context, + const std::vector &sentences, + const std::vector &static_scores, + const RerankOptions &options); + + // 上屏、换方案、换前文：存着的顺序对新语境没有意义了。 + void clear(); + + // 停工，并等当前这一批算完。 + // + // 必须有：进程退出时先跑静态析构、再终止别的线程，中间这一段里后台线程还活着，手上却握着已 + // 经析构掉的模型表和回调目标（回调一路通到服务端的任务队列）。第一次起线程时用 atexit 挂上 + // 它——atexit 是后注册先执行，而模型表、任务队列都是更早注册的，正好排在它后面。 + void shutdown(); + + private: + RescoreWorker() = default; + ~RescoreWorker() = delete; + RescoreWorker(const RescoreWorker &) = delete; + RescoreWorker &operator=(const RescoreWorker &) = delete; + + struct Job + { + const SentenceModel *model = nullptr; + std::string key; + std::string context; + std::vector sentences; + std::vector static_scores; + RerankOptions options; + }; + + void ensure_threads(); + void run(); + + std::mutex mutex_; + std::condition_variable wake_; // 有活干了，或者该收工了 + std::condition_variable idle_; // 手上这一批算完了 + bool threads_started_ = false; + std::size_t busy_workers_ = 0; + bool stopping_ = false; + std::map pending_; + std::set active_models_; + std::map> done_; + Ready ready_; +}; + +} // namespace neural diff --git a/engine/neural/sentence_model.cpp b/engine/neural/sentence_model.cpp new file mode 100644 index 00000000..921289c5 --- /dev/null +++ b/engine/neural/sentence_model.cpp @@ -0,0 +1,1007 @@ +#include "sentence_model.h" + +#include "json.h" + +#include +#include +#include +#include +#include +#include + +namespace neural +{ +namespace +{ + +// The exact GELU, matching torch.nn.functional.gelu with its default settings. The tanh +// approximation is a different function and the model was not trained with it. +float erf_approx(float x) +{ + // Abramowitz and Stegun 7.1.26; error stays below 1.5e-7 across the range. + float sign = x < 0.0f ? -1.0f : 1.0f; + x = std::fabs(x); + float t = 1.0f / (1.0f + 0.327591100f * x); + float poly = t * (0.254829600f + t * (-0.284496740f + t * (1.421413700f + t * (-1.453152000f + t * 1.061405400f)))); + return sign * (1.0f - poly * std::exp(-x * x)); +} + +float gelu(float x) +{ + constexpr float kInvSqrt2 = 0.70710678118654752440f; + return 0.5f * x * (1.0f + erf_approx(x * kInvSqrt2)); +} + +// Four independent accumulators, not one. +// +// This is where essentially all of the model's time goes, so the shape of the loop matters. A single +// running sum makes every iteration wait on the previous add: the FP unit has a ~4-cycle add latency +// but can start one per cycle, so a serial chain leaves it idle three cycles out of four, and the +// compiler cannot fix it — reassociating floating-point addition changes the result, which /fp:precise +// forbids it from doing on its own. Splitting the sum into four lanes that only meet at the end is +// that same reassociation, written out in the source where it is allowed, and it also gives the +// vectorizer a shape it can fold into packed adds. +// +// The trade is that the result differs from a strict left-to-right sum in the last bits. That is fine +// here — these are log probabilities compared against each other, and the parity test against the +// per-character reference path has room to spare — but it does mean this is not bit-identical to a +// naive reference implementation. +float dot(const float *a, const float *b, std::size_t n) +{ + float s0 = 0.0f; + float s1 = 0.0f; + float s2 = 0.0f; + float s3 = 0.0f; + std::size_t i = 0; + for (; i + 4 <= n; i += 4) + { + s0 += a[i] * b[i]; + s1 += a[i + 1] * b[i + 1]; + s2 += a[i + 2] * b[i + 2]; + s3 += a[i + 3] * b[i + 3]; + } + float total = (s0 + s1) + (s2 + s3); + for (; i < n; ++i) + { + total += a[i] * b[i]; + } + return total; +} + +float dot_i8(const float *a, const std::int8_t *b, std::size_t n) +{ + float total = 0.0f; + for (std::size_t i = 0; i < n; ++i) + { + total += a[i] * static_cast(b[i]); + } + return total; +} + +// out[m][o] = dot(x[m], w[o]) + bias[o], both operands row-major. +// +// Output-major on purpose: the weights are far larger than cache (the desktop model's are ~25 MB +// total), so with the rows on the outside every input row restreams the entire matrix from memory +// and the whole forward pass runs at DRAM bandwidth. Hoisting the output loop reads each weight row +// once and reuses it across the batch, which is the reason a batch of candidates costs barely more +// than one. +std::vector linear(const std::vector &x, const Matrix &w, const std::vector &bias, + std::size_t inputs, std::size_t outputs) +{ + std::size_t rows = inputs == 0 ? 0 : x.size() / inputs; + std::vector out(rows * outputs, 0.0f); + std::vector scratch(w.quantized ? inputs : 0, 0.0f); + for (std::size_t o = 0; o < outputs; ++o) + { + // Widen the row's int8 once per output row instead of once per (output row, input row). + // The conversion is what keeps the inner product from vectorizing, and across a batch it + // would otherwise be repeated for every row of the batch. + // + // The per-row scale still multiplies the finished dot product rather than the weights, the + // same way Matrix::dot_row does it: widening an int8 is exact, but folding the scale into + // each weight beforehand would round every element. + const float *weights; + float scale = 1.0f; + if (w.quantized) + { + const std::int8_t *source = &w.ints[o * inputs]; + for (std::size_t i = 0; i < inputs; ++i) + { + scratch[i] = static_cast(source[i]); + } + weights = scratch.data(); + scale = w.scales[o]; + } + else + { + weights = &w.floats[o * inputs]; + } + + const float b = bias.empty() ? 0.0f : bias[o]; + for (std::size_t row = 0; row < rows; ++row) + { + out[row * outputs + o] = dot(&x[row * inputs], weights, inputs) * scale + b; + } + } + return out; +} + +std::vector layer_norm(const float *x, std::size_t rows, const std::vector &weight, + const std::vector &bias, std::size_t width) +{ + std::vector out(rows * width, 0.0f); + for (std::size_t row = 0; row < rows; ++row) + { + const float *source = &x[row * width]; + float *target = &out[row * width]; + float mean = 0.0f; + for (std::size_t i = 0; i < width; ++i) + { + mean += source[i]; + } + mean /= static_cast(width); + float variance = 0.0f; + for (std::size_t i = 0; i < width; ++i) + { + float d = source[i] - mean; + variance += d * d; + } + variance /= static_cast(width); + float inverse = 1.0f / std::sqrt(variance + 1e-5f); + for (std::size_t i = 0; i < width; ++i) + { + target[i] = (source[i] - mean) * inverse * weight[i] + bias[i]; + } + } + return out; +} + +void split_qkv(const std::vector &qkv, std::size_t rows, std::size_t n_embd, std::vector &q, + std::vector &k, std::vector &v) +{ + q.assign(rows * n_embd, 0.0f); + k.assign(rows * n_embd, 0.0f); + v.assign(rows * n_embd, 0.0f); + for (std::size_t row = 0; row < rows; ++row) + { + const float *source = &qkv[row * 3 * n_embd]; + std::memcpy(&q[row * n_embd], source, n_embd * sizeof(float)); + std::memcpy(&k[row * n_embd], source + n_embd, n_embd * sizeof(float)); + std::memcpy(&v[row * n_embd], source + 2 * n_embd, n_embd * sizeof(float)); + } +} + +// IEEE 754 binary16 to binary32, with subnormal handling rather than flush-to-zero. +float half_to_f32(std::uint16_t bits) +{ + std::uint32_t sign = static_cast(bits & 0x8000u) << 16; + std::uint32_t exponent = static_cast((bits >> 10) & 0x1fu); + std::uint32_t mantissa = static_cast(bits & 0x03ffu); + std::uint32_t out; + if (exponent == 0) + { + if (mantissa == 0) + { + out = sign; + } + else + { + unsigned leading = 0; + std::uint32_t m = mantissa; + // Count leading zeros of a 32-bit value. + while ((m & 0x80000000u) == 0) + { + m <<= 1; + ++leading; + } + std::uint32_t exp = 134u - leading; + std::uint32_t mant = (mantissa << (leading - 8)) & 0x007fffffu; + out = sign | (exp << 23) | mant; + } + } + else if (exponent == 0x1f) + { + out = sign | 0x7f800000u | (mantissa << 13); + } + else + { + out = sign | ((exponent + 127u - 15u) << 23) | (mantissa << 13); + } + float result; + std::memcpy(&result, &out, sizeof(result)); + return result; +} + +// Decodes UTF-8 text into Unicode code points. Malformed bytes map to U+FFFD. +std::vector decode_utf8(const std::string &text) +{ + std::vector out; + std::size_t i = 0; + const std::size_t n = text.size(); + while (i < n) + { + unsigned char c = static_cast(text[i]); + char32_t code; + std::size_t extra; + if (c < 0x80) + { + code = c; + extra = 0; + } + else if ((c & 0xE0) == 0xC0) + { + code = c & 0x1F; + extra = 1; + } + else if ((c & 0xF0) == 0xE0) + { + code = c & 0x0F; + extra = 2; + } + else if ((c & 0xF8) == 0xF0) + { + code = c & 0x07; + extra = 3; + } + else + { + out.push_back(0xFFFD); + ++i; + continue; + } + if (i + extra >= n) + { + out.push_back(0xFFFD); + break; + } + bool ok = true; + for (std::size_t j = 1; j <= extra; ++j) + { + unsigned char cc = static_cast(text[i + j]); + if ((cc & 0xC0) != 0x80) + { + ok = false; + break; + } + code = (code << 6) | (cc & 0x3F); + } + if (!ok) + { + out.push_back(0xFFFD); + ++i; + continue; + } + out.push_back(code); + i += extra + 1; + } + return out; +} + +// A tensor descriptor from the safetensors header. +struct Entry +{ + std::string dtype; + std::vector shape; + std::size_t off0 = 0; + std::size_t off1 = 0; +}; + +std::vector read_f32(const std::uint8_t *data, std::size_t data_len, const Entry &entry) +{ + if (entry.off1 > data_len || entry.off0 > entry.off1) + { + throw std::runtime_error("tensor offsets out of range"); + } + std::size_t bytes = entry.off1 - entry.off0; + std::size_t count = bytes / 4; + std::vector out(count); + const std::uint8_t *p = data + entry.off0; + for (std::size_t i = 0; i < count; ++i) + { + std::memcpy(&out[i], p + i * 4, 4); + } + return out; +} + +std::vector read_f16(const std::uint8_t *data, std::size_t data_len, const Entry &entry) +{ + if (entry.off1 > data_len || entry.off0 > entry.off1) + { + throw std::runtime_error("tensor offsets out of range"); + } + std::size_t bytes = entry.off1 - entry.off0; + std::size_t count = bytes / 2; + std::vector out(count); + const std::uint8_t *p = data + entry.off0; + for (std::size_t i = 0; i < count; ++i) + { + std::uint16_t bits; + std::memcpy(&bits, p + i * 2, 2); + out[i] = half_to_f32(bits); + } + return out; +} + +std::size_t product(const std::vector &shape) +{ + std::size_t total = 1; + for (std::size_t d : shape) + { + total *= d; + } + return total; +} + +} // namespace + +float Matrix::dot_row(const float *x, std::size_t row, std::size_t width) const +{ + if (quantized) + { + return dot_i8(x, &ints[row * width], width) * scales[row]; + } + return dot(x, &floats[row * width], width); +} + +void Matrix::row_into(std::size_t index, std::size_t width, float *out) const +{ + if (quantized) + { + float scale = scales[index]; + const std::int8_t *p = &ints[index * width]; + for (std::size_t i = 0; i < width; ++i) + { + out[i] = static_cast(p[i]) * scale; + } + } + else + { + std::memcpy(out, &floats[index * width], width * sizeof(float)); + } +} + +std::unique_ptr SentenceModel::load(const std::vector &bytes) +{ + try + { + if (bytes.size() < 8) + { + return nullptr; + } + std::uint64_t header_len = 0; + std::memcpy(&header_len, bytes.data(), 8); + std::size_t header_end = 8 + static_cast(header_len); + if (header_len == 0 || header_end > bytes.size()) + { + return nullptr; + } + std::string header_text(reinterpret_cast(bytes.data() + 8), static_cast(header_len)); + json::Value header = json::parse(header_text); + if (!header.is_object()) + { + return nullptr; + } + const std::uint8_t *data = bytes.data() + header_end; + std::size_t data_len = bytes.size() - header_end; + + // Metadata. + const json::Value *metadata = header.find("__metadata__"); + if (metadata == nullptr || !metadata->is_object()) + { + return nullptr; + } + auto metadata_string = [&](const std::string &key) -> const std::string * { + const json::Value *value = metadata->find(key); + if (value == nullptr || !value->is_string()) + { + return nullptr; + } + return &value->string; + }; + const std::string *config_text = metadata_string("config"); + const std::string *vocab_text = metadata_string("vocab"); + if (config_text == nullptr || vocab_text == nullptr) + { + return nullptr; + } + + auto model = std::unique_ptr(new SentenceModel()); + + json::Value config_json = json::parse(*config_text); + auto config_uint = [&](const std::string &key) -> std::size_t { + const json::Value *value = config_json.find(key); + if (value == nullptr || !value->is_number()) + { + throw std::runtime_error("config missing " + key); + } + return static_cast(value->number); + }; + model->config_.vocab = config_uint("vocab"); + model->config_.n_layer = config_uint("n_layer"); + model->config_.n_head = config_uint("n_head"); + model->config_.n_embd = config_uint("n_embd"); + model->config_.context = config_uint("context"); + const Config &config = model->config_; + if (config.n_head == 0 || config.n_embd == 0 || config.n_embd % config.n_head != 0 || config.vocab == 0 || + config.n_layer == 0 || config.context == 0) + { + return nullptr; + } + + json::Value vocab_json = json::parse(*vocab_text); + if (!vocab_json.is_array() || vocab_json.array.size() != config.vocab) + { + return nullptr; + } + for (std::size_t id = 0; id < vocab_json.array.size(); ++id) + { + const json::Value &item = vocab_json.array[id]; + if (!item.is_string()) + { + continue; + } + std::vector chars = decode_utf8(item.string); + // Reserved tokens ( etc.) are addressed by constant, not by text; single code + // points are the character vocabulary. + if (chars.size() == 1) + { + model->index_[chars[0]] = static_cast(id); + } + } + + if (const std::string *attribution = metadata_string("attribution")) + { + model->attribution_ = *attribution; + } + if (const std::string *license = metadata_string("license")) + { + model->license_ = *license; + } + + // Collect every tensor descriptor (including the .scale companions). + std::unordered_map entries; + for (const auto &field : header.object) + { + if (field.first == "__metadata__") + { + continue; + } + const json::Value &desc = field.second; + if (!desc.is_object()) + { + continue; + } + Entry entry; + const json::Value *dtype = desc.find("dtype"); + const json::Value *shape = desc.find("shape"); + const json::Value *offsets = desc.find("data_offsets"); + if (dtype == nullptr || !dtype->is_string() || offsets == nullptr || !offsets->is_array() || + offsets->array.size() != 2) + { + return nullptr; + } + entry.dtype = dtype->string; + if (shape != nullptr && shape->is_array()) + { + for (const json::Value &dim : shape->array) + { + entry.shape.push_back(static_cast(dim.number)); + } + } + entry.off0 = static_cast(offsets->array[0].number); + entry.off1 = static_cast(offsets->array[1].number); + entries.emplace(field.first, std::move(entry)); + } + + auto materialize = [&](const std::string &name) -> Matrix { + auto it = entries.find(name); + if (it == entries.end()) + { + throw std::runtime_error("missing tensor " + name); + } + const Entry &entry = it->second; + Matrix matrix; + std::size_t expected = product(entry.shape); + if (entry.dtype == "F32") + { + matrix.floats = read_f32(data, data_len, entry); + matrix.quantized = false; + if (matrix.floats.size() != expected) + { + throw std::runtime_error(name + ": shape does not match data"); + } + } + else if (entry.dtype == "F16") + { + matrix.floats = read_f16(data, data_len, entry); + matrix.quantized = false; + if (matrix.floats.size() != expected) + { + throw std::runtime_error(name + ": shape does not match data"); + } + } + else if (entry.dtype == "I8") + { + auto scale_it = entries.find(name + ".scale"); + if (scale_it == entries.end()) + { + throw std::runtime_error(name + ": missing scale"); + } + matrix.scales = read_f32(data, data_len, scale_it->second); + std::size_t rows = entry.shape.empty() ? 1 : entry.shape.front(); + if (rows == 0 || matrix.scales.size() != rows) + { + throw std::runtime_error(name + ": scale length does not match rows"); + } + if (entry.off1 > data_len || entry.off0 > entry.off1) + { + throw std::runtime_error(name + ": offsets out of range"); + } + std::size_t count = entry.off1 - entry.off0; + matrix.ints.resize(count); + std::memcpy(matrix.ints.data(), data + entry.off0, count); + matrix.quantized = true; + if (matrix.ints.size() != expected) + { + throw std::runtime_error(name + ": shape does not match data"); + } + } + else + { + throw std::runtime_error(name + ": unsupported dtype " + entry.dtype); + } + return matrix; + }; + + auto materialize_floats = [&](const std::string &name) -> std::vector { + Matrix matrix = materialize(name); + if (matrix.quantized) + { + throw std::runtime_error(name + " is quantized but must be float"); + } + return std::move(matrix.floats); + }; + + model->token_ = materialize("tok.weight"); + model->position_ = materialize_floats("pos.weight"); + model->blocks_.reserve(config.n_layer); + for (std::size_t layer = 0; layer < config.n_layer; ++layer) + { + std::string prefix = "blocks." + std::to_string(layer) + "."; + Block block; + block.ln1_weight = materialize_floats(prefix + "ln1.weight"); + block.ln1_bias = materialize_floats(prefix + "ln1.bias"); + block.qkv_weight = materialize(prefix + "qkv.weight"); + block.qkv_bias = materialize_floats(prefix + "qkv.bias"); + block.proj_weight = materialize(prefix + "proj.weight"); + block.proj_bias = materialize_floats(prefix + "proj.bias"); + block.ln2_weight = materialize_floats(prefix + "ln2.weight"); + block.ln2_bias = materialize_floats(prefix + "ln2.bias"); + block.fc_weight = materialize(prefix + "fc.weight"); + block.fc_bias = materialize_floats(prefix + "fc.bias"); + block.out_weight = materialize(prefix + "out.weight"); + block.out_bias = materialize_floats(prefix + "out.bias"); + model->blocks_.push_back(std::move(block)); + } + model->final_weight_ = materialize_floats("ln_f.weight"); + model->final_bias_ = materialize_floats("ln_f.bias"); + + // Geometry checks: any mismatch means a malformed file, so refuse it. + const std::size_t n_embd = config.n_embd; + if (model->token_.len() != config.vocab * n_embd || model->position_.size() != config.context * n_embd || + model->final_weight_.size() != n_embd || model->final_bias_.size() != n_embd) + { + return nullptr; + } + for (const Block &block : model->blocks_) + { + if (block.qkv_weight.len() != 3 * n_embd * n_embd || block.proj_weight.len() != n_embd * n_embd || + block.fc_weight.len() != 4 * n_embd * n_embd || block.out_weight.len() != 4 * n_embd * n_embd) + { + return nullptr; + } + } + return model; + } + catch (const std::exception &) + { + return nullptr; + } +} + +std::unique_ptr SentenceModel::load_file(const std::string &path) +{ + std::ifstream file(path, std::ios::binary); + if (!file) + { + return nullptr; + } + file.seekg(0, std::ios::end); + std::streamoff size = file.tellg(); + if (size <= 0) + { + return nullptr; + } + file.seekg(0, std::ios::beg); + std::vector bytes(static_cast(size)); + if (!file.read(reinterpret_cast(bytes.data()), size)) + { + return nullptr; + } + return load(bytes); +} + +std::optional SentenceModel::token(char32_t character) const +{ + auto it = index_.find(character); + if (it == index_.end()) + { + return std::nullopt; + } + return it->second; +} + +std::vector SentenceModel::encode(const std::string &text) const +{ + std::vector chars = decode_utf8(text); + std::vector out; + out.reserve(chars.size()); + for (char32_t ch : chars) + { + auto it = index_.find(ch); + out.push_back(it == index_.end() ? kUnk : it->second); + } + return out; +} + +std::vector SentenceModel::embed(const std::vector &tokens, std::size_t width, + std::size_t offset) const +{ + const std::size_t n_embd = config_.n_embd; + std::vector x(tokens.size() * n_embd, 0.0f); + std::vector token_row(n_embd, 0.0f); + for (std::size_t step = 0; step < tokens.size(); ++step) + { + std::size_t token = std::min(tokens[step], config_.vocab - 1); + // Position restarts at `offset` on every batch row: each row is a separate continuation of + // the same cached prefix, not a continuation of the row before it. + std::size_t within = width == 0 ? step : step % width; + std::size_t position = std::min(offset + within, config_.context - 1); + token_.row_into(token, n_embd, token_row.data()); + const float *position_row = &position_[position * n_embd]; + float *row = &x[step * n_embd]; + for (std::size_t i = 0; i < n_embd; ++i) + { + row[i] = token_row[i] + position_row[i]; + } + } + return x; +} + +void SentenceModel::attention(const std::vector &q, const std::vector &k, const std::vector &v, + std::size_t batch, std::size_t width, const std::vector *past_k, + const std::vector *past_v, std::size_t past_len, std::vector &out) const +{ + const std::size_t n_embd = config_.n_embd; + const std::size_t heads = config_.n_head; + const std::size_t head_width = n_embd / heads; + const float scale = 1.0f / std::sqrt(static_cast(head_width)); + if (past_k == nullptr || past_v == nullptr) + { + past_len = 0; + } + out.assign(batch * width * n_embd, 0.0f); + std::vector scores(past_len + width, 0.0f); + for (std::size_t row = 0; row < batch; ++row) + { + const std::size_t base_row = row * width; + for (std::size_t step = 0; step < width; ++step) + { + // Causal, in two runs: the whole shared prefix, then this row's own tokens up to and + // including `step`. Rows never see each other. + const std::size_t visible_self = step + 1; + const std::size_t visible = past_len + visible_self; + for (std::size_t head = 0; head < heads; ++head) + { + const std::size_t base = head * head_width; + const float *query = &q[(base_row + step) * n_embd + base]; + float highest = -std::numeric_limits::infinity(); + for (std::size_t j = 0; j < past_len; ++j) + { + float s = dot(query, &(*past_k)[j * n_embd + base], head_width) * scale; + scores[j] = s; + highest = std::max(highest, s); + } + for (std::size_t j = 0; j < visible_self; ++j) + { + float s = dot(query, &k[(base_row + j) * n_embd + base], head_width) * scale; + scores[past_len + j] = s; + highest = std::max(highest, s); + } + float sum = 0.0f; + for (std::size_t j = 0; j < visible; ++j) + { + scores[j] = std::exp(scores[j] - highest); + sum += scores[j]; + } + float *target = &out[(base_row + step) * n_embd + base]; + for (std::size_t j = 0; j < past_len; ++j) + { + float weight = scores[j] / sum; + const float *value = &(*past_v)[j * n_embd + base]; + for (std::size_t i = 0; i < head_width; ++i) + { + target[i] += weight * value[i]; + } + } + for (std::size_t j = 0; j < visible_self; ++j) + { + float weight = scores[past_len + j] / sum; + const float *value = &v[(base_row + j) * n_embd + base]; + for (std::size_t i = 0; i < head_width; ++i) + { + target[i] += weight * value[i]; + } + } + } + } + } +} + +void SentenceModel::feed_forward(std::vector &x, const Block &block) const +{ + const std::size_t n_embd = config_.n_embd; + std::size_t rows = x.size() / n_embd; + std::vector normed = layer_norm(x.data(), rows, block.ln2_weight, block.ln2_bias, n_embd); + std::vector hidden = linear(normed, block.fc_weight, block.fc_bias, n_embd, 4 * n_embd); + for (float &value : hidden) + { + value = gelu(value); + } + std::vector out = linear(hidden, block.out_weight, block.out_bias, 4 * n_embd, n_embd); + for (std::size_t i = 0; i < x.size(); ++i) + { + x[i] += out[i]; + } +} + +std::vector SentenceModel::forward(const std::vector &tokens, std::size_t batch, + std::size_t width, const PrefixCache *past, PrefixCache *record) const +{ + const std::size_t n_embd = config_.n_embd; + const std::size_t rows = batch * width; + const std::size_t past_len = past == nullptr ? 0 : past->length; + std::vector x = embed(tokens, width, past_len); + std::vector q; + std::vector k; + std::vector v; + std::vector attended; + if (record != nullptr) + { + // A cache is shared by every row that later attends to it, so it can only ever describe one + // sequence. build_prefix is the only caller and always passes batch == 1 from position 0. + record->keys.assign(blocks_.size(), {}); + record->values.assign(blocks_.size(), {}); + record->length = past_len + width; + } + for (std::size_t layer = 0; layer < blocks_.size(); ++layer) + { + const Block &block = blocks_[layer]; + std::vector normed = layer_norm(x.data(), rows, block.ln1_weight, block.ln1_bias, n_embd); + std::vector qkv = linear(normed, block.qkv_weight, block.qkv_bias, n_embd, 3 * n_embd); + split_qkv(qkv, rows, n_embd, q, k, v); + const std::vector *past_k = past_len == 0 ? nullptr : &past->keys[layer]; + const std::vector *past_v = past_len == 0 ? nullptr : &past->values[layer]; + attention(q, k, v, batch, width, past_k, past_v, past_len, attended); + std::vector projected = linear(attended, block.proj_weight, block.proj_bias, n_embd, n_embd); + for (std::size_t i = 0; i < x.size(); ++i) + { + x[i] += projected[i]; + } + feed_forward(x, block); + if (record != nullptr) + { + record->keys[layer] = k; + record->values[layer] = v; + } + } + return layer_norm(x.data(), rows, final_weight_, final_bias_, n_embd); +} + +PrefixCache SentenceModel::build_prefix(const std::vector &tokens) const +{ + PrefixCache cache; + if (tokens.empty()) + { + return cache; + } + forward(tokens, 1, tokens.size(), nullptr, &cache); + return cache; +} + +std::vector SentenceModel::distribution(const std::vector &hidden) const +{ + const std::size_t n_embd = config_.n_embd; + const std::size_t vocab = config_.vocab; + std::vector logits(vocab, 0.0f); + float highest = -std::numeric_limits::infinity(); + for (std::size_t row = 0; row < vocab; ++row) + { + float logit = token_.dot_row(hidden.data(), row, n_embd); + logits[row] = logit; + highest = std::max(highest, logit); + } + float sum = 0.0f; + for (float logit : logits) + { + sum += std::exp(logit - highest); + } + float offset = highest + std::log(sum); + for (float &logit : logits) + { + logit -= offset; + } + return logits; +} + +std::vector SentenceModel::target_log_probs(const std::vector &hidden, std::size_t rows, + const std::vector &targets) const +{ + const std::size_t n_embd = config_.n_embd; + const std::size_t vocab = config_.vocab; + std::vector logits(rows * vocab, 0.0f); + // Output-major over the vocabulary, and widening each row's int8 once, for the same reasons as + // linear(): the tied embedding table is the single largest matrix in the model. + std::vector scratch(token_.quantized ? n_embd : 0, 0.0f); + for (std::size_t id = 0; id < vocab; ++id) + { + const float *weights; + float scale = 1.0f; + if (token_.quantized) + { + const std::int8_t *source = &token_.ints[id * n_embd]; + for (std::size_t i = 0; i < n_embd; ++i) + { + scratch[i] = static_cast(source[i]); + } + weights = scratch.data(); + scale = token_.scales[id]; + } + else + { + weights = &token_.floats[id * n_embd]; + } + for (std::size_t row = 0; row < rows; ++row) + { + logits[row * vocab + id] = dot(&hidden[row * n_embd], weights, n_embd) * scale; + } + } + + std::vector out(rows, 0.0f); + for (std::size_t row = 0; row < rows; ++row) + { + if (targets[row] == kNoTarget) + { + continue; + } + const float *source = &logits[row * vocab]; + float highest = -std::numeric_limits::infinity(); + for (std::size_t id = 0; id < vocab; ++id) + { + highest = std::max(highest, source[id]); + } + float sum = 0.0f; + for (std::size_t id = 0; id < vocab; ++id) + { + sum += std::exp(source[id] - highest); + } + out[row] = source[targets[row]] - (highest + std::log(sum)); + } + return out; +} + +std::vector SentenceModel::next_log_probabilities(const std::string &text) const +{ + std::vector tokens; + tokens.push_back(kBos); + std::vector encoded = encode(text); + std::size_t room = config_.context - 1; + std::size_t start = encoded.size() > room ? encoded.size() - room : 0; + tokens.insert(tokens.end(), encoded.begin() + static_cast(start), encoded.end()); + std::vector hidden = forward(tokens, 1, tokens.size(), nullptr, nullptr); + const std::size_t n_embd = config_.n_embd; + std::vector last(hidden.end() - static_cast(n_embd), hidden.end()); + return distribution(last); +} + +std::vector SentenceModel::score_sentences(const std::string &context, + const std::vector &texts) const +{ + if (texts.empty()) + { + return {}; + } + const std::size_t limit = config_.context; + if (limit < 2) + { + return std::vector(texts.size(), 0.0); + } + + // The sequence each candidate is scored under is + context + candidate. Every candidate + // shares all of that but its own tail, so the shared part is split at its last token: that token + // leads each batch row, because the distribution it produces is what scores the candidate's + // first character. Only what precedes it goes in the cache, since only that stays fixed. + std::vector full; + full.push_back(kBos); + for (std::uint32_t id : encode(context)) + { + full.push_back(id); + } + + std::vector> tails; + tails.reserve(texts.size()); + std::size_t longest = 0; + for (const std::string &text : texts) + { + std::vector ids = encode(text); + // A candidate longer than the window minus the leading token keeps its tail, which is the + // part the remaining context can still inform. + if (ids.size() > limit - 1) + { + ids.erase(ids.begin(), ids.begin() + static_cast(ids.size() - (limit - 1))); + } + longest = std::max(longest, ids.size()); + tails.push_back(std::move(ids)); + } + const std::size_t width = longest + 1; // the shared leading token, then the longest candidate + + const std::uint32_t last = full.back(); + std::vector head(full.begin(), full.end() - 1); + // Head plus the widest row must fit the position table; the oldest context goes first. + const std::size_t room = limit > width ? limit - width : 0; + if (head.size() > room) + { + head.erase(head.begin(), head.begin() + static_cast(head.size() - room)); + } + + const std::size_t batch = tails.size(); + // Short rows are padded with token 0; their positions get kNoTarget and are never read. + std::vector flat(batch * width, 0u); + std::vector targets(batch * width, kNoTarget); + for (std::size_t row = 0; row < batch; ++row) + { + flat[row * width] = last; + for (std::size_t i = 0; i < tails[row].size(); ++i) + { + flat[row * width + 1 + i] = tails[row][i]; + // Position i predicts token i+1, so the leading token's slot carries the first character. + targets[row * width + i] = tails[row][i]; + } + } + + std::vector scored; + { + std::lock_guard lock(cache_mutex_); + if (cached_tokens_ != head) + { + cached_prefix_ = build_prefix(head); + cached_tokens_ = head; + } + std::vector hidden = forward(flat, batch, width, &cached_prefix_, nullptr); + scored = target_log_probs(hidden, batch * width, targets); + } + + std::vector out(batch, 0.0); + for (std::size_t row = 0; row < batch; ++row) + { + double total = 0.0; + for (std::size_t i = 0; i < tails[row].size(); ++i) + { + total += static_cast(scored[row * width + i]); + } + out[row] = total; + } + return out; +} + +} // namespace neural diff --git a/engine/neural/sentence_model.h b/engine/neural/sentence_model.h new file mode 100644 index 00000000..41b2dd3a --- /dev/null +++ b/engine/neural/sentence_model.h @@ -0,0 +1,185 @@ +#pragma once + +// Character-level Transformer "sentence model" from the chinese-ime-lm project, ported to C++. +// +// This is a faithful port of the Apache-2.0 Rust reference implementation +// (chinese-ime-lm/reference/src/{weights.rs,lib.rs}) and the format spec (docs/format.md). The +// model is a decoder-only Pre-LN Transformer with tied input/output embeddings; given some +// committed context it produces a log-probability distribution over the next character. +// +// The primitive the IME actually uses is score_sentences: log P(sentence | context) for a handful +// of already-formed sentences at once, which neural_decoder.* turns into a reranking of the word +// lattice's n-best paths. next_log_probabilities remains as the single-step reference API that the +// parity tests check against the Rust implementation. +// +// Two things make the scorer fast enough to run on the typing path, and both matter far more than +// the arithmetic itself. The weights are ~25 MB, which does not fit in cache, so a forward pass is +// memory-bandwidth bound: every matmul must stream the whole matrix. Therefore (1) the context's +// per-layer K/V is cached (PrefixCache) so an unchanged context is never recomputed, and (2) all +// candidates are run as one batch through each layer, so each weight matrix is streamed once for +// the whole batch instead of once per candidate per token. +// +// Numerical details that MUST match the reference or results are garbage: exact erf-based GELU +// (not the tanh approximation), LayerNorm eps 1e-5, attention scale 1/sqrt(n_embd/n_head), int8 +// per-output-row symmetric dequant (scale applied once after the dot product), and the rule that +// position i predicts token i+1. + +#include +#include +#include +#include +#include +#include +#include + +namespace neural +{ + +// A weight matrix kept in whatever precision the file stored it in. Quantized matrices are not +// expanded on load: the per-output-row scale is a constant factor for the whole row, so it lifts +// out of the dot product and multiplies the result once. +struct Matrix +{ + bool quantized = false; + std::vector floats; // used when !quantized + std::vector ints; // used when quantized + std::vector scales; // per output row, used when quantized + + std::size_t len() const + { + return quantized ? ints.size() : floats.size(); + } + + // dot(x, row) * scale[row], where the scale is 1 for an unquantized matrix. + float dot_row(const float *x, std::size_t row, std::size_t width) const; + + // One row as f32, for the token table read as an embedding as well as an output projection. + void row_into(std::size_t index, std::size_t width, float *out) const; +}; + +struct Config +{ + std::size_t vocab = 0; + std::size_t n_layer = 0; + std::size_t n_head = 0; + std::size_t n_embd = 0; + std::size_t context = 0; +}; + +struct Block +{ + std::vector ln1_weight; + std::vector ln1_bias; + Matrix qkv_weight; + std::vector qkv_bias; + Matrix proj_weight; + std::vector proj_bias; + std::vector ln2_weight; + std::vector ln2_bias; + Matrix fc_weight; + std::vector fc_bias; + Matrix out_weight; + std::vector out_bias; +}; + +// Per-layer keys and values for a fixed run of leading tokens, each [length, n_embd] row-major. +// Every batch row of a later segment attends to the same cache, so a context shared by all the +// candidates is paid for once. +struct PrefixCache +{ + std::vector> keys; // one entry per layer + std::vector> values; // one entry per layer + std::size_t length = 0; // tokens covered, i.e. the position the next token takes +}; + +class SentenceModel +{ + public: + // Loads a model from an in-memory safetensors file. Returns nullptr on any malformed input + // rather than throwing, so a missing or corrupt model leaves the IME working without it. + static std::unique_ptr load(const std::vector &bytes); + + // Loads a model from disk. Returns nullptr when the file is absent or unreadable/invalid. + static std::unique_ptr load_file(const std::string &path); + + std::size_t context_length() const + { + return config_.context; + } + const std::string &attribution() const + { + return attribution_; + } + const std::string &license() const + { + return license_; + } + + // The token id this model uses for a character (a Unicode code point), if any. + std::optional token(char32_t character) const; + + // log P(next character | text), over the whole vocabulary. Index by token id. + std::vector next_log_probabilities(const std::string &text) const; + + // Sum over characters of log P(character | context and everything before it in the text), one + // entry per text. Characters outside the vocabulary score as , so a text is never dropped. + // + // The whole batch is one pass: `context` (minus its last token) becomes a PrefixCache reused + // across calls while the context is unchanged, and the texts are padded to equal length and run + // together. An empty `texts` returns empty. A text longer than the model's context window is + // truncated from the left, as is the context when the two together do not fit. + // + // Thread-safe: the prefix cache is guarded, and concurrent callers with different contexts are + // correct but will evict each other. + std::vector score_sentences(const std::string &context, const std::vector &texts) const; + + private: + static constexpr std::uint32_t kUnk = 1; + static constexpr std::uint32_t kBos = 2; + + Config config_; + std::unordered_map index_; + Matrix token_; + std::vector position_; + std::vector blocks_; + std::vector final_weight_; + std::vector final_bias_; + std::string attribution_; + std::string license_; + + // Guards the cached prefix below. mutable so scoring stays const. + mutable std::mutex cache_mutex_; + mutable std::vector cached_tokens_; + mutable PrefixCache cached_prefix_; + + std::vector encode(const std::string &text) const; + std::vector embed(const std::vector &tokens, std::size_t width, std::size_t offset) const; + + // Runs `batch` rows of `width` tokens (row-major in `tokens`, padded to equal length by the + // caller) at positions [past_len, past_len + width), where every row additionally attends to + // `past`. Returns the final-LayerNorm'd hidden states, [batch * width, n_embd]. + // + // When `record` is non-null the K/V produced here is stored into it, which requires batch == 1 + // (a cache is shared by later rows, so it may only ever describe one sequence). + std::vector forward(const std::vector &tokens, std::size_t batch, std::size_t width, + const PrefixCache *past, PrefixCache *record) const; + + // Builds the K/V cache for `tokens`, which are taken to start at position 0. + PrefixCache build_prefix(const std::vector &tokens) const; + + // log-softmax over the vocabulary for one hidden state. + std::vector distribution(const std::vector &hidden) const; + + // For each of `rows` hidden states, log P(targets[row]) under that state's distribution. Rows + // with target kNoTarget are skipped and report 0. + static constexpr std::uint32_t kNoTarget = 0xFFFFFFFFu; + std::vector target_log_probs(const std::vector &hidden, std::size_t rows, + const std::vector &targets) const; + + void attention(const std::vector &q, const std::vector &k, const std::vector &v, + std::size_t batch, std::size_t width, const std::vector *past_k, + const std::vector *past_v, std::size_t past_len, std::vector &out) const; + void feed_forward(std::vector &x, const Block &block) const; +}; + +} // namespace neural diff --git a/engine/quanpin/engine.cpp b/engine/quanpin/engine.cpp index e5cd0fdd..994b54da 100644 --- a/engine/quanpin/engine.cpp +++ b/engine/quanpin/engine.cpp @@ -38,6 +38,11 @@ std::vector QuanpinEngine::query(const QueryRequest &request) return {}; } + // 整句联想开关随请求下发到词典层，四条整句来源据此各自门控（见 query_series）。 + dictionary_.set_sentence_association(request.sentence_association); + // 前文同理随请求下发；换了前文，按旧前文重排出来的缓存顺序作废。 + dictionary_.set_rescoring_context(request.rescoring_context); + const size_t helpcode_length = request.enable_quanpin_helpcode ? quanpin::detect_active_helpcode_length(request.raw_input, request.raw_input_with_cases) diff --git a/engine/quanpin/lattice_rerank.cpp b/engine/quanpin/lattice_rerank.cpp new file mode 100644 index 00000000..3bc00f57 --- /dev/null +++ b/engine/quanpin/lattice_rerank.cpp @@ -0,0 +1,56 @@ +#include "lattice_rerank.h" + +#include "../neural/rescore_worker.h" + +#include +#include + +namespace quanpin +{ + +LatticeReranker make_neural_reranker(const neural::SentenceModel *model, const std::string &context, + const neural::RerankOptions &options) +{ + if (model == nullptr) + { + return {}; + } + // 在这里就把前文裁到模型真正会看的那几十个字：后台重排的结果表以「前文 + 候选句」为键， + // 拿没裁过的整段历史当键的话，每上屏一个字键就变了，缓存永远打不中。 + const std::string trimmed = neural::last_characters(context, options.context_chars); + return [model, context = trimmed, options](std::vector &paths) -> bool { + if (paths.size() < 2) + { + return false; + } + std::vector sentences; + std::vector scores; + sentences.reserve(paths.size()); + scores.reserve(paths.size()); + for (const LatticePath &path : paths) + { + sentences.push_back(path.sentence); + scores.push_back(path.log_prob); + } + + // 不在这里打分：交给后台线程，第一次必然落空，按词格的静态顺序先出候选。线程算完会叫 + // 醒会话重查一次，那一次才在结果表里拿到顺序。 + const std::optional> order = + neural::RescoreWorker::instance().order_for(*model, context, sentences, scores, options); + if (!order || order->size() != paths.size()) + { + return false; // 还没算出来，或者模型自己弃权了：词格的顺序原样留着 + } + + std::vector reordered; + reordered.reserve(paths.size()); + for (std::size_t index : *order) + { + reordered.push_back(paths[index]); + } + paths = std::move(reordered); + return true; + }; +} + +} // namespace quanpin diff --git a/engine/quanpin/lattice_rerank.h b/engine/quanpin/lattice_rerank.h new file mode 100644 index 00000000..197fde4c --- /dev/null +++ b/engine/quanpin/lattice_rerank.h @@ -0,0 +1,27 @@ +#pragma once + +// Bridges the word lattice's n-best paths to the neural sentence model. +// +// word_lattice.* knows nothing about the model and engine/neural knows nothing about lattices; this +// is the one place that speaks both. Quanpin and shuangpin share it because their sentence pipelines +// are the same lattice over the same syllables. + +#include "../neural/neural_decoder.h" +#include "word_lattice.h" + +#include + +namespace quanpin +{ + +// A reranker for merge_lattice_candidates that reorders paths by the neural model, or an empty +// std::function when `model` is null (the model file is absent, or the feature is switched off). +// An empty reranker is not an error: the lattice's own order stands. +// +// `context` is the text already committed before the composition, used as the model's conditioning +// prefix. Empty is fine and simply scores each sentence on its own. Pass the whole history: it is +// trimmed here to the last `options.context_chars` characters, which is what actually gets scored. +LatticeReranker make_neural_reranker(const neural::SentenceModel *model, const std::string &context, + const neural::RerankOptions &options = {}); + +} // namespace quanpin diff --git a/engine/quanpin/quanpin_dictionary.cpp b/engine/quanpin/quanpin_dictionary.cpp index 15bae78c..8c737e33 100644 --- a/engine/quanpin/quanpin_dictionary.cpp +++ b/engine/quanpin/quanpin_dictionary.cpp @@ -239,6 +239,10 @@ QuanpinDictionary::QuanpinDictionary(std::string db_path, metasequoia::RuntimePa decoder_(paths_.resource(metasequoia::assets::pinyin_model), paths_.user(metasequoia::assets::pinyin_user_dictionary)), language_model_(&ngram::shared_language_model(paths_.resource(metasequoia::assets::language_model))), + neural_desktop_model_(neural::shared_sentence_model( + metasequoia::path_to_utf8(paths_.resource(metasequoia::assets::neural_model_desktop)))), + neural_keyboard_model_(neural::shared_sentence_model( + metasequoia::path_to_utf8(paths_.resource(metasequoia::assets::neural_model_keyboard)))), db_path_(db_path.empty() ? metasequoia::path_to_utf8(paths_.dictionary(metasequoia::assets::main_dictionary)) : std::move(db_path)) { @@ -522,33 +526,58 @@ std::vector QuanpinDictionary::query_series(const std::string &raw_inp // // 不管走哪条，送进去之前都要把 ü 换成它认的写法：它的音节表里只有 nue/lue， // nve 会被拆成 nv + e。见 quanpin::to_google_spelling。 - const std::string google_input = - raw_input.find('\'') != std::string::npos - ? quanpin::to_google_spelling(raw_input) - : remove_delimiters(quanpin::to_google_spelling(segmentation.empty() ? raw_input : segmentation)); - const std::string google_sentence = search_sentence_from_ime_engine(google_input); - if (!google_sentence.empty()) + // Google 整句（Fallback）：仅在开关打开时才解码、才出候选。 + if (sentence_association_.google) { - const auto duplicate = std::find_if(result.begin(), result.end(), - [&](const WordItem &item) { return item.word == google_sentence; }); - if (duplicate == result.end()) + const std::string google_input = + raw_input.find('\'') != std::string::npos + ? quanpin::to_google_spelling(raw_input) + : remove_delimiters(quanpin::to_google_spelling(segmentation.empty() ? raw_input : segmentation)); + const std::string google_sentence = search_sentence_from_ime_engine(google_input); + if (!google_sentence.empty()) { - // 整句 fallback 必须带上 canonical quanpin，否则以它结尾的造词无法落库： - // update_creating_word_progress 依赖 canonical_pinyin 才能拼出完整读音。 - // segmentation 为空时保持为空，交由既有逻辑判定为不可落库。 - const size_t insert_at = quanpin::generated_sentence_insert_position(result, segments); - result.insert(result.begin() + static_cast(insert_at), - WordItem(segmentation.empty() ? raw_input : segmentation, google_sentence, 1, - CandidateSource::Fallback, segmentation)); + const auto duplicate = std::find_if(result.begin(), result.end(), + [&](const WordItem &item) { return item.word == google_sentence; }); + if (duplicate == result.end()) + { + // 整句 fallback 必须带上 canonical quanpin，否则以它结尾的造词无法落库： + // update_creating_word_progress 依赖 canonical_pinyin 才能拼出完整读音。 + // segmentation 为空时保持为空，交由既有逻辑判定为不可落库。 + const size_t insert_at = quanpin::generated_sentence_insert_position(result, segments); + result.insert(result.begin() + static_cast(insert_at), + WordItem(segmentation.empty() ? raw_input : segmentation, google_sentence, 1, + CandidateSource::Fallback, segmentation)); + } } } + // 词格给 Trigram 候选和神经模型共用：只开神经时仍在内部解出 n-best 供它重排， + // 但不把词格自己的首选显示出来。神经模型不自己造句，只在这批路径中选一句。 + std::vector neural_rerankers; + if (sentence_association_.neural_keyboard && neural_keyboard_model_ != nullptr) + { + neural_rerankers.push_back({quanpin::make_neural_reranker(neural_keyboard_model_, rescoring_context_), + CandidateSource::NeuralKeyboard}); + } + if (sentence_association_.neural_desktop && neural_desktop_model_ != nullptr) + { + neural_rerankers.push_back({quanpin::make_neural_reranker(neural_desktop_model_, rescoring_context_), + CandidateSource::NeuralDesktop}); + } quanpin::WordLatticeOptions lattice_options; - lattice_options.nbest = 1; + // 神经模型和去重补位都需要完整候选池；否则只解最优的一条。 + const bool needs_alternatives = !neural_rerankers.empty() || (sentence_association_.word_lattice && + sentence_association_.show_next_on_duplicate); + lattice_options.nbest = needs_alternatives ? static_cast(neural::RerankOptions{}.max_paths) : 1; + lattice_options.include_lattice_best = sentence_association_.word_lattice; + lattice_options.show_next_on_duplicate = sentence_association_.show_next_on_duplicate; lattice_options.language_model = language_model_; - quanpin::merge_lattice_candidates( - result, segments, quanpin::make_lattice_db_lookup(db_, statement_cache_, lattice_options.span_limit), - segmentation.empty() ? raw_input : segmentation, lattice_options); + if (sentence_association_.word_lattice || !neural_rerankers.empty()) + { + quanpin::merge_lattice_candidates( + result, segments, quanpin::make_lattice_db_lookup(db_, statement_cache_, lattice_options.span_limit), + segmentation.empty() ? raw_input : segmentation, lattice_options, neural_rerankers); + } } if (result.size() < kSparsePinyinFallbackThreshold) @@ -770,6 +799,11 @@ std::vector QuanpinDictionary::append_ime_fallback(const std::string & { return result; } + // 这条兜底整句也来自 Google 解码器，随「Google 整句联想」开关一起关闭。 + if (!sentence_association_.google) + { + return result; + } const std::string normalized = remove_delimiters(segmentation.empty() ? raw_input : segmentation); const std::string sentence = search_sentence_from_ime_engine(normalized); @@ -1068,6 +1102,29 @@ void QuanpinDictionary::reset_cache() segmentation_cache_.clear(); } +void QuanpinDictionary::set_sentence_association(const SentenceAssociationOptions &options) +{ + if (sentence_association_ == options) + { + return; + } + sentence_association_ = options; + // 开关变了，之前缓存的候选列表按旧开关算出，必须清掉，否则打开/关闭后要到缓存过期 + // 才见效。整句候选都进这几个缓存，见 query_series / query_single_path。 + reset_cache(); +} + +void QuanpinDictionary::set_rescoring_context(const std::string &context) +{ + if (rescoring_context_ == context) + { + return; + } + rescoring_context_ = context; + // 上文变了，缓存里的顺序是按旧上文重排出来的，同 set_sentence_association。 + reset_cache(); +} + void QuanpinDictionary::reset_cache_if_database_changed() { if (db_ == nullptr) diff --git a/engine/quanpin/quanpin_dictionary.h b/engine/quanpin/quanpin_dictionary.h index f9ada8ef..203f371e 100644 --- a/engine/quanpin/quanpin_dictionary.h +++ b/engine/quanpin/quanpin_dictionary.h @@ -7,6 +7,9 @@ #include "../common/cache.h" #include "../core/key_event.h" #include "../core/word_item.h" +#include "../core/sentence_association_options.h" +#include "../neural/neural_decoder.h" +#include "lattice_rerank.h" #include "quanpin_query.h" #include "engine/ngram/language_model.h" #include "../core/fuzzy_pinyin_options.h" @@ -70,6 +73,14 @@ class QuanpinDictionary std::string search_sentence_from_ime_engine(const std::string &user_pinyin); + // 设置整句候选来源与去重补位选项。值发生变化时清一次缓存，避免旧候选被继续复用。 + void set_sentence_association(const SentenceAssociationOptions &options); + + // 光标前已上屏的文本，作为神经重排给模型看的上下文。宿主能拿到应用里的前文就给它， + // 拿不到就不设——空上下文照样能排，只是模型只看句子本身、看不到上文。 + // 变化时清一次缓存：同一串拼音在不同上文下的排序可以不同。 + void set_rescoring_context(const std::string &context); + void reset_state(); void reset_cache(); @@ -144,6 +155,14 @@ class QuanpinDictionary // 词格整句的打分模型（kenlm 三元），按资源路径进程内共享、只读。缺模型时 // valid() 为假，word_lattice 退回旧的启发式打分：整句候选只是变差，不会消失。 const ngram::LanguageModel *language_model_ = nullptr; + // 神经整句模型（chinese-ime-lm），按资源路径进程内共享、只读。两档可同时参与打分； + // 缺文件时对应的整句候选不出，不影响其它候选。 + const neural::SentenceModel *neural_desktop_model_ = nullptr; + const neural::SentenceModel *neural_keyboard_model_ = nullptr; + // 整句候选来源与去重补位选项，默认全关；由 set_sentence_association 随请求更新。 + SentenceAssociationOptions sentence_association_; + // 神经重排的上下文（光标前已上屏的文本）。宿主没给就是空串。 + std::string rescoring_context_; std::unordered_map statement_cache_; std::string db_path_; diff --git a/engine/quanpin/word_lattice.cpp b/engine/quanpin/word_lattice.cpp index 152f5eab..130b919c 100644 --- a/engine/quanpin/word_lattice.cpp +++ b/engine/quanpin/word_lattice.cpp @@ -3,7 +3,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -314,30 +316,160 @@ std::vector decode_word_lattice(const Segments &syllables, const Wo void merge_lattice_candidates(std::vector &candidates, const Segments &syllables, const WordLatticeLookup &lookup, const std::string &typed_pinyin, - const WordLatticeOptions &options) + const WordLatticeOptions &options, const std::vector &rerankers) { if (!lookup || syllables.size() < 2) return; if (!has_only_complete_pinyin_segments(syllables)) return; - const auto paths = decode_word_lattice(syllables, lookup, options); + auto paths = decode_word_lattice(syllables, lookup, options); if (paths.empty()) return; + // n-best 是解码的中间产物，不是要给用户看的东西：整个 n-best 铺出来，候选区会被一批 + // 只差一个字的整句占满。这里只留各来源的首选，n-best 的其余部分仅供重排挑选。 + // 每个重排器都从相同的静态顺序开始，不能让前一个模型的结论变成后一个的输入。保留完整顺序， + // 这样首选与已有候选重复时可以从该来源自己的次选继续补位。 + std::optional> desktop_paths; + std::optional> keyboard_paths; + std::vector, CandidateSource>> other_ranked_sources; + bool keyboard_enabled = false; + for (const SourcedLatticeReranker &entry : rerankers) + { + if (entry.source == CandidateSource::NeuralKeyboard) + keyboard_enabled = true; + + std::vector reranked_paths = paths; + if (!entry.rerank || !entry.rerank(reranked_paths)) + continue; + + if (entry.source == CandidateSource::NeuralDesktop) + desktop_paths = std::move(reranked_paths); + else if (entry.source == CandidateSource::NeuralKeyboard) + keyboard_paths = std::move(reranked_paths); + else + other_ranked_sources.emplace_back(std::move(reranked_paths), entry.source); + } + + const auto source_prefers = [](const std::optional> &ranked_paths, + const std::string &sentence) { + return ranked_paths && !ranked_paths->empty() && ranked_paths->front().sentence == sentence; + }; + const auto unigram_position = std::find_if(candidates.begin(), candidates.end(), [](const WordItem &item) { + return item.source == CandidateSource::Fallback; + }); + const bool unigram_consensus = + unigram_position != candidates.end() && + ((options.include_lattice_best && paths.front().sentence == unigram_position->word) || + source_prefers(keyboard_paths, unigram_position->word) || + source_prefers(desktop_paths, unigram_position->word)); + std::unordered_set already; for (const auto &item : candidates) already.insert(item.word); - std::vector extra; - for (const auto &path : paths) - { - if (!already.insert(path.sentence).second) - continue; + const auto first_distinct = [&](const std::vector &ranked_paths) -> const LatticePath * { + for (const LatticePath &path : ranked_paths) + { + if (already.find(path.sentence) != already.end()) + { + if (options.show_next_on_duplicate) + continue; + break; + } + return &path; + } + return nullptr; + }; + const auto select_distinct = [&](const std::vector &ranked_paths, + CandidateSource source) -> std::optional { + const LatticePath *path = first_distinct(ranked_paths); + if (path == nullptr) + return std::nullopt; + already.insert(path->sentence); // Often negative (log-space). Ranking is insert order, not weight. - const auto weight = static_cast(path.log_prob * 1000.0); - extra.emplace_back(typed_pinyin, path.sentence, weight, CandidateSource::Generated, path.key); + const auto weight = static_cast(path->log_prob * 1000.0); + return WordItem(typed_pinyin, path->sentence, weight, source, path->key); + }; + + // Existing rows include Unigram, so Trigram avoids it before the neural models choose their rows. + std::optional trigram_item; + std::optional keyboard_item; + std::optional desktop_item; + if (options.include_lattice_best) + trigram_item = select_distinct(paths, CandidateSource::Generated); + const bool trigram_consensus = trigram_item && trigram_item->word == paths.front().sentence && + (source_prefers(keyboard_paths, paths.front().sentence) || + source_prefers(desktop_paths, paths.front().sentence)); + bool keyboard_before_desktop = false; + if (keyboard_paths && desktop_paths) + { + // Compare each model's best row after excluding Unigram and Trigram, but before the neural + // models exclude each other. Keyboard owns a shared best; otherwise desktop keeps its own + // best and is displayed first. + const LatticePath *keyboard_best = first_distinct(*keyboard_paths); + const LatticePath *desktop_best = first_distinct(*desktop_paths); + keyboard_before_desktop = + keyboard_best != nullptr && desktop_best != nullptr && keyboard_best->sentence == desktop_best->sentence; + if (keyboard_before_desktop) + { + keyboard_item = select_distinct(*keyboard_paths, CandidateSource::NeuralKeyboard); + desktop_item = select_distinct(*desktop_paths, CandidateSource::NeuralDesktop); + } + else + { + desktop_item = select_distinct(*desktop_paths, CandidateSource::NeuralDesktop); + keyboard_item = select_distinct(*keyboard_paths, CandidateSource::NeuralKeyboard); + } } + else if (keyboard_paths) + { + keyboard_item = select_distinct(*keyboard_paths, CandidateSource::NeuralKeyboard); + } + // When both neural modes are enabled, desktop waits for keyboard to finish. Otherwise desktop + // could briefly claim keyboard's best distinct row and then visibly change on the next query. + else if (desktop_paths && !keyboard_enabled) + desktop_item = select_distinct(*desktop_paths, CandidateSource::NeuralDesktop); + + std::vector other_items; + for (const auto &[ranked_paths, source] : other_ranked_sources) + { + if (auto item = select_distinct(ranked_paths, source)) + other_items.push_back(std::move(*item)); + } + + std::optional promoted_unigram; + if (unigram_consensus) + { + promoted_unigram = std::move(*unigram_position); + candidates.erase(unigram_position); + } + + // Consensus from another source promotes Unigram first and Trigram ahead of neural rows. A + // shared neural best belongs to keyboard; when the neural best rows differ, desktop goes first. + std::vector extra; + if (promoted_unigram) + extra.push_back(std::move(*promoted_unigram)); + if (trigram_consensus) + extra.push_back(std::move(*trigram_item)); + if (keyboard_before_desktop) + { + if (keyboard_item) + extra.push_back(std::move(*keyboard_item)); + if (desktop_item) + extra.push_back(std::move(*desktop_item)); + } + else + { + if (desktop_item) + extra.push_back(std::move(*desktop_item)); + if (keyboard_item) + extra.push_back(std::move(*keyboard_item)); + } + extra.insert(extra.end(), std::make_move_iterator(other_items.begin()), std::make_move_iterator(other_items.end())); + if (trigram_item && !trigram_consensus) + extra.push_back(std::move(*trigram_item)); if (extra.empty()) return; @@ -345,4 +477,15 @@ void merge_lattice_candidates(std::vector &candidates, const Segments candidates.insert(candidates.begin() + static_cast(insert_at), extra.begin(), extra.end()); } +void merge_lattice_candidates(std::vector &candidates, const Segments &syllables, + const WordLatticeLookup &lookup, const std::string &typed_pinyin, + const WordLatticeOptions &options, const LatticeReranker &rerank, + CandidateSource rerank_source) +{ + std::vector rerankers; + if (rerank) + rerankers.push_back({rerank, rerank_source}); + merge_lattice_candidates(candidates, syllables, lookup, typed_pinyin, options, rerankers); +} + } // namespace quanpin diff --git a/engine/quanpin/word_lattice.h b/engine/quanpin/word_lattice.h index b1066024..2f5ff234 100644 --- a/engine/quanpin/word_lattice.h +++ b/engine/quanpin/word_lattice.h @@ -48,6 +48,11 @@ namespace quanpin // // Lattice WordItem.weight is log_prob * 1000 and is often negative. List // order is the ranking; do not sort these rows by weight. +// +// merge_lattice_candidates inserts at most three rows, not the whole n-best: each +// reranker's pick and, when enabled, the lattice's own pick. The +// rest of the n-best exists so the reranker has something to choose from -- several +// near-identical sentences are a decoder detail, not a candidate list. struct LatticeLexeme { @@ -68,6 +73,12 @@ struct WordLatticeOptions { int beam = 32; int nbest = 5; + // Whether the lattice's own best path should become a visible candidate. Neural-only mode sets + // this to false: the lattice still supplies n-best paths to the reranker, but only its pick is shown. + bool include_lattice_best = true; + // Each visible source contributes at most one row. When its first choice duplicates an earlier + // candidate, continue down that source's own ranking until a distinct sentence is found. + bool show_next_on_duplicate = false; // Cap on lexemes per span (injected lookups). The DB lookup applies the // same cap itself, inside query_exact_segmentations_keyed_flat. int span_limit = 32; @@ -103,6 +114,23 @@ struct WordLatticeOptions using WordLatticeLookup = std::function(const Segments &span)>; +// Optional second opinion on the decoded n-best, applied before the paths become candidates. +// engine/neural supplies one that reorders them by a character-level Transformer; anything else +// that can rank finished sentences would plug in the same way. Reordering only -- a reranker must +// not add, drop, or edit paths, because their WordItem fields are derived from them afterwards. +// +// Returns whether it actually ranked the paths. False means it declined -- the model abstained, or +// (in the async case) its answer is not back yet -- and the order on return is still the lattice's. +// The caller needs this to attribute the surviving row: a reranker that ran and agreed leaves the +// same order as one that never ran, and those two are not the same thing to report. +using LatticeReranker = std::function &paths)>; + +struct SourcedLatticeReranker +{ + LatticeReranker rerank; + CandidateSource source = CandidateSource::Generated; +}; + // Index where a whole-sentence candidate belongs: just past the leading run of // Database/UserDatabase hits whose key equals the typed syllables exactly. // Prefix-range and fuzzy rows are not exact hits even when they happen to have @@ -114,8 +142,21 @@ size_t generated_sentence_insert_position(const std::vector &candidate std::vector decode_word_lattice(const Segments &syllables, const WordLatticeLookup &lookup, const WordLatticeOptions &options = {}); +// `rerank_source` is the source stamped on the reranker's pick, so the candidate list can say which +// component chose that row. When options.include_lattice_best is true, the lattice's own pick stays +// CandidateSource::Generated; when false, no row is inserted until the reranker has produced a result. void merge_lattice_candidates(std::vector &candidates, const Segments &syllables, const WordLatticeLookup &lookup, const std::string &typed_pinyin, - const WordLatticeOptions &options = {}); + const WordLatticeOptions &options = {}, const LatticeReranker &rerank = {}, + CandidateSource rerank_source = CandidateSource::Generated); + +// Runs every enabled reranker independently over the same original n-best. Each completed reranker +// may contribute one visible row. Existing candidates (including Unigram) and Trigram reserve their +// rows first. If both neural models have the same best remaining row, NeuralKeyboard owns it and is +// displayed first; otherwise NeuralDesktop keeps its best remaining row and is displayed first. +// A Unigram or Trigram best shared by another enabled source is promoted ahead of the neural rows. +void merge_lattice_candidates(std::vector &candidates, const Segments &syllables, + const WordLatticeLookup &lookup, const std::string &typed_pinyin, + const WordLatticeOptions &options, const std::vector &rerankers); } // namespace quanpin diff --git a/engine/shuangpin/engine.cpp b/engine/shuangpin/engine.cpp index 3f08e5db..b6712bc9 100644 --- a/engine/shuangpin/engine.cpp +++ b/engine/shuangpin/engine.cpp @@ -117,6 +117,11 @@ std::vector ShuangpinEngine::query(const QueryRequest &request) return {}; } + // 整句联想开关随请求下发到词典层，四条整句来源据此各自门控（见 generateSeries）。 + dictionary_.set_sentence_association(request.sentence_association); + // 前文同理随请求下发；换了前文，按旧前文重排出来的缓存顺序作废。 + dictionary_.set_rescoring_context(request.rescoring_context); + const std::string &raw_input = request.raw_input; const std::string &raw_input_with_cases = request.raw_input_with_cases.empty() ? request.raw_input : request.raw_input_with_cases; diff --git a/engine/shuangpin/shuangpin_dictionary.cpp b/engine/shuangpin/shuangpin_dictionary.cpp index 5ff8db48..7cf5c1bf 100644 --- a/engine/shuangpin/shuangpin_dictionary.cpp +++ b/engine/shuangpin/shuangpin_dictionary.cpp @@ -53,6 +53,10 @@ ShuangpinDictionary::ShuangpinDictionary(const ShuangpinProfile &profile, metase : profile_(profile), paths_(std::move(paths)), decoder_(paths_.resource(metasequoia::assets::pinyin_model), paths_.user(metasequoia::assets::pinyin_user_dictionary)), language_model_(&ngram::shared_language_model(paths_.resource(metasequoia::assets::language_model))), + neural_desktop_model_(neural::shared_sentence_model( + metasequoia::path_to_utf8(paths_.resource(metasequoia::assets::neural_model_desktop)))), + neural_keyboard_model_(neural::shared_sentence_model( + metasequoia::path_to_utf8(paths_.resource(metasequoia::assets::neural_model_keyboard)))), helpcodes_(HelpcodeUtils::load_helpcode_keymap(paths_.resources, HelpcodeUtils::selected_helpcode_schema())), _kb_input_sequence(100), _cached_buffer(128), _cached_buffer_sgl(128), _cached_buffer_sgl_reversed(128), _cached_buffer_dbl(128), _cached_buffer_series(128) @@ -160,7 +164,7 @@ vector ShuangpinDictionary::generateSeries( // } else { /* 可能数据库查询的结果是空，这时就需要联想，这个只适合在此处联想 */ - if (candidate_list.size() == 0) + if (candidate_list.size() == 0 && sentence_association_.google) { string quanpin_str = ShuangpinUtil::convert_seg_shuangpin_to_seg_complete_pinyin(pinyin_segmentation, profile_); @@ -203,7 +207,8 @@ vector ShuangpinDictionary::generateSeries( // // 命中词库的候选之后；merge_lattice_candidates 算出的位置就是这条 Fallback // 所在的下标，词格随后会落在它上面。 const auto quanpin_syllables = quanpin::split_segments(quanpin_segmentation); - if (quanpin_syllables.size() >= 2 && quanpin_segmentation.find('\'') != std::string::npos) + if (sentence_association_.google && quanpin_syllables.size() >= 2 && + quanpin_segmentation.find('\'') != std::string::npos) { // 同上：解码器认 nue/lue，词库与下面的 canonical 读音认 nve/lve。 const std::string google_sentence = @@ -219,13 +224,33 @@ vector ShuangpinDictionary::generateSeries( // WordItem(_pinyin_sequence, google_sentence, 1, CandidateSource::Fallback, quanpin_segmentation)); } } + // 词格给 Trigram 候选和神经模型共用；只开神经时仍在内部解出 n-best，但不显示词格首选。 + // 与全拼同构，见 QuanpinDictionary::query_series。 + std::vector neural_rerankers; + if (sentence_association_.neural_keyboard && neural_keyboard_model_ != nullptr) + { + neural_rerankers.push_back({quanpin::make_neural_reranker(neural_keyboard_model_, rescoring_context_), + CandidateSource::NeuralKeyboard}); + } + if (sentence_association_.neural_desktop && neural_desktop_model_ != nullptr) + { + neural_rerankers.push_back({quanpin::make_neural_reranker(neural_desktop_model_, rescoring_context_), + CandidateSource::NeuralDesktop}); + } quanpin::WordLatticeOptions lattice_options; - lattice_options.nbest = 1; + const bool needs_alternatives = !neural_rerankers.empty() || (sentence_association_.word_lattice && + sentence_association_.show_next_on_duplicate); + lattice_options.nbest = needs_alternatives ? static_cast(neural::RerankOptions{}.max_paths) : 1; + lattice_options.include_lattice_best = sentence_association_.word_lattice; + lattice_options.show_next_on_duplicate = sentence_association_.show_next_on_duplicate; lattice_options.language_model = language_model_; - quanpin::merge_lattice_candidates( - candidate_list, quanpin_syllables, - quanpin::make_lattice_db_lookup(quanpin_db_, quanpin_statement_cache_, lattice_options.span_limit), - pinyin_sequence, lattice_options); + if (sentence_association_.word_lattice || !neural_rerankers.empty()) + { + quanpin::merge_lattice_candidates( + candidate_list, quanpin_syllables, + quanpin::make_lattice_db_lookup(quanpin_db_, quanpin_statement_cache_, lattice_options.span_limit), + pinyin_sequence, lattice_options, neural_rerankers); + } /* 缓存起来 */ _cached_buffer_series.insert(effective_cache_key, candidate_list); @@ -1115,6 +1140,27 @@ void ShuangpinDictionary::reset_cache() _cached_buffer_series.clear(); } +void ShuangpinDictionary::set_sentence_association(const SentenceAssociationOptions &options) +{ + if (sentence_association_ == options) + { + return; + } + sentence_association_ = options; + // 开关变了就清缓存，否则打开/关闭要等缓存过期才见效。见 generateSeries 的整句块。 + reset_cache(); +} + +void ShuangpinDictionary::set_rescoring_context(const std::string &context) +{ + if (rescoring_context_ == context) + { + return; + } + rescoring_context_ = context; + reset_cache(); +} + void ShuangpinDictionary::reset_cache_if_database_changed() { if (quanpin_db_ == nullptr) diff --git a/engine/shuangpin/shuangpin_dictionary.h b/engine/shuangpin/shuangpin_dictionary.h index c3dbc681..39bb7d8f 100644 --- a/engine/shuangpin/shuangpin_dictionary.h +++ b/engine/shuangpin/shuangpin_dictionary.h @@ -8,6 +8,9 @@ #include "../common/cache.h" #include "../core/key_event.h" #include "../core/word_item.h" +#include "../core/sentence_association_options.h" +#include "../neural/neural_decoder.h" +#include "../quanpin/lattice_rerank.h" #include "../quanpin/quanpin_query.h" #include "engine/ngram/language_model.h" #include "shuangpin_profile.h" @@ -66,6 +69,12 @@ class ShuangpinDictionary std::string search_sentence_from_ime_engine(const std::string &user_pinyin); + // 设置整句候选来源与去重补位选项。值变化时清缓存，见 QuanpinDictionary::set_sentence_association。 + void set_sentence_association(const SentenceAssociationOptions &options); + + // 神经重排的上下文，见 QuanpinDictionary::set_rescoring_context。 + void set_rescoring_context(const std::string &context); + explicit ShuangpinDictionary(const ShuangpinProfile &profile = GetXiaoheShuangpinProfile(), metasequoia::RuntimePaths paths = metasequoia::RuntimePaths::legacy()); ~ShuangpinDictionary(); @@ -84,6 +93,14 @@ class ShuangpinDictionary metasequoia::PinyinDecoder decoder_; // 与全拼共用的词格打分模型，见 QuanpinDictionary::language_model_。 const ngram::LanguageModel *language_model_ = nullptr; + // 神经整句模型，见 QuanpinDictionary::neural_desktop_model_。与全拼共用同一份 + // 进程内缓存（按资源路径），双拼这侧只是各自持一份指针。 + const neural::SentenceModel *neural_desktop_model_ = nullptr; + const neural::SentenceModel *neural_keyboard_model_ = nullptr; + // 整句候选来源与去重补位选项，默认全关；由 set_sentence_association 随请求更新。 + SentenceAssociationOptions sentence_association_; + // 神经重排的上下文（光标前已上屏的文本）。宿主没给就是空串。 + std::string rescoring_context_; HelpcodeUtils::SharedKeymap helpcodes_; std::unordered_map quanpin_statement_cache_; void reset_cache_if_database_changed(); diff --git a/engine/tests/CMakeLists.txt b/engine/tests/CMakeLists.txt index edf0878c..6a492556 100644 --- a/engine/tests/CMakeLists.txt +++ b/engine/tests/CMakeLists.txt @@ -71,6 +71,9 @@ set( "../local_modes/kaomoji_query.cpp" "../local_modes/quick_phrase_query.cpp" "../local_modes/unicode_query.cpp" + "../neural/sentence_model.cpp" + "../neural/neural_decoder.cpp" + "../neural/rescore_worker.cpp" "../providers/pinyin_candidate_provider.cpp" "../providers/wubi_candidate_provider.cpp" "../providers/japanese_candidate_provider.cpp" @@ -83,6 +86,7 @@ set( "../quanpin/quanpin_dictionary.cpp" "../quanpin/quanpin_query.cpp" "../quanpin/quanpin_utils.cpp" + "../quanpin/lattice_rerank.cpp" "../quanpin/word_lattice.cpp" "../shuangpin/engine.cpp" "../common/helpcode_utils.cpp" diff --git a/engine/tests/src/test_engine_smoke.cpp b/engine/tests/src/test_engine_smoke.cpp index d3ee6a6f..aa7943a1 100644 --- a/engine/tests/src/test_engine_smoke.cpp +++ b/engine/tests/src/test_engine_smoke.cpp @@ -2,6 +2,8 @@ #include "../../core/data_path.h" #include "../../english/english_dictionary.h" #include "../../japanese/romaji_converter.h" +#include "../../neural/neural_decoder.h" +#include "../../quanpin/word_lattice.h" #include "../../user_dictionary/user_dictionary_journal.h" #include "test_directory_cleanup.h" @@ -14,6 +16,8 @@ #include #include #include +#include +#include namespace { @@ -84,6 +88,152 @@ int run_test() throw std::runtime_error("Kana conversion changed during the platform port."); } + { + // Neural-only mode still asks the lattice for twelve internal alternatives, but it must not + // expose the lattice's own pick while the asynchronous neural result is pending. + std::unordered_map> table; + table["ni"] = {{"ni", "你", 4000}, {"ni", "拟", 3000}, {"ni", "泥", 2000}, {"ni", "逆", 1000}}; + table["hao"] = {{"hao", "好", 4000}, {"hao", "号", 3000}, {"hao", "浩", 2000}, {"hao", "耗", 1000}}; + const auto lookup = [&](const quanpin::Segments &span) { + std::string key; + for (const std::string &syllable : span) + { + if (!key.empty()) + key.push_back('\''); + key += syllable; + } + const auto found = table.find(key); + return found == table.end() ? std::vector{} : found->second; + }; + + quanpin::WordLatticeOptions options; + options.nbest = static_cast(neural::RerankOptions{}.max_paths); + options.include_lattice_best = false; + options.show_next_on_duplicate = true; + const std::vector neural_paths = + quanpin::decode_word_lattice({"ni", "hao"}, lookup, options); + std::size_t desktop_alternatives = 0; + std::size_t keyboard_alternatives = 0; + std::vector rerankers; + rerankers.push_back({[&](std::vector &paths) { + desktop_alternatives = paths.size(); + return true; + }, + CandidateSource::NeuralDesktop}); + rerankers.push_back({[&](std::vector &paths) { + keyboard_alternatives = paths.size(); + return true; + }, + CandidateSource::NeuralKeyboard}); + std::vector selected; + quanpin::merge_lattice_candidates(selected, {"ni", "hao"}, lookup, "ni'hao", options, rerankers); + if (desktop_alternatives != 12 || keyboard_alternatives != 12 || selected.size() != 2 || + selected[0].source != CandidateSource::NeuralKeyboard || selected[0].word != neural_paths[0].sentence || + selected[1].source != CandidateSource::NeuralDesktop || selected[1].word != neural_paths[1].sentence) + { + throw std::runtime_error("Keyboard did not keep the shared neural best before desktop's next choice."); + } + + std::vector different_neural_best; + const std::vector different_best_rerankers = { + {[](std::vector &paths) { + std::swap(paths[0], paths[1]); + return true; + }, + CandidateSource::NeuralDesktop}, + {[](std::vector &) { return true; }, CandidateSource::NeuralKeyboard}}; + quanpin::merge_lattice_candidates(different_neural_best, {"ni", "hao"}, lookup, "ni'hao", options, + different_best_rerankers); + if (different_neural_best.size() != 2 || different_neural_best[0].source != CandidateSource::NeuralDesktop || + different_neural_best[0].word != neural_paths[1].sentence || + different_neural_best[1].source != CandidateSource::NeuralKeyboard || + different_neural_best[1].word != neural_paths[0].sentence) + { + throw std::runtime_error("Desktop did not precede keyboard when their neural best rows differed."); + } + + std::vector pending; + const std::vector pending_rerankers = { + {[](std::vector &) { return false; }, CandidateSource::NeuralDesktop}, + {[](std::vector &) { return false; }, CandidateSource::NeuralKeyboard}}; + quanpin::merge_lattice_candidates(pending, {"ni", "hao"}, lookup, "ni'hao", options, pending_rerankers); + if (!pending.empty()) + { + throw std::runtime_error("Neural-only sentence selection exposed a lattice row before scoring finished."); + } + + std::vector keyboard_ready; + const std::vector keyboard_ready_rerankers = { + {[](std::vector &) { return false; }, CandidateSource::NeuralDesktop}, + {[](std::vector &) { return true; }, CandidateSource::NeuralKeyboard}}; + quanpin::merge_lattice_candidates(keyboard_ready, {"ni", "hao"}, lookup, "ni'hao", options, + keyboard_ready_rerankers); + if (keyboard_ready.size() != 1 || keyboard_ready[0].source != CandidateSource::NeuralKeyboard || + keyboard_ready[0].word != neural_paths[0].sentence) + { + throw std::runtime_error("Keyboard neural output waited for the desktop model."); + } + + quanpin::WordLatticeOptions trigram_consensus_options = options; + trigram_consensus_options.include_lattice_best = true; + std::vector trigram_consensus; + const std::vector trigram_consensus_rerankers = { + {[](std::vector &) { return true; }, CandidateSource::NeuralKeyboard}}; + quanpin::merge_lattice_candidates(trigram_consensus, {"ni", "hao"}, lookup, "ni'hao", trigram_consensus_options, + trigram_consensus_rerankers); + if (trigram_consensus.size() != 2 || trigram_consensus[0].source != CandidateSource::Generated || + trigram_consensus[0].word != neural_paths[0].sentence || + trigram_consensus[1].source != CandidateSource::NeuralKeyboard || + trigram_consensus[1].word != neural_paths[1].sentence) + { + throw std::runtime_error("Trigram consensus was not promoted ahead of the neural rows."); + } + + // Simulate an earlier Unigram row equal to every later source's first choice. Ownership must + // remain Unigram, Trigram, keyboard, desktop regardless of reranker input or completion order. + const std::vector static_paths = + quanpin::decode_word_lattice({"ni", "hao"}, lookup, options); + std::vector deduplicated; + deduplicated.emplace_back("ni'hao", static_paths.front().sentence, 1, CandidateSource::Fallback, "ni'hao"); + options.include_lattice_best = true; + options.show_next_on_duplicate = true; + const std::vector duplicate_rerankers = { + {[](std::vector &) { return true; }, CandidateSource::NeuralDesktop}, + {[](std::vector &paths) { + std::swap(paths[0], paths[1]); + return true; + }, + CandidateSource::NeuralKeyboard}}; + quanpin::merge_lattice_candidates(deduplicated, {"ni", "hao"}, lookup, "ni'hao", options, duplicate_rerankers); + if (deduplicated.size() != 4 || deduplicated[0].source != CandidateSource::Fallback || + deduplicated[0].word != static_paths[0].sentence || + deduplicated[1].source != CandidateSource::NeuralKeyboard || + deduplicated[1].word != static_paths[2].sentence || + deduplicated[2].source != CandidateSource::NeuralDesktop || + deduplicated[2].word != static_paths[3].sentence || deduplicated[3].source != CandidateSource::Generated || + deduplicated[3].word != static_paths[1].sentence) + { + throw std::runtime_error("Unigram consensus was not promoted before distinct fallback rows."); + } + + // Desktop must wait while an enabled keyboard model is pending. Trigram can already appear + // because its ownership was fixed before either neural model. + std::vector partially_ready; + partially_ready.emplace_back("ni'hao", static_paths.front().sentence, 1, CandidateSource::Fallback, "ni'hao"); + const std::vector partially_ready_rerankers = { + {[](std::vector &) { return true; }, CandidateSource::NeuralDesktop}, + {[](std::vector &) { return false; }, CandidateSource::NeuralKeyboard}}; + quanpin::merge_lattice_candidates(partially_ready, {"ni", "hao"}, lookup, "ni'hao", options, + partially_ready_rerankers); + if (partially_ready.size() != 2 || partially_ready[0].source != CandidateSource::Fallback || + partially_ready[0].word != static_paths[0].sentence || + partially_ready[1].source != CandidateSource::Generated || + partially_ready[1].word != static_paths[1].sentence) + { + throw std::runtime_error("Desktop neural output appeared before keyboard ownership was fixed."); + } + } + const auto unique_suffix = std::to_string(std::chrono::high_resolution_clock::now().time_since_epoch().count()); const std::filesystem::path data_directory = std::filesystem::temp_directory_path() / std::filesystem::u8path("metasequoia-engine-词库-" + unique_suffix); diff --git a/engine/tests/src/test_input_session.cpp b/engine/tests/src/test_input_session.cpp index 43357f35..ae513c28 100644 --- a/engine/tests/src/test_input_session.cpp +++ b/engine/tests/src/test_input_session.cpp @@ -638,6 +638,11 @@ int run_test() database.execute("CREATE TABLE tbl_3_n(key TEXT,jp TEXT,value TEXT,weight INTEGER)"); { metasequoia::InputSession sentence(SchemeType::Quanpin); + // 整句联想默认全关，这一段考的就是整句候选，得先把造句那两条打开。 + SentenceAssociationOptions association; + association.word_lattice = true; + association.google = true; + sentence.set_sentence_association(association); type(sentence, "na'yi'tiao"); const auto guessed = candidate_index(sentence, "那一条"); const auto guessed_source = sentence.candidates()[guessed].source; @@ -660,6 +665,10 @@ int run_test() { // 学习整句是造词，不是调频：关掉候选学习的会话一条也不该落库。 metasequoia::InputSession unlearned(SchemeType::Quanpin, 0, true, true, false); + SentenceAssociationOptions association; + association.word_lattice = true; + association.google = true; + unlearned.set_sentence_association(association); type(unlearned, "na'yi'na"); const auto guessed = candidate_index(unlearned, "那一那"); (void)unlearned.select_candidate(guessed); diff --git a/engine/tests/src/test_neural_sentence.cpp b/engine/tests/src/test_neural_sentence.cpp new file mode 100644 index 00000000..59ef7aee --- /dev/null +++ b/engine/tests/src/test_neural_sentence.cpp @@ -0,0 +1,358 @@ +// 神经整句重排的正确性与延迟测试。 +// +// 模型现在只重排词格的 n-best，不再自己造句（见 engine/neural/neural_decoder.h）。快是靠两件 +// 事换来的：前文的 K/V 缓存下来不重算，以及所有候选拼成一批一起过每一层。这两件事都可能算错 +// 而不报错——位置编码串行、padding 行漏掉、缓存没失效，都只会让分数悄悄变样。所以主检查项是 +// 对拍：score_sentences 的结果必须等于逐字调用 next_log_probabilities 累加出来的参考值，后者 +// 是与 Rust 参考实现对齐过的那条老路径。 +// +// 第三件事是打分挪到了后台线程（engine/neural/rescore_worker.h），所以走词典的那段集成检查只看得 +// 到静态顺序——重排要等下一次重查。后台那条路单独测：请求一次必落空，回调到了再请求必命中。 +// +// 缺 sentence-model*.safetensors 时与其它集成测试一样优雅跳过（返回 0）。用 +// METASEQUOIA_IME_DATA_DIR 指到备好资源的目录来实际跑通。 + +#include "core/sentence_association_options.h" +#include "core/word_item.h" +#include "neural/neural_decoder.h" +#include "neural/rescore_worker.h" +#include "neural/sentence_model.h" +#include "quanpin/quanpin_dictionary.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace +{ +int failures = 0; + +void expect(bool condition, const std::string &message) +{ + if (!condition) + { + std::fprintf(stderr, "FAIL: %s\n", message.c_str()); + ++failures; + } +} + +void expect_near(double actual, double expected, double tolerance, const std::string &message) +{ + if (!(std::fabs(actual - expected) <= tolerance)) + { + std::fprintf(stderr, "FAIL: %s (got %.6f, want %.6f)\n", message.c_str(), actual, expected); + ++failures; + } +} + +std::vector characters(const std::string &text) +{ + std::vector out; + for (std::size_t i = 0; i < text.size();) + { + std::size_t width = 1; + const unsigned char c = static_cast(text[i]); + if ((c & 0xF8) == 0xF0) + width = 4; + else if ((c & 0xF0) == 0xE0) + width = 3; + else if ((c & 0xE0) == 0xC0) + width = 2; + width = std::min(width, text.size() - i); + out.push_back(text.substr(i, width)); + i += width; + } + return out; +} + +char32_t code_point(const std::string &character) +{ + const unsigned char c = static_cast(character[0]); + if (c < 0x80) + return c; + char32_t code; + std::size_t extra; + if ((c & 0xE0) == 0xC0) + { + code = c & 0x1F; + extra = 1; + } + else if ((c & 0xF0) == 0xE0) + { + code = c & 0x0F; + extra = 2; + } + else + { + code = c & 0x07; + extra = 3; + } + for (std::size_t j = 1; j <= extra && j < character.size(); ++j) + code = (code << 6) | (static_cast(character[j]) & 0x3F); + return code; +} + +// 参考实现：一个字一个字地问「下一个字的分布」，取出实际那个字的 log 概率累加。每一步都重跑 +// 整个前缀，正是重构掉的那条慢路径——留在这里当标尺。 +double reference_score(const neural::SentenceModel &model, const std::string &context, const std::string &text) +{ + double total = 0.0; + std::string prefix = context; + for (const std::string &character : characters(text)) + { + const std::vector distribution = model.next_log_probabilities(prefix); + // 词表外的字按 记，与 SentenceModel::encode 一致。 + const std::optional id = model.token(code_point(character)); + total += static_cast(distribution[id ? *id : 1u]); + prefix += character; + } + return total; +} +} // namespace + +int main() +{ + const metasequoia::RuntimePaths paths = metasequoia::RuntimePaths::legacy(); + const neural::SentenceModel *model = neural::shared_sentence_model( + metasequoia::path_to_utf8(paths.resource(metasequoia::assets::neural_model_keyboard))); + const neural::SentenceModel *desktop_model = neural::shared_sentence_model( + metasequoia::path_to_utf8(paths.resource(metasequoia::assets::neural_model_desktop))); + if (model == nullptr) + { + std::printf("Skipped: neural model (sentence-model.safetensors) not present in data directory.\n"); + return 0; + } + + // ---- 对拍：批量 + 缓存的结果必须等于逐字参考值 ---- + const std::vector>> cases = { + {"", {"你好世界", "拟好世界"}}, + {"我今天想去", {"上海", "伤害", "吃饭"}}, + // 长短不一，逼出 padding 行；最短的一条只有一个字。 + {"明天的会议", {"推迟到下周一", "推迟", "退", "推迟到下周义"}}, + }; + for (const auto &[context, texts] : cases) + { + const std::vector actual = model->score_sentences(context, texts); + expect(actual.size() == texts.size(), "score_sentences should return one score per text"); + if (actual.size() != texts.size()) + continue; + for (std::size_t i = 0; i < texts.size(); ++i) + { + // 容差放在 1e-3：两条路径的浮点累加次序不同，但结构上必须是同一个数。 + expect_near(actual[i], reference_score(*model, context, texts[i]), 1e-3, + "batched score should match the per-character reference for '" + texts[i] + "'"); + expect(actual[i] < 0.0, "a log probability should be negative for '" + texts[i] + "'"); + } + } + + // ---- 批量与单条一致：padding 行不能污染同批的其它行 ---- + { + const std::string context = "我今天想去"; + const std::vector texts = {"上海", "伤害", "吃饭"}; + const std::vector batched = model->score_sentences(context, texts); + for (std::size_t i = 0; i < texts.size(); ++i) + { + const std::vector alone = model->score_sentences(context, {texts[i]}); + expect_near(alone[0], batched[i], 1e-4, "scoring alone should match scoring in a batch"); + } + } + + // ---- 前文缓存必须随前文失效，且换回来后结果不变 ---- + { + const std::vector texts = {"上海"}; + const double first = model->score_sentences("我今天想去", texts)[0]; + const double other = model->score_sentences("我昨天去了", texts)[0]; + const double again = model->score_sentences("我今天想去", texts)[0]; + expect_near(again, first, 1e-6, "the same context must score the same after the cache was evicted"); + expect(std::fabs(other - first) > 1e-6, "a different context must actually change the score"); + } + + // ---- 超长输入不能崩：前文与候选都会被截断 ---- + { + std::string long_context; + for (int i = 0; i < 80; ++i) + long_context += "很长的前文。"; + expect(model->score_sentences(long_context, {"上海"})[0] < 0.0, "an over-long context should still score"); + std::string long_text; + for (int i = 0; i < 200; ++i) + long_text += "字"; + expect(model->score_sentences("", {long_text})[0] < 0.0, "an over-long candidate should still score"); + } + + // ---- rerank_order 的契约 ---- + { + const std::vector sentences = {"伤害", "上海"}; + // 词格分（log10）打平，让神经分说了算。 + const std::vector flat = {-5.0, -5.0}; + const std::vector order = neural::rerank_order(*model, "我今天想去", sentences, flat); + expect(order.size() == sentences.size(), "rerank_order should return a full permutation"); + if (order.size() == sentences.size()) + { + expect(order[0] == 1, "with the lattice tied, the model should prefer 上海 after 我今天想去"); + } + // lambda 为 0 等于不重排，调用方必须能看出「没排」。 + neural::RerankOptions off; + off.lambda = 0.0; + expect(neural::rerank_order(*model, "", sentences, flat, off).empty(), + "lambda of zero should report that no reranking happened"); + expect(neural::rerank_order(*model, "", {"只有一条"}, {-5.0}).empty(), "a single path cannot be reranked"); + } + + // ---- 前文确实在起作用，且按字数（不是字节数）裁剪 ---- + { + // 前文不是摆设：同一个候选换个前文，分数必须跟着变。上屏历史一路传到这里的 + // 链路（InputSession -> QueryRequest -> 词典 -> make_neural_reranker）断在任何 + // 一环，模型拿到的都是空串，这一条就会挂。 + const double after_go = model->score_sentences("我今天想去", {"上海"})[0]; + const double after_hurt = model->score_sentences("他对我造成了很大的", {"上海"})[0]; + expect(after_go != after_hurt, "the committed text must reach the model, not be dropped on the way"); + expect(after_go > after_hurt, "上海 should fit 我今天想去 better than 他对我造成了很大的"); + + // 上屏历史整段给进来，模型只看末尾若干字；裁剪按 UTF-8 字符边界，不能切出半个汉字。 + expect(neural::last_characters("我今天想去上海", 2) == "上海", "the tail should be counted in characters"); + expect(neural::last_characters("上海", 99) == "上海", "a short context should come back whole"); + expect(neural::last_characters("上海", 0).empty(), "a zero-character window should yield no context"); + } + + // ---- 后台线程：第一次必落空，算完之后同一批必命中 ---- + { + const std::vector sentences = {"伤害", "上海"}; + const std::vector flat = {-5.0, -5.0}; + std::mutex mutex; + std::condition_variable ready; + bool fired = false; + neural::RescoreWorker &worker = neural::RescoreWorker::instance(); + worker.clear(); + worker.set_ready_callback([&] { + { + std::lock_guard lock(mutex); + fired = true; + } + ready.notify_all(); + }); + + // 按键路径上的那一次：不等模型，立刻返回空，调用方照词格顺序出候选。 + expect(!worker.order_for(*model, "我今天想去", sentences, flat, {}).has_value(), + "the first request must not block the caller waiting for the model"); + + { + std::unique_lock lock(mutex); + expect(ready.wait_for(lock, std::chrono::seconds(30), [&] { return fired; }), + "the background thread should finish and call back"); + } + + // 重查那一次：结果表里已经有顺序了。 + const auto order = worker.order_for(*model, "我今天想去", sentences, flat, {}); + expect(order.has_value(), "the finished order should be waiting for the requery"); + if (order && order->size() == 2) + { + expect((*order)[0] == 1, "the background result should be the same ordering as the direct call"); + } + worker.set_ready_callback(nullptr); + worker.clear(); + } + + // 两个模型共用一个后台线程，但待办和结果必须按模型隔离，不能后提交的覆盖前一个。 + if (desktop_model != nullptr) + { + const std::vector sentences = {"伤害", "上海"}; + const std::vector flat = {-5.0, -5.0}; + std::mutex mutex; + std::condition_variable ready; + int fired = 0; + neural::RescoreWorker &worker = neural::RescoreWorker::instance(); + worker.clear(); + worker.set_ready_callback([&] { + { + std::lock_guard lock(mutex); + ++fired; + } + ready.notify_all(); + }); + + expect(!worker.order_for(*desktop_model, "我今天想去", sentences, flat, {}).has_value(), + "the desktop model should enqueue its own first request"); + expect(!worker.order_for(*model, "我今天想去", sentences, flat, {}).has_value(), + "the keyboard model should enqueue without replacing the desktop request"); + { + std::unique_lock lock(mutex); + expect(ready.wait_for(lock, std::chrono::seconds(30), [&] { return fired >= 2; }), + "both neural models should finish independently"); + } + expect(worker.order_for(*desktop_model, "我今天想去", sentences, flat, {}).has_value(), + "the desktop result should remain in its own cache entry"); + expect(worker.order_for(*model, "我今天想去", sentences, flat, {}).has_value(), + "the keyboard result should remain in its own cache entry"); + worker.set_ready_callback(nullptr); + worker.clear(); + } + + // ---- 延迟：这次重构的全部意义 ---- + { + const std::string context = "今天下午的会议讨论了输入法的排序问题，大家觉得整句转换还可以再准一些，"; + const std::vector texts = {"我们明天再讨论一下", "我们明天在讨论一下", "我们名天再讨论一下", + "我门明天再讨论一下", "我们明天再讨论以下", "我们明天再讨论一夏", + "我们明天再讨论一吓", "我们明天再讨论议下", "我们明天在讨论以下", + "我们名天在讨论一下", "我们明天再讨论一霞", "我们明天再讨论一侠"}; + model->score_sentences(context, texts); // 预热，把前文灌进缓存 + const auto start = std::chrono::steady_clock::now(); + constexpr int kRuns = 10; + for (int i = 0; i < kRuns; ++i) + model->score_sentences(context, texts); + const double cached_ms = + std::chrono::duration(std::chrono::steady_clock::now() - start).count() / kRuns; + + const auto cold_start = std::chrono::steady_clock::now(); + for (int i = 0; i < kRuns; ++i) + model->score_sentences(context + std::to_string(i), texts); // 每次换前文，逼它重算缓存 + const double cold_ms = + std::chrono::duration(std::chrono::steady_clock::now() - cold_start).count() / kRuns; + + std::printf("12 candidates, 34-char context: %.2f ms cached, %.2f ms recomputed\n", cached_ms, cold_ms); + } + + // ---- 集成：两种神经模式可同时提交重排，且不泄漏词格自己的首选 ---- + { + QuanpinDictionary dictionary; + SentenceAssociationOptions options; + options.neural_desktop = true; + options.neural_keyboard = true; + dictionary.set_sentence_association(options); + for (const char *query : {"nihaoshijie", "womenzaijianmian", "jintiantianqizhenhao"}) + { + const std::vector result = dictionary.query(query); + expect(!result.empty(), std::string("reranked query should still produce candidates for ") + query); + int lattice_rows = 0; + int desktop_rows = 0; + int keyboard_rows = 0; + for (const WordItem &item : result) + { + expect(!item.word.empty(), std::string("no candidate should be empty for ") + query); + if (item.source == CandidateSource::Generated) + ++lattice_rows; + else if (item.source == CandidateSource::NeuralDesktop) + ++desktop_rows; + else if (item.source == CandidateSource::NeuralKeyboard) + ++keyboard_rows; + } + // Trigram 显示开关没开，内部生成的路径只能交给神经模型，不能先冒充可见候选。 + expect(lattice_rows == 0, std::string("neural-only mode exposed a lattice row for ") + query); + expect(desktop_rows <= 1, std::string("reranking must not add duplicate desktop rows for ") + query); + expect(keyboard_rows <= 1, std::string("reranking must not add duplicate keyboard rows for ") + query); + } + } + + if (failures != 0) + { + std::fprintf(stderr, "%d assertion(s) failed.\n", failures); + return 1; + } + std::printf("Neural sentence reranking test passed.\n"); + return 0; +} diff --git a/engine/tests/src/test_online_input_session.cpp b/engine/tests/src/test_online_input_session.cpp index 9803fcac..a041ad21 100644 --- a/engine/tests/src/test_online_input_session.cpp +++ b/engine/tests/src/test_online_input_session.cpp @@ -134,6 +134,22 @@ int main() "A replacement cloud candidate was rejected."); require(quanpin.candidates()[1].word == "呢" && !has_word(quanpin, "泥"), "Replacing a cloud candidate left the previous cloud row visible."); + // 神经整句重排落地时宿主会 reset_cache + recompute 重查一次：云/AI 行只在 series cache 里， + // 会被一并清掉，宿主要拿当前 online_query 回填。重查不能推进请求 generation，否则回填会被拒。 + quanpin.reset_cache(); + quanpin.recompute_candidates(); + require(!has_word(quanpin, "呢") && !has_word(quanpin, "拟"), + "reset_cache no longer drops online rows; the host re-apply in ApplyRescoredOrder is redundant."); + const auto recomputed_query = quanpin.online_query(); + require(recomputed_query && + quanpin.apply_online_candidate(*recomputed_query, "呢", CandidateSource::CloudSuggestion) && + quanpin.apply_online_candidate(*recomputed_query, "拟", CandidateSource::AiSuggestion), + "Online rows could not be restored after a cache reset and recompute."); + require(quanpin.apply_online_candidate(*query, "尼", CandidateSource::CloudSuggestion), + "Recomputing candidates invalidated an in-flight online query."); + require(quanpin.apply_online_candidate(*query, "呢", CandidateSource::CloudSuggestion) && + quanpin.candidates()[1].word == "呢" && quanpin.candidates()[2].word == "拟", + "Restored online rows did not return to slots two and three."); const auto candidate_count = quanpin.candidates().size(); require(!quanpin.apply_online_candidate(*query, "你", CandidateSource::CloudSuggestion) && quanpin.candidates().size() == candidate_count, diff --git a/engine/tests/src/test_pinyin.cpp b/engine/tests/src/test_pinyin.cpp index f15927c1..dd1f1570 100644 --- a/engine/tests/src/test_pinyin.cpp +++ b/engine/tests/src/test_pinyin.cpp @@ -801,6 +801,61 @@ void test_word_lattice() "A leading exact-key hit still outranks the sentence; the prefix row behind it does not."); } + { + // 整句只出各来源的首选，n-best 的其余部分是解码中间产物，不进候选区。 + // 没有重排器时就一条；重排器改了主意时，把它的首选和词格的首选各留一条。 + std::unordered_map> table; + table["ni"] = {{"ni", "你", 12000}, {"ni", "拟", 3000}, {"ni", "泥", 2000}}; + table["hao"] = {{"hao", "好", 12000}, {"hao", "号", 3000}, {"hao", "耗", 2000}}; + + std::vector alone; + quanpin::merge_lattice_candidates(alone, {"ni", "hao"}, make_table_lattice_lookup(table), "ni'hao"); + expect(alone.size() == 1, fmt::format("Expected a single lattice sentence, got {}", alone.size())); + + // 重排器把第二条顶到首位：两条首选都要留下，重排的在前。 + std::vector reranked; + quanpin::merge_lattice_candidates( + reranked, {"ni", "hao"}, make_table_lattice_lookup(table), "ni'hao", {}, + [](std::vector &paths) { + if (paths.size() >= 2) + std::swap(paths[0], paths[1]); + return true; + }, + CandidateSource::NeuralDesktop); + expect(reranked.size() == 2, + fmt::format("Expected the reranker's pick and the lattice's, got {}", reranked.size())); + if (reranked.size() == 2 && alone.size() == 1) + { + expect(reranked.front().word != alone.front().word && reranked[1].word == alone.front().word, + "Expected the reranked pick first and the lattice's own pick behind it."); + // 来源要分得开：候选框据此标注是谁挑的。 + expect(reranked.front().source == CandidateSource::NeuralDesktop && + reranked[1].source == CandidateSource::Generated, + "Expected the reranker's pick tagged with its own source and the lattice's with Generated."); + } + + // 重排器同意词格的排序时不该凭空多出一行，但来源要记在重排那边：模型看过并认可了它。 + std::vector agreed; + quanpin::merge_lattice_candidates( + agreed, {"ni", "hao"}, make_table_lattice_lookup(table), "ni'hao", {}, + [](std::vector &) { return true; }, CandidateSource::NeuralDesktop); + expect(agreed.size() == 1, + fmt::format("A reranker that agrees should add no extra row, got {}", agreed.size())); + if (agreed.size() == 1) + { + expect(agreed.front().source == CandidateSource::NeuralDesktop, + "A reranker that ran and agreed should still be credited for the row."); + } + + // 弃权（模型没加载、后台还没算完）与「跑了且同意」不是一回事：来源仍是词格。 + std::vector declined; + quanpin::merge_lattice_candidates( + declined, {"ni", "hao"}, make_table_lattice_lookup(table), "ni'hao", {}, + [](std::vector &) { return false; }, CandidateSource::NeuralDesktop); + expect(declined.size() == 1 && declined.front().source == CandidateSource::Generated, + "A reranker that declined must leave the row credited to the lattice."); + } + { std::unordered_map> table; table["g"] = {{"g", "个", 100}}; diff --git a/neural-model/NOTICE.md b/neural-model/NOTICE.md new file mode 100644 index 00000000..bb39b654 --- /dev/null +++ b/neural-model/NOTICE.md @@ -0,0 +1,20 @@ +# chinese-ime-lm sentence models + +The two neural sentence models shipped here are from the **chinese-ime-lm** project +(), release `model-v1`: + +- `sentence-model.safetensors` — "keyboard" preset (vocab 8192, 6 layers, n_embd 192, context 64). +- `sentence-model-desktop.safetensors` — "desktop" preset (vocab 12288, 8 layers, n_embd 448, context 128). + +Both are character-level Transformer language models used to decode pinyin into whole sentences. + +## License + +The model **weights** are licensed under **Apache-2.0** (recorded in each file's +`__metadata__["license"]`). The C++ inference code in `engine/neural/` is a port of the project's +Apache-2.0 reference implementation (`reference/`). + +## Attribution + +Trained on the Chinese portion of C4 (ODC-BY, ) and +LCCC (MIT, ). diff --git a/neural-model/lock.json b/neural-model/lock.json new file mode 100644 index 00000000..b2a248dc --- /dev/null +++ b/neural-model/lock.json @@ -0,0 +1,21 @@ +{ + "source": "https://github.com/metasequoiaime/chinese-ime-lm/releases/tag/model-v1", + "release_tag": "model-v1", + "license": "Apache-2.0", + "models": [ + { + "file": "sentence-model.safetensors", + "preset": "keyboard", + "url": "https://github.com/metasequoiaime/chinese-ime-lm/releases/download/model-v1/sentence-model.safetensors", + "sha256": "86ac529510cb3b4968a5e6ade83ec8080f5b34a0a681e75c362accbbd387d1c1", + "size": 4486280 + }, + { + "file": "sentence-model-desktop.safetensors", + "preset": "desktop", + "url": "https://github.com/metasequoiaime/chinese-ime-lm/releases/download/model-v1/sentence-model-desktop.safetensors", + "sha256": "0a6ecba69bf1d39c7eb49549c716477dd6c03fdf05757773e1f50435262fb469", + "size": 25480184 + } + ] +} diff --git a/scripts/fetch-neural-model.ps1 b/scripts/fetch-neural-model.ps1 new file mode 100644 index 00000000..c71c5c3d --- /dev/null +++ b/scripts/fetch-neural-model.ps1 @@ -0,0 +1,53 @@ +<# +.SYNOPSIS + 下载 chinese-ime-lm 的神经整句模型（safetensors）到 neural-model/，并按 lock.json 校验 sha256。 + +.DESCRIPTION + 与 build-language-model.ps1 对 sc.lm 的处理同理：模型二进制不进仓库（见 .gitignore）， + 只有 lock.json 和 NOTICE.md 进仓库。打包（installer/Prepare-PackageFiles.ps1）前跑一次 + 本脚本把两个模型取到位。已存在且校验通过的文件跳过重复下载（除非 -Force）。 + +.PARAMETER Force + 忽略已存在文件，强制重新下载。 +#> +[CmdletBinding()] +param( + [switch]$Force +) + +$ErrorActionPreference = 'Stop' +Set-StrictMode -Version Latest + +$here = Split-Path -Parent $PSScriptRoot +$modelDir = Join-Path $here 'neural-model' +$lockPath = Join-Path $modelDir 'lock.json' +if (-not (Test-Path -LiteralPath $lockPath)) { + throw "找不到 lock.json：$lockPath" +} +$lock = Get-Content -LiteralPath $lockPath -Raw | ConvertFrom-Json + +foreach ($model in $lock.models) { + $dest = Join-Path $modelDir $model.file + $needsDownload = $true + if ((Test-Path -LiteralPath $dest) -and -not $Force) { + $hash = (Get-FileHash -LiteralPath $dest -Algorithm SHA256).Hash.ToLowerInvariant() + if ($hash -eq $model.sha256.ToLowerInvariant()) { + Write-Host "已存在且校验通过，跳过：$($model.file)" + $needsDownload = $false + } + else { + Write-Host "校验不符，重新下载：$($model.file)" + } + } + if ($needsDownload) { + Write-Host "下载 $($model.file) ..." + Invoke-WebRequest -Uri $model.url -OutFile $dest + $hash = (Get-FileHash -LiteralPath $dest -Algorithm SHA256).Hash.ToLowerInvariant() + if ($hash -ne $model.sha256.ToLowerInvariant()) { + throw "$($model.file) sha256 校验失败：期望 $($model.sha256)，实际 $hash" + } + Write-Host " 校验通过：$($model.file)" + } +} + +Write-Host '神经整句模型准备完毕。'