mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-23 08:56:42 +08:00
Port PR14140 and PR16881 to GO (#17102)
### Summary Port https://github.com/infiniflow/ragflow/pull/14140/ https://github.com/infiniflow/ragflow/pull/16881
This commit is contained in:
@@ -28,6 +28,8 @@
|
||||
#include <iostream>
|
||||
#include <cmath>
|
||||
#include <fstream>
|
||||
#include <algorithm>
|
||||
#include <unordered_set>
|
||||
// import :term;
|
||||
// import :stemmer;
|
||||
// import :analyzer;
|
||||
@@ -47,6 +49,28 @@ static const std::string WORDNET_PATH = "wordnet";
|
||||
|
||||
static const std::string OPENCC_PATH = "opencc";
|
||||
|
||||
// Map language names (lowercase) to Stemmer Language enum.
|
||||
// Used by SetLanguage() to configure language-specific stemming.
|
||||
static const std::pair<std::string, Language> SNOWBALL_LANGUAGE_MAP[] = {
|
||||
{"english", STEM_LANG_ENGLISH},
|
||||
{"dutch", STEM_LANG_DUTCH},
|
||||
{"german", STEM_LANG_GERMAN},
|
||||
{"french", STEM_LANG_FRENCH},
|
||||
{"spanish", STEM_LANG_SPANISH},
|
||||
{"italian", STEM_LANG_ITALIAN},
|
||||
{"portuguese", STEM_LANG_PORTUGUESE},
|
||||
{"portuguese br", STEM_LANG_PORTUGUESE},
|
||||
{"russian", STEM_LANG_RUSSIAN},
|
||||
{"arabic", STEM_LANG_UNKNOWN},
|
||||
{"danish", STEM_LANG_DANISH},
|
||||
{"finnish", STEM_LANG_FINNISH},
|
||||
{"hungarian", STEM_LANG_HUNGARIAN},
|
||||
{"norwegian", STEM_LANG_NORWEGIAN},
|
||||
{"romanian", STEM_LANG_ROMANIAN},
|
||||
{"swedish", STEM_LANG_SWEDISH},
|
||||
{"turkish", STEM_LANG_TURKISH},
|
||||
};
|
||||
|
||||
static const std::string REGEX_SPLIT_CHAR =
|
||||
R"#(([ ,\.<>/?;'\[\]\`!@#$%^&*$$\{\}\|_+=《》,。?、;‘’:“”【】~!¥%……()——-]+|[a-zA-Z\.-]+|[0-9,\.-]+))#";
|
||||
|
||||
@@ -79,11 +103,6 @@ static inline int32_t DecodeFreq(int32_t value) {
|
||||
return v1;
|
||||
}
|
||||
|
||||
static inline int32_t DecodePOSIndex(int32_t value) {
|
||||
// POS index is stored in the high 8 bits (bits 24-31)
|
||||
return static_cast<int32_t>(static_cast<uint32_t>(value) >> 24);
|
||||
}
|
||||
|
||||
void Split(const std::string &input, const std::string &split_pattern, std::vector<std::string> &result, bool keep_delim = false) {
|
||||
re2::RE2 pattern(split_pattern);
|
||||
re2::StringPiece leftover(input.data());
|
||||
@@ -364,17 +383,16 @@ struct CompiledRegex {
|
||||
class NLTKWordTokenizer {
|
||||
MacIntyreContractions contractions_;
|
||||
|
||||
// Static singleton instance
|
||||
static std::unique_ptr<NLTKWordTokenizer> instance_;
|
||||
static std::once_flag init_flag_;
|
||||
|
||||
public:
|
||||
// Static method to get the singleton instance
|
||||
// Magic Static singleton
|
||||
static NLTKWordTokenizer &GetInstance() {
|
||||
std::call_once(init_flag_, []() { instance_ = std::make_unique<NLTKWordTokenizer>(); });
|
||||
return *instance_;
|
||||
static NLTKWordTokenizer instance;
|
||||
return instance;
|
||||
}
|
||||
|
||||
NLTKWordTokenizer(const NLTKWordTokenizer &) = delete;
|
||||
NLTKWordTokenizer &operator=(const NLTKWordTokenizer &) = delete;
|
||||
|
||||
// Starting quotes.
|
||||
std::vector<std::pair<std::string, std::string>> STARTING_QUOTES = {
|
||||
{std::string(R"(([«“‘„]|[`]+))"), std::string(R"( $1 )")},
|
||||
@@ -618,10 +636,6 @@ private:
|
||||
}
|
||||
};
|
||||
|
||||
// Static member definitions for NLTKWordTokenizer singleton
|
||||
std::unique_ptr<NLTKWordTokenizer> NLTKWordTokenizer::instance_ = nullptr;
|
||||
std::once_flag NLTKWordTokenizer::init_flag_;
|
||||
|
||||
void SentenceSplitter(const std::string &text, std::vector<std::string> &result) {
|
||||
int error_code;
|
||||
PCRE2_SIZE error_offset;
|
||||
@@ -682,6 +696,31 @@ RAGAnalyzer::~RAGAnalyzer() {
|
||||
}
|
||||
}
|
||||
|
||||
void RAGAnalyzer::InitStemmer(Language language) {
|
||||
stemmer_->Init(language);
|
||||
use_lemmatizer_ = (language == STEM_LANG_ENGLISH);
|
||||
}
|
||||
|
||||
void RAGAnalyzer::SetLanguage(const std::string &language) {
|
||||
std::string lang_key = language;
|
||||
std::transform(lang_key.begin(), lang_key.end(), lang_key.begin(), [](unsigned char c) { return std::tolower(c); });
|
||||
lang_key.erase(lang_key.find_last_not_of(" \t") + 1);
|
||||
lang_key.erase(0, lang_key.find_first_not_of(" \t"));
|
||||
|
||||
Language stem_lang = STEM_LANG_UNKNOWN;
|
||||
for (const auto &pair : SNOWBALL_LANGUAGE_MAP) {
|
||||
if (pair.first == lang_key) {
|
||||
stem_lang = pair.second;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (stem_lang != STEM_LANG_UNKNOWN) {
|
||||
stemmer_->Init(stem_lang);
|
||||
use_lemmatizer_ = (stem_lang == STEM_LANG_ENGLISH);
|
||||
}
|
||||
}
|
||||
|
||||
int32_t RAGAnalyzer::Load() {
|
||||
fs::path root(dict_path_);
|
||||
|
||||
@@ -733,12 +772,6 @@ int32_t RAGAnalyzer::Load() {
|
||||
int32_t freq = std::stoi(results[1]);
|
||||
freq = int32_t(std::log(float(freq) / DENOMINATOR) + 0.5);
|
||||
int32_t pos_idx = pos_table_->GetPOSIndex(results[2]);
|
||||
// If the POS tag is not in the POS table (e.g. "eng" for
|
||||
// English words), default to index 0. Using -1 would cause
|
||||
// Encode() to produce a negative int32_t which the darts
|
||||
// trie rejects.
|
||||
if (pos_idx < 0)
|
||||
pos_idx = 0;
|
||||
int value = Encode(freq, pos_idx);
|
||||
trie_->Add(results[0], value);
|
||||
std::string rkey = RKey(results[0]);
|
||||
@@ -860,17 +893,14 @@ int32_t RAGAnalyzer::Freq(const std::string_view key) const {
|
||||
return static_cast<int32_t>(std::exp(v) * DENOMINATOR + 0.5);
|
||||
}
|
||||
|
||||
std::string RAGAnalyzer::Tag(std::string_view key) const {
|
||||
std::string lower_key = Key(std::string(key));
|
||||
int32_t encoded_value = trie_->Get(lower_key);
|
||||
if (encoded_value == -1) {
|
||||
return "";
|
||||
}
|
||||
int32_t pos_idx = DecodePOSIndex(encoded_value);
|
||||
if (pos_table_ == nullptr) {
|
||||
return "";
|
||||
}
|
||||
const char* pos_tag = pos_table_->GetPOS(pos_idx);
|
||||
std::string RAGAnalyzer::Tag(const std::string_view key) const {
|
||||
std::string lower_key = Key(key);
|
||||
int32_t v = trie_->Get(lower_key);
|
||||
if (v == -1) return "";
|
||||
uint32_t pos_idx = (static_cast<uint32_t>(v) >> 24) & 0xFF;
|
||||
if (pos_idx == 0) return "";
|
||||
if (pos_table_ == nullptr) return "";
|
||||
const char *pos_tag = pos_table_->GetPOS(static_cast<int32_t>(pos_idx));
|
||||
return pos_tag ? std::string(pos_tag) : "";
|
||||
}
|
||||
|
||||
@@ -1347,7 +1377,7 @@ std::string RAGAnalyzer::Merge(const std::string &tks_str) const {
|
||||
void RAGAnalyzer::MergeWithPosition(const std::vector<std::string> &tokens,
|
||||
const std::vector<std::pair<unsigned, unsigned>> &positions,
|
||||
std::vector<std::string> &merged_tokens,
|
||||
std::vector<std::pair<unsigned, unsigned>> &merged_positions) const {
|
||||
std::vector<std::pair<unsigned, unsigned>> &merged_positions) {
|
||||
// Filter out empty tokens first (like spaces) to match Merge behavior
|
||||
std::vector<std::string> filtered_tokens;
|
||||
std::vector<std::pair<unsigned, unsigned>> filtered_positions;
|
||||
@@ -1391,15 +1421,19 @@ void RAGAnalyzer::MergeWithPosition(const std::vector<std::string> &tokens,
|
||||
merged_positions = std::move(res_positions);
|
||||
}
|
||||
|
||||
void RAGAnalyzer::EnglishNormalize(const std::vector<std::string> &tokens, std::vector<std::string> &res) const {
|
||||
void RAGAnalyzer::EnglishNormalize(const std::vector<std::string> &tokens, std::vector<std::string> &res) {
|
||||
for (auto &t : tokens) {
|
||||
if (re2::RE2::PartialMatch(t, pattern1_)) { //"[a-zA-Z_-]+$"
|
||||
// Apply lowercase before lemmatization to match Python NLTK behavior
|
||||
char *lowercase_term = lowercase_string_buffer_.data();
|
||||
ToLower(t.c_str(), t.size(), lowercase_term, term_string_buffer_limit_);
|
||||
std::string lemma_term = wordnet_lemma_->Lemmatize(lowercase_term);
|
||||
std::string term_to_stem;
|
||||
if (use_lemmatizer_) {
|
||||
term_to_stem = wordnet_lemma_->Lemmatize(lowercase_term);
|
||||
} else {
|
||||
term_to_stem = lowercase_term;
|
||||
}
|
||||
std::string stem_term;
|
||||
stemmer_->Stem(lemma_term, stem_term);
|
||||
stemmer_->Stem(term_to_stem, stem_term);
|
||||
res.push_back(stem_term);
|
||||
} else {
|
||||
res.push_back(t);
|
||||
@@ -1733,7 +1767,7 @@ std::string PCRE2GlobalReplace(const std::string &text, const std::string &patte
|
||||
return result;
|
||||
}
|
||||
|
||||
std::string RAGAnalyzer::Tokenize(const std::string &line) const {
|
||||
std::string RAGAnalyzer::Tokenize(const std::string &line) {
|
||||
// Python-style simple tokenization: re.sub(r"\\W+", " ", line)
|
||||
std::string processed_line = PCRE2GlobalReplace(line, R"#(\W+)#", " ");
|
||||
std::string str1 = StrQ2B(processed_line);
|
||||
@@ -1756,12 +1790,16 @@ std::string RAGAnalyzer::Tokenize(const std::string &line) const {
|
||||
NLTKWordTokenizer::GetInstance().Tokenize(sentence, term_list);
|
||||
}
|
||||
for (unsigned i = 0; i < term_list.size(); ++i) {
|
||||
// Apply lowercase before lemmatization to match Python NLTK behavior
|
||||
char *lowercase_term = lowercase_string_buffer_.data();
|
||||
ToLower(term_list[i].c_str(), term_list[i].size(), lowercase_term, term_string_buffer_limit_);
|
||||
std::string lemma_term = wordnet_lemma_->Lemmatize(lowercase_term);
|
||||
std::string term_to_stem;
|
||||
if (use_lemmatizer_) {
|
||||
term_to_stem = wordnet_lemma_->Lemmatize(lowercase_term);
|
||||
} else {
|
||||
term_to_stem = lowercase_term;
|
||||
}
|
||||
std::string stem_term;
|
||||
stemmer_->Stem(lemma_term, stem_term);
|
||||
stemmer_->Stem(term_to_stem, stem_term);
|
||||
res.push_back(stem_term);
|
||||
}
|
||||
continue;
|
||||
@@ -1793,7 +1831,7 @@ std::string RAGAnalyzer::Tokenize(const std::string &line) const {
|
||||
return ret;
|
||||
}
|
||||
|
||||
std::pair<std::vector<std::string>, std::vector<std::pair<unsigned, unsigned>>> RAGAnalyzer::TokenizeWithPosition(const std::string &line) const {
|
||||
std::pair<std::vector<std::string>, std::vector<std::pair<unsigned, unsigned>>> RAGAnalyzer::TokenizeWithPosition(const std::string &line) {
|
||||
// Python-style simple tokenization: re.sub(r"\W+", " ", line)
|
||||
// Get processed line and position mapping from PCRE2GlobalReplace
|
||||
auto [processed_line, pcre2_pos_mapping] = PCRE2GlobalReplaceWithPosition(line, R"#(\W+)#", " ");
|
||||
@@ -1877,9 +1915,14 @@ std::pair<std::vector<std::string>, std::vector<std::pair<unsigned, unsigned>>>
|
||||
// Apply lowercase before lemmatization to match Python NLTK behavior
|
||||
char *lowercase_term = lowercase_string_buffer_.data();
|
||||
ToLower(term.c_str(), term.size(), lowercase_term, term_string_buffer_limit_);
|
||||
std::string lemma_term = wordnet_lemma_->Lemmatize(lowercase_term);
|
||||
std::string term_to_stem;
|
||||
if (use_lemmatizer_) {
|
||||
term_to_stem = wordnet_lemma_->Lemmatize(lowercase_term);
|
||||
} else {
|
||||
term_to_stem = lowercase_term;
|
||||
}
|
||||
std::string stem_term;
|
||||
stemmer_->Stem(lemma_term, stem_term);
|
||||
stemmer_->Stem(term_to_stem, stem_term);
|
||||
|
||||
tokens.push_back(stem_term);
|
||||
|
||||
@@ -1943,7 +1986,7 @@ std::pair<std::vector<std::string>, std::vector<std::pair<unsigned, unsigned>>>
|
||||
return {std::move(tokens), std::move(positions)};
|
||||
}
|
||||
|
||||
unsigned RAGAnalyzer::MapToOriginalPosition(unsigned processed_pos, const std::vector<std::pair<unsigned, unsigned>> &mapping) const {
|
||||
unsigned RAGAnalyzer::MapToOriginalPosition(unsigned processed_pos, const std::vector<std::pair<unsigned, unsigned>> &mapping) {
|
||||
for (const auto &[orig, proc] : mapping) {
|
||||
if (proc == processed_pos) {
|
||||
return orig;
|
||||
@@ -1964,7 +2007,7 @@ void RAGAnalyzer::TokenizeInnerWithPosition(const std::string &L,
|
||||
std::vector<std::string> &tokens,
|
||||
std::vector<std::pair<unsigned, unsigned>> &positions,
|
||||
unsigned base_pos,
|
||||
const std::vector<unsigned> *pos_mapping) const {
|
||||
const std::vector<unsigned> *pos_mapping) {
|
||||
auto [tks, s] = MaxForward(L);
|
||||
auto [tks1, s1] = MaxBackward(L);
|
||||
|
||||
@@ -2193,18 +2236,22 @@ void RAGAnalyzer::TokenizeInnerWithPosition(const std::string &L,
|
||||
void RAGAnalyzer::EnglishNormalizeWithPosition(const std::vector<std::string> &tokens,
|
||||
const std::vector<std::pair<unsigned, unsigned>> &positions,
|
||||
std::vector<std::string> &normalize_tokens,
|
||||
std::vector<std::pair<unsigned, unsigned>> &normalize_positions) const {
|
||||
std::vector<std::pair<unsigned, unsigned>> &normalize_positions) {
|
||||
for (size_t i = 0; i < tokens.size(); ++i) {
|
||||
const auto &token = tokens[i];
|
||||
const auto &[start_pos, end_pos] = positions[i];
|
||||
|
||||
if (re2::RE2::PartialMatch(token, pattern1_)) { //"[a-zA-Z_-]+$"
|
||||
// Apply lowercase before lemmatization to match Python NLTK behavior
|
||||
char *lowercase_term = lowercase_string_buffer_.data();
|
||||
ToLower(token.c_str(), token.size(), lowercase_term, term_string_buffer_limit_);
|
||||
std::string lemma_term = wordnet_lemma_->Lemmatize(lowercase_term);
|
||||
std::string term_to_stem;
|
||||
if (use_lemmatizer_) {
|
||||
term_to_stem = wordnet_lemma_->Lemmatize(lowercase_term);
|
||||
} else {
|
||||
term_to_stem = lowercase_term;
|
||||
}
|
||||
std::string stem_term;
|
||||
stemmer_->Stem(lemma_term, stem_term);
|
||||
stemmer_->Stem(term_to_stem, stem_term);
|
||||
|
||||
normalize_tokens.push_back(stem_term);
|
||||
normalize_positions.emplace_back(start_pos, end_pos);
|
||||
@@ -2218,7 +2265,7 @@ void RAGAnalyzer::EnglishNormalizeWithPosition(const std::vector<std::string> &t
|
||||
void RAGAnalyzer::FineGrainedTokenizeWithPosition(const std::string &tokens_str,
|
||||
const std::vector<std::pair<unsigned, unsigned>> &positions,
|
||||
std::vector<std::string> &fine_tokens,
|
||||
std::vector<std::pair<unsigned, unsigned>> &fine_positions) const {
|
||||
std::vector<std::pair<unsigned, unsigned>> &fine_positions) {
|
||||
std::vector<std::string> tks;
|
||||
Split(tokens_str, blank_pattern_, tks);
|
||||
|
||||
@@ -2331,7 +2378,7 @@ void RAGAnalyzer::FineGrainedTokenizeWithPosition(const std::string &tokens_str,
|
||||
// fine_tokens already contains the correct Chinese tokens
|
||||
}
|
||||
|
||||
void RAGAnalyzer::FineGrainedTokenize(const std::string &tokens, std::vector<std::string> &result) const {
|
||||
void RAGAnalyzer::FineGrainedTokenize(const std::string &tokens, std::vector<std::string> &result) {
|
||||
std::vector<std::string> tks;
|
||||
Split(tokens, blank_pattern_, tks);
|
||||
std::vector<std::string> res;
|
||||
@@ -2422,11 +2469,11 @@ void RAGAnalyzer::FineGrainedTokenize(const std::string &tokens, std::vector<std
|
||||
// return ret;
|
||||
}
|
||||
|
||||
int RAGAnalyzer::AnalyzeImpl(const Term &input, void *data, bool fine_grained, bool enable_position, HookType func) const {
|
||||
if (enable_position) {
|
||||
int RAGAnalyzer::AnalyzeImpl(const Term &input, void *data, HookType func) {
|
||||
if (enable_position_) {
|
||||
auto [tokens, positions] = TokenizeWithPosition(input.text_);
|
||||
|
||||
if (fine_grained) {
|
||||
if (fine_grained_) {
|
||||
std::vector<std::string> fine_tokens;
|
||||
std::vector<std::pair<unsigned, unsigned>> fine_positions;
|
||||
FineGrainedTokenizeWithPosition(Join(tokens, 0), positions, fine_tokens, fine_positions);
|
||||
@@ -2437,13 +2484,15 @@ int RAGAnalyzer::AnalyzeImpl(const Term &input, void *data, bool fine_grained, b
|
||||
for (size_t i = 0; i < tokens.size(); ++i) {
|
||||
if (tokens[i].empty())
|
||||
continue;
|
||||
if (IsStopword(tokens[i]))
|
||||
continue;
|
||||
const auto &[start_pos, end_pos] = positions[i];
|
||||
func(data, tokens[i].c_str(), tokens[i].size(), start_pos, end_pos, false, 0);
|
||||
}
|
||||
} else {
|
||||
std::string result = Tokenize(input.text_);
|
||||
std::vector<std::string> tokens;
|
||||
if (fine_grained) {
|
||||
if (fine_grained_) {
|
||||
FineGrainedTokenize(result, tokens);
|
||||
} else {
|
||||
Split(result, blank_pattern_, tokens);
|
||||
@@ -2452,8 +2501,20 @@ int RAGAnalyzer::AnalyzeImpl(const Term &input, void *data, bool fine_grained, b
|
||||
for (auto &t : tokens) {
|
||||
if (t.empty())
|
||||
continue;
|
||||
if (IsStopword(t))
|
||||
continue;
|
||||
func(data, t.c_str(), t.size(), offset++, 0, false, 0);
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
bool RAGAnalyzer::IsStopword(const std::string &term) {
|
||||
// Mirrors rag/nlp/term_weight.py Dealer.stop_words in RAGFlow.
|
||||
// Kept hard-coded (not loaded from a resource file) for parity with RAGFlow.
|
||||
static const std::unordered_set<std::string> kStopwords = {
|
||||
"请问", "您", "你", "我", "他", "是", "的", "就", "有", "于", "及", "即", "在", "为", "最",
|
||||
"从", "以", "了", "将", "与", "吗", "吧", "中", "#", "什么", "怎么", "哪个", "哪些", "啥", "相关",
|
||||
};
|
||||
return kStopwords.find(term) != kStopwords.end();
|
||||
}
|
||||
|
||||
@@ -50,7 +50,9 @@ public:
|
||||
|
||||
~RAGAnalyzer();
|
||||
|
||||
void InitStemmer(Language language) { stemmer_->Init(language); }
|
||||
void InitStemmer(Language language);
|
||||
|
||||
void SetLanguage(const std::string& language);
|
||||
|
||||
int32_t Load();
|
||||
|
||||
@@ -59,38 +61,41 @@ public:
|
||||
void SetEnablePosition(bool enable_position) { enable_position_ = enable_position; }
|
||||
|
||||
std::pair<std::vector<std::string>, std::vector<std::pair<unsigned, unsigned>>> TokenizeWithPosition(
|
||||
const std::string& line) const;
|
||||
std::string Tokenize(const std::string& line) const;
|
||||
const std::string& line);
|
||||
std::string Tokenize(const std::string& line);
|
||||
|
||||
void FineGrainedTokenize(const std::string& tokens, std::vector<std::string>& result) const;
|
||||
void FineGrainedTokenize(const std::string& tokens, std::vector<std::string>& result);
|
||||
|
||||
void TokenizeInnerWithPosition(const std::string& L,
|
||||
std::vector<std::string>& tokens,
|
||||
std::vector<std::pair<unsigned, unsigned>>& positions,
|
||||
unsigned base_pos,
|
||||
const std::vector<unsigned>* pos_mapping = nullptr) const;
|
||||
const std::vector<unsigned>* pos_mapping = nullptr);
|
||||
void FineGrainedTokenizeWithPosition(const std::string& tokens_str,
|
||||
const std::vector<std::pair<unsigned, unsigned>>& positions,
|
||||
std::vector<std::string>& fine_tokens,
|
||||
std::vector<std::pair<unsigned, unsigned>>& fine_positions) const;
|
||||
std::vector<std::pair<unsigned, unsigned>>& fine_positions);
|
||||
void EnglishNormalizeWithPosition(const std::vector<std::string>& tokens,
|
||||
const std::vector<std::pair<unsigned, unsigned>>& positions,
|
||||
std::vector<std::string>& normalize_tokens,
|
||||
std::vector<std::pair<unsigned, unsigned>>& normalize_positions) const;
|
||||
std::vector<std::pair<unsigned, unsigned>>& normalize_positions);
|
||||
unsigned MapToOriginalPosition(unsigned processed_pos,
|
||||
const std::vector<std::pair<unsigned, unsigned>>& mapping) const;
|
||||
const std::vector<std::pair<unsigned, unsigned>>& mapping);
|
||||
void MergeWithPosition(const std::vector<std::string>& tokens,
|
||||
const std::vector<std::pair<unsigned, unsigned>>& positions,
|
||||
std::vector<std::string>& merged_tokens,
|
||||
std::vector<std::pair<unsigned, unsigned>>& merged_positions) const;
|
||||
std::vector<std::pair<unsigned, unsigned>>& merged_positions);
|
||||
|
||||
void SplitByLang(const std::string& line, std::vector<std::pair<std::string, bool>>& txt_lang_pairs) const;
|
||||
|
||||
int32_t Freq(std::string_view key) const;
|
||||
|
||||
std::string Tag(std::string_view key) const;
|
||||
|
||||
static bool IsStopword(const std::string& term);
|
||||
|
||||
protected:
|
||||
int AnalyzeImpl(const Term& input, void* data, bool fine_grained, bool enable_position, HookType func) const;
|
||||
int AnalyzeImpl(const Term& input, void* data, HookType func);
|
||||
|
||||
private:
|
||||
static constexpr float DENOMINATOR = 1000000;
|
||||
@@ -130,7 +135,7 @@ private:
|
||||
|
||||
[[nodiscard]] std::string Merge(const std::string& tokens) const;
|
||||
|
||||
void EnglishNormalize(const std::vector<std::string>& tokens, std::vector<std::string>& res) const;
|
||||
void EnglishNormalize(const std::vector<std::string>& tokens, std::vector<std::string>& res);
|
||||
|
||||
public:
|
||||
[[nodiscard]] std::vector<std::pair<std::vector<std::string_view>, double>> GetBestTokensTopN(
|
||||
@@ -154,6 +159,8 @@ public:
|
||||
|
||||
mutable std::vector<char> lowercase_string_buffer_;
|
||||
|
||||
bool use_lemmatizer_{true};
|
||||
|
||||
bool fine_grained_{false};
|
||||
|
||||
bool enable_position_{false};
|
||||
|
||||
@@ -44,6 +44,12 @@ void RAGAnalyzer_SetEnablePosition(RAGAnalyzerHandle handle, bool enable_positio
|
||||
analyzer->SetEnablePosition(enable_position);
|
||||
}
|
||||
|
||||
void RAGAnalyzer_SetLanguage(RAGAnalyzerHandle handle, const char* language) {
|
||||
if (!handle || !language) return;
|
||||
RAGAnalyzer* analyzer = static_cast<RAGAnalyzer*>(handle);
|
||||
analyzer->SetLanguage(std::string(language));
|
||||
}
|
||||
|
||||
int RAGAnalyzer_Analyze(RAGAnalyzerHandle handle, const char* text, RAGTokenCallback callback) {
|
||||
if (!handle || !text || !callback) return -1;
|
||||
|
||||
@@ -88,43 +94,47 @@ RAGTokenList* RAGAnalyzer_TokenizeWithPosition(RAGAnalyzerHandle handle, const c
|
||||
|
||||
RAGAnalyzer* analyzer = static_cast<RAGAnalyzer*>(handle);
|
||||
|
||||
Term input;
|
||||
input.text_ = std::string(text);
|
||||
auto [tokens, positions] = analyzer->TokenizeWithPosition(std::string(text));
|
||||
if (analyzer->fine_grained_) {
|
||||
std::string joined_tokens;
|
||||
for (size_t i = 0; i < tokens.size(); ++i) {
|
||||
if (i > 0) joined_tokens += " ";
|
||||
joined_tokens += tokens[i];
|
||||
}
|
||||
|
||||
TermList output;
|
||||
// Pass fine_grained and enable_position=true to get position information
|
||||
analyzer->Analyze(input, output, analyzer->fine_grained_, true);
|
||||
std::vector<std::string> fine_tokens;
|
||||
std::vector<std::pair<unsigned, unsigned>> fine_positions;
|
||||
analyzer->FineGrainedTokenizeWithPosition(joined_tokens, positions, fine_tokens, fine_positions);
|
||||
tokens = std::move(fine_tokens);
|
||||
positions = std::move(fine_positions);
|
||||
}
|
||||
|
||||
// Allocate memory for the token list structure
|
||||
RAGTokenList* token_list = static_cast<RAGTokenList*>(malloc(sizeof(RAGTokenList)));
|
||||
if (!token_list) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
// Allocate memory for the tokens array
|
||||
token_list->tokens = nullptr;
|
||||
token_list->count = static_cast<uint32_t>(tokens.size());
|
||||
if (tokens.empty()) {
|
||||
return token_list;
|
||||
}
|
||||
|
||||
token_list->tokens = static_cast<RAGTokenWithPosition*>(
|
||||
malloc(sizeof(RAGTokenWithPosition) * output.size())
|
||||
malloc(sizeof(RAGTokenWithPosition) * tokens.size())
|
||||
);
|
||||
if (!token_list->tokens) {
|
||||
free(token_list);
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
token_list->count = static_cast<uint32_t>(output.size());
|
||||
|
||||
// Fill in the tokens
|
||||
for (size_t i = 0; i < output.size(); ++i) {
|
||||
// Allocate memory for the text and copy it
|
||||
token_list->tokens[i].text = static_cast<char*>(
|
||||
malloc(output[i].text_.size() + 1)
|
||||
);
|
||||
for (size_t i = 0; i < tokens.size(); ++i) {
|
||||
token_list->tokens[i].text = static_cast<char*>(malloc(tokens[i].size() + 1));
|
||||
if (token_list->tokens[i].text) {
|
||||
std::memcpy(token_list->tokens[i].text,
|
||||
output[i].text_.c_str(),
|
||||
output[i].text_.size() + 1);
|
||||
std::memcpy(token_list->tokens[i].text, tokens[i].c_str(), tokens[i].size() + 1);
|
||||
}
|
||||
token_list->tokens[i].offset = output[i].word_offset_;
|
||||
token_list->tokens[i].end_offset = output[i].end_offset_;
|
||||
token_list->tokens[i].offset = positions[i].first;
|
||||
token_list->tokens[i].end_offset = positions[i].second;
|
||||
}
|
||||
|
||||
return token_list;
|
||||
|
||||
@@ -40,6 +40,9 @@ void RAGAnalyzer_SetFineGrained(RAGAnalyzerHandle handle, bool fine_grained);
|
||||
// Set enable position tracking
|
||||
void RAGAnalyzer_SetEnablePosition(RAGAnalyzerHandle handle, bool enable_position);
|
||||
|
||||
// Set the language for stemming (e.g. "English", "Dutch", "German")
|
||||
void RAGAnalyzer_SetLanguage(RAGAnalyzerHandle handle, const char* language);
|
||||
|
||||
// Analyze text and call callback for each token
|
||||
// Returns: 0 on success, negative value on failure
|
||||
int RAGAnalyzer_Analyze(
|
||||
|
||||
@@ -112,6 +112,13 @@ void RAGAnalyzer_SetEnablePosition(RAGAnalyzerHandle handle, bool enable_positio
|
||||
fprintf(stderr, "[C_API] SetEnablePosition: %d\n", enable_position);
|
||||
}
|
||||
|
||||
void RAGAnalyzer_SetLanguage(RAGAnalyzerHandle handle, const char* language) {
|
||||
if (!handle || !language) return;
|
||||
RAGAnalyzer* analyzer = static_cast<RAGAnalyzer*>(handle);
|
||||
analyzer->SetLanguage(std::string(language));
|
||||
fprintf(stderr, "[C_API] SetLanguage: %s\n", language);
|
||||
}
|
||||
|
||||
int RAGAnalyzer_Analyze(RAGAnalyzerHandle handle, const char* text, RAGTokenCallback callback) {
|
||||
if (!handle || !text || !callback) return -1;
|
||||
|
||||
|
||||
@@ -96,6 +96,17 @@ func (a *Analyzer) SetEnablePosition(enablePosition bool) {
|
||||
C.RAGAnalyzer_SetEnablePosition(a.handle, C.bool(enablePosition))
|
||||
}
|
||||
|
||||
// SetLanguage configures the Snowball stemmer for the given language (e.g. "English", "Dutch").
|
||||
// Falls back to the English Porter stemmer for unmapped languages.
|
||||
func (a *Analyzer) SetLanguage(language string) {
|
||||
if a.handle == nil {
|
||||
return
|
||||
}
|
||||
cLang := C.CString(language)
|
||||
defer C.free(unsafe.Pointer(cLang))
|
||||
C.RAGAnalyzer_SetLanguage(a.handle, cLang)
|
||||
}
|
||||
|
||||
// Analyze analyzes the input text and returns all tokens
|
||||
func (a *Analyzer) Analyze(text string) ([]Token, error) {
|
||||
if a.handle == nil {
|
||||
|
||||
@@ -48,6 +48,8 @@ func (a *Analyzer) SetFineGrained(bool) {}
|
||||
|
||||
func (a *Analyzer) SetEnablePosition(bool) {}
|
||||
|
||||
func (a *Analyzer) SetLanguage(string) {}
|
||||
|
||||
func (a *Analyzer) Analyze(text string) ([]Token, error) {
|
||||
return nil, fmt.Errorf("rag_analyzer: cgo required for Analyze(%q)", text)
|
||||
}
|
||||
|
||||
@@ -118,9 +118,11 @@ func (c *QAChunkerComponent) invoke(_ context.Context, inputs map[string]any) (m
|
||||
}
|
||||
|
||||
chunks := make([]schema.ChunkDoc, 0, len(qaPairs))
|
||||
lang, _ := inputs["lang"].(string)
|
||||
tok := tokenizer.New(lang)
|
||||
for _, pair := range qaPairs {
|
||||
contentLTKS, _ := tokenizer.Tokenize(pair.Question)
|
||||
contentSMLTKS, _ := tokenizer.FineGrainedTokenize(contentLTKS)
|
||||
contentLTKS, _ := tok.Tokenize(pair.Question)
|
||||
contentSMLTKS, _ := tok.FineGrainedTokenize(contentLTKS)
|
||||
answer := rmQAPrefix(pair.Answer)
|
||||
if isMarkdown {
|
||||
answer = renderMarkdown(answer)
|
||||
@@ -258,9 +260,17 @@ func extractQAText(text string) []qaPair {
|
||||
return nil
|
||||
}
|
||||
lines := strings.Split(text, "\n")
|
||||
|
||||
delimiter := detectDelimiter(lines)
|
||||
|
||||
if delimiter == "\t" {
|
||||
return extractQATextTab(lines)
|
||||
}
|
||||
return extractQATextCSV(text, lines)
|
||||
}
|
||||
|
||||
// extractQATextTab handles tab-delimited Q&A where no CSV quoting
|
||||
// rules apply and physical lines always map 1:1 to records.
|
||||
func extractQATextTab(lines []string) []qaPair {
|
||||
var pairs []qaPair
|
||||
var question, answer string
|
||||
|
||||
@@ -268,7 +278,7 @@ func extractQAText(text string) []qaPair {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
parts := splitQA(line, delimiter)
|
||||
parts := strings.Split(line, "\t")
|
||||
if len(parts) != 2 {
|
||||
if question != "" {
|
||||
answer += "\n" + line
|
||||
@@ -287,6 +297,66 @@ func extractQAText(text string) []qaPair {
|
||||
return pairs
|
||||
}
|
||||
|
||||
// extractQATextCSV uses a full-text csv.Reader so that quoted fields
|
||||
// that span multiple physical lines are parsed correctly (mirrors the
|
||||
// Python fix in infiniflow/ragflow#16881).
|
||||
//
|
||||
// Because csv.Reader can merge several physical lines into one record,
|
||||
// we track the byte offset via InputOffset() and map it back to the
|
||||
// original lines slice so that malformed rows append the correct raw
|
||||
// continuation text.
|
||||
func extractQATextCSV(text string, lines []string) []qaPair {
|
||||
// Pre‑compute the byte offset where each physical line starts.
|
||||
lineStarts := make([]int, len(lines)+1)
|
||||
off := 0
|
||||
for i, l := range lines {
|
||||
lineStarts[i] = off
|
||||
off += len(l) + 1 // +1 for '\n'
|
||||
}
|
||||
lineStarts[len(lines)] = off // sentinel
|
||||
|
||||
r := csv.NewReader(strings.NewReader(text))
|
||||
r.LazyQuotes = true
|
||||
r.FieldsPerRecord = -1
|
||||
|
||||
var pairs []qaPair
|
||||
var question, answer string
|
||||
prevLine := 0
|
||||
|
||||
for {
|
||||
record, err := r.Read()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
|
||||
// Map InputOffset back to the physical lines consumed.
|
||||
endOff := int(r.InputOffset())
|
||||
curLine := prevLine
|
||||
for curLine < len(lineStarts) && lineStarts[curLine] < endOff {
|
||||
curLine++
|
||||
}
|
||||
|
||||
raw := strings.Join(lines[prevLine:curLine], "\n")
|
||||
prevLine = curLine
|
||||
|
||||
if len(record) != 2 {
|
||||
if question != "" {
|
||||
answer += "\n" + raw
|
||||
}
|
||||
continue
|
||||
}
|
||||
if question != "" && answer != "" {
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)})
|
||||
}
|
||||
question = record[0]
|
||||
answer = record[1]
|
||||
}
|
||||
if question != "" {
|
||||
pairs = append(pairs, qaPair{Question: strings.TrimSpace(question), Answer: strings.TrimSpace(answer)})
|
||||
}
|
||||
return pairs
|
||||
}
|
||||
|
||||
func detectDelimiter(lines []string) string {
|
||||
comma, tab := 0, 0
|
||||
for _, line := range lines {
|
||||
@@ -303,24 +373,6 @@ func detectDelimiter(lines []string) string {
|
||||
return ","
|
||||
}
|
||||
|
||||
func splitQA(line, delimiter string) []string {
|
||||
if delimiter == "\t" {
|
||||
parts := strings.Split(line, "\t")
|
||||
if len(parts) == 2 {
|
||||
return parts
|
||||
}
|
||||
return []string{line}
|
||||
}
|
||||
r := csv.NewReader(strings.NewReader(line))
|
||||
r.Comma = ','
|
||||
r.LazyQuotes = true
|
||||
records, err := r.Read()
|
||||
if err != nil || len(records) != 2 {
|
||||
return []string{line}
|
||||
}
|
||||
return records
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// JSON / structured QA extraction
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
@@ -38,6 +38,7 @@ package chunker
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/csv"
|
||||
"fmt"
|
||||
"html"
|
||||
"strings"
|
||||
@@ -107,13 +108,15 @@ func (c *TagChunkerComponent) invoke(_ context.Context, inputs map[string]any) (
|
||||
}
|
||||
|
||||
chunks := make([]schema.ChunkDoc, 0, len(pairs))
|
||||
lang, _ := inputs["lang"].(string)
|
||||
tok := tokenizer.New(lang)
|
||||
for _, pair := range pairs {
|
||||
content := strings.TrimSpace(pair.Content)
|
||||
if content == "" {
|
||||
continue
|
||||
}
|
||||
contentLTKS, _ := tokenizer.Tokenize(content)
|
||||
contentSMLTKS, _ := tokenizer.FineGrainedTokenize(contentLTKS)
|
||||
contentLTKS, _ := tok.Tokenize(content)
|
||||
contentSMLTKS, _ := tok.FineGrainedTokenize(contentLTKS)
|
||||
chunk := schema.ChunkDoc{
|
||||
ContentWithWeight: content,
|
||||
DocType: "text",
|
||||
@@ -133,11 +136,7 @@ type tagPair struct {
|
||||
Tags string
|
||||
}
|
||||
|
||||
// extractTagText ports tag.py:60-89 (txt) and tag.py:91-113 (csv):
|
||||
// every deformed line is accumulated into the running content; a line that
|
||||
// splits into exactly two fields is emitted as a (content+prefix, tags)
|
||||
// pair and the accumulator is reset. Delimiter selection (tab vs comma)
|
||||
// reuses the same detectDelimiter/splitQA helpers as QAChunker.
|
||||
// extractTagText ports tag.py:60-89 (txt) and tag.py:91-113 (csv).
|
||||
func extractTagText(text string) []tagPair {
|
||||
if text == "" {
|
||||
return nil
|
||||
@@ -145,13 +144,22 @@ func extractTagText(text string) []tagPair {
|
||||
lines := strings.Split(text, "\n")
|
||||
delimiter := detectDelimiter(lines)
|
||||
|
||||
if delimiter == "\t" {
|
||||
return extractTagTextTab(lines)
|
||||
}
|
||||
return extractTagTextCSV(text, lines)
|
||||
}
|
||||
|
||||
// extractTagTextTab handles tab-delimited text where no CSV quoting rules
|
||||
// apply; one physical line always maps to one record.
|
||||
func extractTagTextTab(lines []string) []tagPair {
|
||||
var pairs []tagPair
|
||||
content := ""
|
||||
for _, line := range lines {
|
||||
if strings.TrimSpace(line) == "" {
|
||||
continue
|
||||
}
|
||||
parts := splitQA(line, delimiter)
|
||||
parts := strings.Split(line, "\t")
|
||||
if len(parts) != 2 {
|
||||
content += "\n" + line
|
||||
continue
|
||||
@@ -163,6 +171,49 @@ func extractTagText(text string) []tagPair {
|
||||
return pairs
|
||||
}
|
||||
|
||||
// extractTagTextCSV uses a full-text csv.Reader so that quoted fields
|
||||
// spanning multiple physical lines are parsed correctly (see #16881).
|
||||
func extractTagTextCSV(text string, lines []string) []tagPair {
|
||||
lineStarts := make([]int, len(lines)+1)
|
||||
off := 0
|
||||
for i, l := range lines {
|
||||
lineStarts[i] = off
|
||||
off += len(l) + 1
|
||||
}
|
||||
lineStarts[len(lines)] = off
|
||||
|
||||
r := csv.NewReader(strings.NewReader(text))
|
||||
r.LazyQuotes = true
|
||||
r.FieldsPerRecord = -1
|
||||
|
||||
var pairs []tagPair
|
||||
content := ""
|
||||
prevLine := 0
|
||||
|
||||
for {
|
||||
record, err := r.Read()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
endOff := int(r.InputOffset())
|
||||
curLine := prevLine
|
||||
for curLine < len(lineStarts) && lineStarts[curLine] < endOff {
|
||||
curLine++
|
||||
}
|
||||
raw := strings.Join(lines[prevLine:curLine], "\n")
|
||||
prevLine = curLine
|
||||
|
||||
if len(record) != 2 {
|
||||
content += "\n" + raw
|
||||
continue
|
||||
}
|
||||
content += "\n" + record[0]
|
||||
pairs = append(pairs, tagPair{Content: content, Tags: record[1]})
|
||||
content = ""
|
||||
}
|
||||
return pairs
|
||||
}
|
||||
|
||||
// extractTagTable ports the spreadsheet (xlsx/csv → html table) path.
|
||||
// Python tag.py reads these via the Excel parser's (q, a) pairs, which
|
||||
// are exactly the first two columns of each row — so we take the same
|
||||
|
||||
@@ -399,6 +399,7 @@ type extractorInputs struct {
|
||||
llmID string
|
||||
systemPrompt string
|
||||
prompt string
|
||||
lang string
|
||||
chunks []map[string]any
|
||||
}
|
||||
|
||||
@@ -428,6 +429,9 @@ func (c *ExtractorComponent) resolveInputs(inputs map[string]any) extractorInput
|
||||
if v, ok := inputs["system_prompt"].(string); ok && v != "" {
|
||||
out.systemPrompt = v
|
||||
}
|
||||
if v, ok := inputs["lang"].(string); ok && v != "" {
|
||||
out.lang = v
|
||||
}
|
||||
for _, key := range extractorChunkInputOrder(inputs) {
|
||||
if chunks, ok := extractorChunkList(inputs[key]); ok {
|
||||
out.chunks = chunks
|
||||
@@ -579,7 +583,8 @@ func (c *ExtractorComponent) runAutoKeywords(ctx context.Context, in extractorIn
|
||||
return nil
|
||||
}
|
||||
ck["important_kwd"] = kwds
|
||||
tks, tkErr := tokenizer.Tokenize(strings.Join(kwds, " "))
|
||||
tok := tokenizer.New(in.lang)
|
||||
tks, tkErr := tok.Tokenize(strings.Join(kwds, " "))
|
||||
if tkErr == nil {
|
||||
ck["important_tks"] = tks
|
||||
}
|
||||
@@ -616,7 +621,8 @@ func (c *ExtractorComponent) runAutoQuestions(ctx context.Context, in extractorI
|
||||
return nil
|
||||
}
|
||||
ck["question_kwd"] = filtered
|
||||
tks, tkErr := tokenizer.Tokenize(strings.Join(filtered, "\n"))
|
||||
tok := tokenizer.New(in.lang)
|
||||
tks, tkErr := tok.Tokenize(strings.Join(filtered, "\n"))
|
||||
if tkErr == nil {
|
||||
ck["question_tks"] = tks
|
||||
}
|
||||
|
||||
@@ -65,6 +65,7 @@ var GlobalMetadataKeys = []string{
|
||||
"file",
|
||||
"tenant_id",
|
||||
"kb_id",
|
||||
"lang",
|
||||
}
|
||||
|
||||
// SeedIngestionGlobals copies the whitelisted run-level metadata from `in`
|
||||
|
||||
@@ -298,6 +298,7 @@ func (c *ParserComponent) Outputs() map[string]string {
|
||||
"pages": "[]schema.Page: parsed pages sorted by PageNumber.",
|
||||
"name": "string: the upstream file/document name (or doc_id when no name is available).",
|
||||
"output_format": "string: the active output format (\"text\" when emitting text pages).",
|
||||
"lang": "string: the language for tokenization (e.g. English, Dutch, Chinese).",
|
||||
"_ERROR": "string: set on short-circuit errors.",
|
||||
}
|
||||
}
|
||||
@@ -310,6 +311,7 @@ func (c *ParserComponent) Outputs() map[string]string {
|
||||
// "pages": []schema.Page (sorted by PageNumber),
|
||||
// "name": string (from inputs["doc_id"]),
|
||||
// "output_format": "text",
|
||||
// "lang": string (from inputs["lang"]; e.g. English, Dutch),
|
||||
// "_created_time": RFC3339Nano (via TrackElapsed),
|
||||
// "_elapsed_time": float64 seconds (via TrackElapsed),
|
||||
// }
|
||||
@@ -470,7 +472,8 @@ func (c *ParserComponent) Invoke(ctx context.Context, inputs map[string]any) (ma
|
||||
return nil, fmt.Errorf("Parser: %w", err)
|
||||
}
|
||||
sortPagesByNumber(parsed)
|
||||
out := buildParserOutputs(parsed, dispatched, filename, fileTypeExt)
|
||||
lang, _ := getString(inputs, "lang")
|
||||
out := buildParserOutputs(parsed, dispatched, filename, fileTypeExt, lang)
|
||||
// Forward the storage references so a downstream chunker can
|
||||
// re-acquire the source PDF and crop section images on demand,
|
||||
// instead of carrying the binary across the component boundary.
|
||||
|
||||
@@ -382,11 +382,14 @@ func pagesFromDispatch(pages []schema.Page) [][]byte {
|
||||
// at rag/flow/parser/parser.py:_invoke — the downstream chunker
|
||||
// / tokenizer / extractor components read the matching family
|
||||
// key, with "pages" as the universal fallback shape.
|
||||
func buildParserOutputs(parsed []schema.Page, dispatched parserDispatchResult, name string, fileType utility.FileType) map[string]any {
|
||||
func buildParserOutputs(parsed []schema.Page, dispatched parserDispatchResult, name string, fileType utility.FileType, lang string) map[string]any {
|
||||
out := map[string]any{
|
||||
"pages": toAnyPages(parsed),
|
||||
"name": name,
|
||||
}
|
||||
if lang != "" {
|
||||
out["lang"] = lang
|
||||
}
|
||||
if dispatched.Err == nil && dispatched.OutputFormat != "" {
|
||||
out["output_format"] = dispatched.OutputFormat
|
||||
switch dispatched.OutputFormat {
|
||||
|
||||
@@ -356,8 +356,10 @@ func (c *TokenizerComponent) Invoke(ctx context.Context, inputs map[string]any)
|
||||
|
||||
normalizeChunkTextFallback(chunks)
|
||||
|
||||
language := globals.GlobalOrInput(ctx, inputs, "lang", "English")
|
||||
|
||||
if contains(c.param.SearchMethod, "full_text") {
|
||||
if err := tokenizeChunks(chunks, titleStem); err != nil {
|
||||
if err := tokenizeChunks(chunks, titleStem, language); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
@@ -629,16 +631,21 @@ func normalizeChunkTextFallback(chunks []schema.ChunkDoc) {
|
||||
|
||||
// tokenizeChunks annotates each chunk with title_tks, content_ltks,
|
||||
// and (when applicable) question_tks / important_tks / summary fields.
|
||||
// Mirrors python tokenizer.py:130-185.
|
||||
func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error {
|
||||
// Mirrors python tokenizer.py:130-185 and rag/nlp/__init__.py tokenize() /
|
||||
// tokenize_chunks().
|
||||
//
|
||||
// language sets the Snowball stemmer language, matching Python's
|
||||
// rag_tokenizer.tokenizer.set_language(language) call inside tokenize().
|
||||
func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string, language string) error {
|
||||
tok := tokenizer.New(language)
|
||||
for i := range chunks {
|
||||
ck := &chunks[i]
|
||||
ck.ChunkOrderInt = intPtr(i)
|
||||
titleTk, err := tokenizer.Tokenize(titleStem)
|
||||
titleTk, err := tok.Tokenize(titleStem)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: title tokenize: %w", err)
|
||||
}
|
||||
titleSmTk, err := tokenizer.FineGrainedTokenize(titleTk)
|
||||
titleSmTk, err := tok.FineGrainedTokenize(titleTk)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: title fine-grain: %w", err)
|
||||
}
|
||||
@@ -651,7 +658,7 @@ func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error {
|
||||
if err := ck.SetExtraValue("question_kwd", strings.Split(q, "\n")); err != nil {
|
||||
return fmt.Errorf("Tokenizer: question keywords marshal: %w", err)
|
||||
}
|
||||
qt, err := tokenizer.Tokenize(q)
|
||||
qt, err := tok.Tokenize(q)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: question tokenize: %w", err)
|
||||
}
|
||||
@@ -663,7 +670,7 @@ func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error {
|
||||
if err := ck.SetExtraValue("important_kwd", utility.SplitKeywords(kw)); err != nil {
|
||||
return fmt.Errorf("Tokenizer: keyword list marshal: %w", err)
|
||||
}
|
||||
it, err := tokenizer.Tokenize(kw)
|
||||
it, err := tok.Tokenize(kw)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: keyword tokenize: %w", err)
|
||||
}
|
||||
@@ -672,7 +679,7 @@ func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error {
|
||||
}
|
||||
}
|
||||
if s := ck.Summary; strings.TrimSpace(s) != "" {
|
||||
st, err := tokenizer.Tokenize(s)
|
||||
st, err := tok.Tokenize(s)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: summary tokenize: %w", err)
|
||||
}
|
||||
@@ -680,7 +687,7 @@ func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error {
|
||||
st = s
|
||||
}
|
||||
ck.ContentLtks = st
|
||||
smt, err := tokenizer.FineGrainedTokenize(st)
|
||||
smt, err := tok.FineGrainedTokenize(st)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: summary fine-grain: %w", err)
|
||||
}
|
||||
@@ -689,7 +696,7 @@ func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error {
|
||||
}
|
||||
ck.ContentSmLtks = smt
|
||||
} else if t := ck.Text; strings.TrimSpace(t) != "" {
|
||||
tt, err := tokenizer.Tokenize(t)
|
||||
tt, err := tok.Tokenize(t)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: text tokenize: %w", err)
|
||||
}
|
||||
@@ -697,7 +704,7 @@ func tokenizeChunks(chunks []schema.ChunkDoc, titleStem string) error {
|
||||
tt = t
|
||||
}
|
||||
ck.ContentLtks = tt
|
||||
smt, err := tokenizer.FineGrainedTokenize(tt)
|
||||
smt, err := tok.FineGrainedTokenize(tt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("Tokenizer: text fine-grain: %w", err)
|
||||
}
|
||||
|
||||
@@ -748,7 +748,7 @@ func TestTokenizeChunks_SymbolOnlyTextFallsBackToRawText(t *testing.T) {
|
||||
{Text: "("},
|
||||
{Text: "*"},
|
||||
}
|
||||
err := tokenizeChunks(chunks, "test")
|
||||
err := tokenizeChunks(chunks, "test", "English")
|
||||
if err != nil {
|
||||
t.Fatalf("tokenizeChunks: %v", err)
|
||||
}
|
||||
@@ -769,7 +769,7 @@ func TestTokenizeChunks_WhitespaceSummaryShadowsTextBug(t *testing.T) {
|
||||
chunks := []schema.ChunkDoc{
|
||||
{Summary: " ", Text: "real content here"},
|
||||
}
|
||||
err := tokenizeChunks(chunks, "test")
|
||||
err := tokenizeChunks(chunks, "test", "English")
|
||||
if err != nil {
|
||||
t.Fatalf("tokenizeChunks: %v", err)
|
||||
}
|
||||
|
||||
@@ -344,6 +344,9 @@ func (s *PipelineExecutor) runPipelineWithDSL(ctx context.Context, dsl string) (
|
||||
}
|
||||
inputs["tenant_id"] = s.taskCtx.Tenant.ID
|
||||
inputs["kb_id"] = s.taskCtx.KB.ID
|
||||
if s.taskCtx.KB.Language != nil {
|
||||
inputs["lang"] = *s.taskCtx.KB.Language
|
||||
}
|
||||
|
||||
// Component params from Doc.ParserConfig — including the tenant LLM id
|
||||
// injected into Extractor components above — are passed to Run as
|
||||
|
||||
@@ -73,6 +73,12 @@ type analyzerPool struct {
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// defaultLanguage is applied to every analyzer instance on pool acquisition
|
||||
// to clear any sticky language state left by a previous task. Ingestion
|
||||
// callers should use the public API variants that accept an explicit
|
||||
// language override.
|
||||
const defaultLanguage = "English"
|
||||
|
||||
var (
|
||||
globalPool *analyzerPool
|
||||
poolOnce sync.Once
|
||||
@@ -395,8 +401,9 @@ func Close() {
|
||||
}
|
||||
}
|
||||
|
||||
// withAnalyzer executes the given function with an exclusive analyzer instance
|
||||
func withAnalyzer(fn func(*rag.Analyzer) error) error {
|
||||
// withAnalyzer acquires an analyzer instance, applies the given language
|
||||
// (with "" mapping to defaultLanguage), and executes fn.
|
||||
func withAnalyzer(lang string, fn func(*rag.Analyzer) error) error {
|
||||
if globalPool == nil {
|
||||
return fmt.Errorf("tokenizer pool not initialized")
|
||||
}
|
||||
@@ -407,11 +414,16 @@ func withAnalyzer(fn func(*rag.Analyzer) error) error {
|
||||
}
|
||||
defer globalPool.release(instance)
|
||||
|
||||
if lang == "" {
|
||||
lang = defaultLanguage
|
||||
}
|
||||
instance.analyzer.SetLanguage(lang)
|
||||
|
||||
return fn(instance.analyzer)
|
||||
}
|
||||
|
||||
// withAnalyzerResult executes the given function with an exclusive analyzer instance and returns a result
|
||||
func withAnalyzerResult[T any](fn func(*rag.Analyzer) (T, error)) (T, error) {
|
||||
// withAnalyzerResult is the result-returning variant of withAnalyzer.
|
||||
func withAnalyzerResult[T any](lang string, fn func(*rag.Analyzer) (T, error)) (T, error) {
|
||||
var result T
|
||||
if globalPool == nil {
|
||||
return result, fmt.Errorf("tokenizer pool not initialized")
|
||||
@@ -423,32 +435,63 @@ func withAnalyzerResult[T any](fn func(*rag.Analyzer) (T, error)) (T, error) {
|
||||
}
|
||||
defer globalPool.release(instance)
|
||||
|
||||
if lang == "" {
|
||||
lang = defaultLanguage
|
||||
}
|
||||
instance.analyzer.SetLanguage(lang)
|
||||
|
||||
return fn(instance.analyzer)
|
||||
}
|
||||
|
||||
// Tokenize tokenizes the text and returns a space-separated string of tokens
|
||||
type Tokenizer struct {
|
||||
lang string
|
||||
}
|
||||
|
||||
// New returns a request-scoped tokenizer. Empty language falls back to English.
|
||||
func New(lang string) Tokenizer {
|
||||
return Tokenizer{lang: lang}
|
||||
}
|
||||
|
||||
var defaultTokenizer = New("")
|
||||
|
||||
// Tokenize tokenizes the text and returns a space-separated string of tokens.
|
||||
// Example: "hello world" -> "hello world"
|
||||
//
|
||||
// NOTE: For Infinity engine, returns input unchanged to match python's behavior
|
||||
// NOTE: For Infinity engine, returns input unchanged to match python's behavior.
|
||||
func Tokenize(text string) (string, error) {
|
||||
return defaultTokenizer.Tokenize(text)
|
||||
}
|
||||
|
||||
// Tokenize tokenizes the text using the tokenizer's request-scoped language.
|
||||
func (t Tokenizer) Tokenize(text string) (string, error) {
|
||||
if engineTypeProvider() == "infinity" {
|
||||
return text, nil
|
||||
}
|
||||
return withAnalyzerResult(func(a *rag.Analyzer) (string, error) {
|
||||
return withAnalyzerResult(t.lang, func(a *rag.Analyzer) (string, error) {
|
||||
return a.Tokenize(text)
|
||||
})
|
||||
}
|
||||
|
||||
// TokenizeWithPosition tokenizes the text and returns a list of tokens with position information
|
||||
// TokenizeWithPosition tokenizes the text and returns a list of tokens with position information.
|
||||
func TokenizeWithPosition(text string) ([]rag.TokenWithPosition, error) {
|
||||
return withAnalyzerResult(func(a *rag.Analyzer) ([]rag.TokenWithPosition, error) {
|
||||
return defaultTokenizer.TokenizeWithPosition(text)
|
||||
}
|
||||
|
||||
// TokenizeWithPosition tokenizes the text using the tokenizer's request-scoped language.
|
||||
func (t Tokenizer) TokenizeWithPosition(text string) ([]rag.TokenWithPosition, error) {
|
||||
return withAnalyzerResult(t.lang, func(a *rag.Analyzer) ([]rag.TokenWithPosition, error) {
|
||||
return a.TokenizeWithPosition(text)
|
||||
})
|
||||
}
|
||||
|
||||
// Analyze analyzes the text and returns all tokens
|
||||
// Analyze analyzes the text and returns all tokens.
|
||||
func Analyze(text string) ([]rag.Token, error) {
|
||||
return withAnalyzerResult(func(a *rag.Analyzer) ([]rag.Token, error) {
|
||||
return defaultTokenizer.Analyze(text)
|
||||
}
|
||||
|
||||
// Analyze analyzes the text using the tokenizer's request-scoped language.
|
||||
func (t Tokenizer) Analyze(text string) ([]rag.Token, error) {
|
||||
return withAnalyzerResult(t.lang, func(a *rag.Analyzer) ([]rag.Token, error) {
|
||||
return a.Analyze(text)
|
||||
})
|
||||
}
|
||||
@@ -462,16 +505,23 @@ func SetFineGrained(fineGrained bool) {
|
||||
common.Debug("SetFineGrained is no-op in pool mode", zap.Bool("fine_grained", fineGrained))
|
||||
}
|
||||
|
||||
// FineGrainedTokenize performs fine-grained tokenization on space-separated tokens
|
||||
// FineGrainedTokenize performs fine-grained tokenization on space-separated
|
||||
// tokens.
|
||||
// Input: space-separated tokens (e.g., "hello world 测试")
|
||||
// Output: space-separated fine-grained tokens (e.g., "hello world 测 试")
|
||||
//
|
||||
// NOTE: For Infinity engine, returns input unchanged to match python's behavior
|
||||
// NOTE: For Infinity engine, returns input unchanged to match python's behavior.
|
||||
func FineGrainedTokenize(tokens string) (string, error) {
|
||||
return defaultTokenizer.FineGrainedTokenize(tokens)
|
||||
}
|
||||
|
||||
// FineGrainedTokenize performs fine-grained tokenization using the tokenizer's
|
||||
// request-scoped language.
|
||||
func (t Tokenizer) FineGrainedTokenize(tokens string) (string, error) {
|
||||
if engineTypeProvider() == "infinity" {
|
||||
return tokens, nil
|
||||
}
|
||||
return withAnalyzerResult(func(a *rag.Analyzer) (string, error) {
|
||||
return withAnalyzerResult(t.lang, func(a *rag.Analyzer) (string, error) {
|
||||
return a.FineGrainedTokenize(tokens)
|
||||
})
|
||||
}
|
||||
@@ -490,7 +540,7 @@ func IsInitialized() bool {
|
||||
// GetTermFreq returns the frequency of a term (matching Python rag_tokenizer.freq)
|
||||
// Returns: frequency value, or 0 if term not found
|
||||
func GetTermFreq(term string) int32 {
|
||||
result, _ := withAnalyzerResult(func(a *rag.Analyzer) (int32, error) {
|
||||
result, _ := withAnalyzerResult("", func(a *rag.Analyzer) (int32, error) {
|
||||
return a.GetTermFreq(term), nil
|
||||
})
|
||||
return result
|
||||
@@ -499,7 +549,7 @@ func GetTermFreq(term string) int32 {
|
||||
// GetTermTag returns the POS tag of a term (matching Python rag_tokenizer.tag)
|
||||
// Returns: POS tag string (e.g., "n", "v", "ns"), or empty string if term not found or no tag
|
||||
func GetTermTag(term string) string {
|
||||
result, _ := withAnalyzerResult(func(a *rag.Analyzer) (string, error) {
|
||||
result, _ := withAnalyzerResult("", func(a *rag.Analyzer) (string, error) {
|
||||
return a.GetTermTag(term), nil
|
||||
})
|
||||
return result
|
||||
|
||||
@@ -172,6 +172,67 @@ func TestConcurrentTokenize(t *testing.T) {
|
||||
t.Log("=== Test completed successfully ===")
|
||||
}
|
||||
|
||||
func TestConcurrentTokenizeLanguageIsolation(t *testing.T) {
|
||||
restore := saveEngineType()
|
||||
defer restore()
|
||||
RegisterEngineType(func() string { return "" })
|
||||
|
||||
cfg := &PoolConfig{
|
||||
DictPath: "",
|
||||
MinSize: 2,
|
||||
MaxSize: 8,
|
||||
IdleTimeout: 3 * time.Second,
|
||||
AcquireTimeout: 5 * time.Second,
|
||||
}
|
||||
|
||||
if err := Init(cfg); err != nil {
|
||||
t.Fatalf("Failed to initialize pool: %v", err)
|
||||
}
|
||||
defer Close()
|
||||
|
||||
sample := findEnglishDutchDifferentiator(t)
|
||||
|
||||
const goroutinesPerLang = 8
|
||||
const requestsPerGoroutine = 20
|
||||
|
||||
var wg sync.WaitGroup
|
||||
start := make(chan struct{})
|
||||
errors := make(chan string, goroutinesPerLang*requestsPerGoroutine*2)
|
||||
|
||||
run := func(tok Tokenizer, lang, want string) {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
for i := 0; i < requestsPerGoroutine; i++ {
|
||||
got, err := tok.Tokenize(sample.input)
|
||||
if err != nil {
|
||||
errors <- fmt.Sprintf("lang=%s req=%d unexpected error: %v", lang, i, err)
|
||||
return
|
||||
}
|
||||
if got != want {
|
||||
errors <- fmt.Sprintf("lang=%s req=%d got %q want %q", lang, i, got, want)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for i := 0; i < goroutinesPerLang; i++ {
|
||||
wg.Add(2)
|
||||
go run(New("English"), "English", sample.english)
|
||||
go run(New("Dutch"), "Dutch", sample.dutch)
|
||||
}
|
||||
|
||||
close(start)
|
||||
wg.Wait()
|
||||
close(errors)
|
||||
|
||||
for err := range errors {
|
||||
t.Error(err)
|
||||
}
|
||||
if t.Failed() {
|
||||
t.Fatalf("concurrent language isolation failed for input %q (English=%q Dutch=%q)", sample.input, sample.english, sample.dutch)
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentTokenizeWithPosition tests concurrent tokenization with position info
|
||||
func TestConcurrentTokenizeWithPosition(t *testing.T) {
|
||||
cfg := &PoolConfig{
|
||||
|
||||
@@ -19,8 +19,15 @@ package tokenizer
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type languageDifferentiator struct {
|
||||
input string
|
||||
english string
|
||||
dutch string
|
||||
}
|
||||
|
||||
// saveEngineType saves the current engineTypeProvider and returns a function
|
||||
// to restore it. Use this when a test modifies the engine type to avoid
|
||||
// leaking global state between tests.
|
||||
@@ -223,6 +230,79 @@ func TestGetTermTag_PoolNotInitialized(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenize_DefaultLanguageResetsAnalyzerState(t *testing.T) {
|
||||
restore := saveEngineType()
|
||||
defer restore()
|
||||
RegisterEngineType(func() string { return "" })
|
||||
|
||||
if err := Init(&PoolConfig{
|
||||
DictPath: "",
|
||||
MinSize: 1,
|
||||
MaxSize: 1,
|
||||
IdleTimeout: 5 * time.Second,
|
||||
AcquireTimeout: 5 * time.Second,
|
||||
}); err != nil {
|
||||
t.Fatalf("Failed to initialize pool: %v", err)
|
||||
}
|
||||
defer Close()
|
||||
|
||||
sample := findEnglishDutchDifferentiator(t)
|
||||
|
||||
dutchGot, err := New("Dutch").Tokenize(sample.input)
|
||||
if err != nil {
|
||||
t.Fatalf("Tokenize(Dutch, %q) unexpected error: %v", sample.input, err)
|
||||
}
|
||||
if dutchGot != sample.dutch {
|
||||
t.Fatalf("Tokenize(Dutch, %q) = %q, want %q", sample.input, dutchGot, sample.dutch)
|
||||
}
|
||||
|
||||
defaultGot, err := Tokenize(sample.input)
|
||||
if err != nil {
|
||||
t.Fatalf("Tokenize(default, %q) unexpected error: %v", sample.input, err)
|
||||
}
|
||||
if defaultGot != sample.english {
|
||||
t.Fatalf("Tokenize(default, %q) = %q, want explicit English result %q", sample.input, defaultGot, sample.english)
|
||||
}
|
||||
if defaultGot == dutchGot {
|
||||
t.Fatalf("Tokenize(default, %q) unexpectedly inherited Dutch analyzer state: %q", sample.input, defaultGot)
|
||||
}
|
||||
}
|
||||
|
||||
func findEnglishDutchDifferentiator(t *testing.T) languageDifferentiator {
|
||||
t.Helper()
|
||||
|
||||
candidates := []string{
|
||||
"running",
|
||||
"jumps",
|
||||
"ponies",
|
||||
"studies",
|
||||
"wolves",
|
||||
"relational",
|
||||
"conditionally",
|
||||
}
|
||||
|
||||
for _, input := range candidates {
|
||||
english, err := New("English").Tokenize(input)
|
||||
if err != nil {
|
||||
t.Fatalf("Tokenize(English, %q) unexpected error: %v", input, err)
|
||||
}
|
||||
dutch, err := New("Dutch").Tokenize(input)
|
||||
if err != nil {
|
||||
t.Fatalf("Tokenize(Dutch, %q) unexpected error: %v", input, err)
|
||||
}
|
||||
if english != dutch {
|
||||
return languageDifferentiator{
|
||||
input: input,
|
||||
english: english,
|
||||
dutch: dutch,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
t.Skip("no differentiating tokenizer sample found for English vs Dutch")
|
||||
return languageDifferentiator{}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Global state tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user