diff --git a/internal/binding/cpp/rag_analyzer.cpp b/internal/binding/cpp/rag_analyzer.cpp index f9d52e7336..042571fde9 100644 --- a/internal/binding/cpp/rag_analyzer.cpp +++ b/internal/binding/cpp/rag_analyzer.cpp @@ -28,6 +28,8 @@ #include #include #include +#include +#include // 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 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(static_cast(value) >> 24); -} - void Split(const std::string &input, const std::string &split_pattern, std::vector &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 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(); }); - return *instance_; + static NLTKWordTokenizer instance; + return instance; } + NLTKWordTokenizer(const NLTKWordTokenizer &) = delete; + NLTKWordTokenizer &operator=(const NLTKWordTokenizer &) = delete; + // Starting quotes. std::vector> STARTING_QUOTES = { {std::string(R"(([«“‘„]|[`]+))"), std::string(R"( $1 )")}, @@ -618,10 +636,6 @@ private: } }; -// Static member definitions for NLTKWordTokenizer singleton -std::unique_ptr NLTKWordTokenizer::instance_ = nullptr; -std::once_flag NLTKWordTokenizer::init_flag_; - void SentenceSplitter(const std::string &text, std::vector &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(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(v) >> 24) & 0xFF; + if (pos_idx == 0) return ""; + if (pos_table_ == nullptr) return ""; + const char *pos_tag = pos_table_->GetPOS(static_cast(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 &tokens, const std::vector> &positions, std::vector &merged_tokens, - std::vector> &merged_positions) const { + std::vector> &merged_positions) { // Filter out empty tokens first (like spaces) to match Merge behavior std::vector filtered_tokens; std::vector> filtered_positions; @@ -1391,15 +1421,19 @@ void RAGAnalyzer::MergeWithPosition(const std::vector &tokens, merged_positions = std::move(res_positions); } -void RAGAnalyzer::EnglishNormalize(const std::vector &tokens, std::vector &res) const { +void RAGAnalyzer::EnglishNormalize(const std::vector &tokens, std::vector &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>> RAGAnalyzer::TokenizeWithPosition(const std::string &line) const { +std::pair, std::vector>> 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>> // 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>> return {std::move(tokens), std::move(positions)}; } -unsigned RAGAnalyzer::MapToOriginalPosition(unsigned processed_pos, const std::vector> &mapping) const { +unsigned RAGAnalyzer::MapToOriginalPosition(unsigned processed_pos, const std::vector> &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 &tokens, std::vector> &positions, unsigned base_pos, - const std::vector *pos_mapping) const { + const std::vector *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 &tokens, const std::vector> &positions, std::vector &normalize_tokens, - std::vector> &normalize_positions) const { + std::vector> &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 &t void RAGAnalyzer::FineGrainedTokenizeWithPosition(const std::string &tokens_str, const std::vector> &positions, std::vector &fine_tokens, - std::vector> &fine_positions) const { + std::vector> &fine_positions) { std::vector 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 &result) const { +void RAGAnalyzer::FineGrainedTokenize(const std::string &tokens, std::vector &result) { std::vector tks; Split(tokens, blank_pattern_, tks); std::vector res; @@ -2422,11 +2469,11 @@ void RAGAnalyzer::FineGrainedTokenize(const std::string &tokens, std::vector fine_tokens; std::vector> 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 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 kStopwords = { + "请问", "您", "你", "我", "他", "是", "的", "就", "有", "于", "及", "即", "在", "为", "最", + "从", "以", "了", "将", "与", "吗", "吧", "中", "#", "什么", "怎么", "哪个", "哪些", "啥", "相关", + }; + return kStopwords.find(term) != kStopwords.end(); +} diff --git a/internal/binding/cpp/rag_analyzer.h b/internal/binding/cpp/rag_analyzer.h index 9b3027ef9c..95da708609 100644 --- a/internal/binding/cpp/rag_analyzer.h +++ b/internal/binding/cpp/rag_analyzer.h @@ -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>> 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& result) const; + void FineGrainedTokenize(const std::string& tokens, std::vector& result); void TokenizeInnerWithPosition(const std::string& L, std::vector& tokens, std::vector>& positions, unsigned base_pos, - const std::vector* pos_mapping = nullptr) const; + const std::vector* pos_mapping = nullptr); void FineGrainedTokenizeWithPosition(const std::string& tokens_str, const std::vector>& positions, std::vector& fine_tokens, - std::vector>& fine_positions) const; + std::vector>& fine_positions); void EnglishNormalizeWithPosition(const std::vector& tokens, const std::vector>& positions, std::vector& normalize_tokens, - std::vector>& normalize_positions) const; + std::vector>& normalize_positions); unsigned MapToOriginalPosition(unsigned processed_pos, - const std::vector>& mapping) const; + const std::vector>& mapping); void MergeWithPosition(const std::vector& tokens, const std::vector>& positions, std::vector& merged_tokens, - std::vector>& merged_positions) const; + std::vector>& merged_positions); void SplitByLang(const std::string& line, std::vector>& 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& tokens, std::vector& res) const; + void EnglishNormalize(const std::vector& tokens, std::vector& res); public: [[nodiscard]] std::vector, double>> GetBestTokensTopN( @@ -154,6 +159,8 @@ public: mutable std::vector lowercase_string_buffer_; + bool use_lemmatizer_{true}; + bool fine_grained_{false}; bool enable_position_{false}; diff --git a/internal/binding/cpp/rag_analyzer_c_api.cpp b/internal/binding/cpp/rag_analyzer_c_api.cpp index 3ed07dc49e..3b2252384e 100644 --- a/internal/binding/cpp/rag_analyzer_c_api.cpp +++ b/internal/binding/cpp/rag_analyzer_c_api.cpp @@ -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(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(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 fine_tokens; + std::vector> 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(malloc(sizeof(RAGTokenList))); if (!token_list) { return nullptr; } - // Allocate memory for the tokens array + token_list->tokens = nullptr; + token_list->count = static_cast(tokens.size()); + if (tokens.empty()) { + return token_list; + } + token_list->tokens = static_cast( - malloc(sizeof(RAGTokenWithPosition) * output.size()) + malloc(sizeof(RAGTokenWithPosition) * tokens.size()) ); if (!token_list->tokens) { free(token_list); return nullptr; } - token_list->count = static_cast(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( - malloc(output[i].text_.size() + 1) - ); + for (size_t i = 0; i < tokens.size(); ++i) { + token_list->tokens[i].text = static_cast(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; diff --git a/internal/binding/cpp/rag_analyzer_c_api.h b/internal/binding/cpp/rag_analyzer_c_api.h index 63f34f4f3a..510a11bb3d 100644 --- a/internal/binding/cpp/rag_analyzer_c_api.h +++ b/internal/binding/cpp/rag_analyzer_c_api.h @@ -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( diff --git a/internal/binding/cpp/rag_analyzer_c_api_debug.cpp b/internal/binding/cpp/rag_analyzer_c_api_debug.cpp index 4601d6407e..7fceb11a99 100644 --- a/internal/binding/cpp/rag_analyzer_c_api_debug.cpp +++ b/internal/binding/cpp/rag_analyzer_c_api_debug.cpp @@ -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(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; diff --git a/internal/binding/rag_analyzer.go b/internal/binding/rag_analyzer.go index 66646e066e..c248ceff14 100644 --- a/internal/binding/rag_analyzer.go +++ b/internal/binding/rag_analyzer.go @@ -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 { diff --git a/internal/binding/rag_analyzer_nocgo.go b/internal/binding/rag_analyzer_nocgo.go index 769f09ad27..8a1092c7d3 100644 --- a/internal/binding/rag_analyzer_nocgo.go +++ b/internal/binding/rag_analyzer_nocgo.go @@ -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) } diff --git a/internal/ingestion/component/chunker/qa.go b/internal/ingestion/component/chunker/qa.go index 866b8e185d..4f42377c0f 100644 --- a/internal/ingestion/component/chunker/qa.go +++ b/internal/ingestion/component/chunker/qa.go @@ -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 // --------------------------------------------------------------------------- diff --git a/internal/ingestion/component/chunker/tag.go b/internal/ingestion/component/chunker/tag.go index 50bc752ac5..864d6ea5bc 100644 --- a/internal/ingestion/component/chunker/tag.go +++ b/internal/ingestion/component/chunker/tag.go @@ -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 diff --git a/internal/ingestion/component/extractor.go b/internal/ingestion/component/extractor.go index d8759f6e62..0a91db48c0 100644 --- a/internal/ingestion/component/extractor.go +++ b/internal/ingestion/component/extractor.go @@ -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 } diff --git a/internal/ingestion/component/globals/globals.go b/internal/ingestion/component/globals/globals.go index cdbbb9e3e7..3b744a8d2d 100644 --- a/internal/ingestion/component/globals/globals.go +++ b/internal/ingestion/component/globals/globals.go @@ -65,6 +65,7 @@ var GlobalMetadataKeys = []string{ "file", "tenant_id", "kb_id", + "lang", } // SeedIngestionGlobals copies the whitelisted run-level metadata from `in` diff --git a/internal/ingestion/component/parser.go b/internal/ingestion/component/parser.go index 589e3b5a90..432c8f3939 100644 --- a/internal/ingestion/component/parser.go +++ b/internal/ingestion/component/parser.go @@ -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. diff --git a/internal/ingestion/component/parser_dispatch.go b/internal/ingestion/component/parser_dispatch.go index a53accf475..e3a0ded7fb 100644 --- a/internal/ingestion/component/parser_dispatch.go +++ b/internal/ingestion/component/parser_dispatch.go @@ -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 { diff --git a/internal/ingestion/component/tokenizer.go b/internal/ingestion/component/tokenizer.go index a66285cb0b..bc6e9488a3 100644 --- a/internal/ingestion/component/tokenizer.go +++ b/internal/ingestion/component/tokenizer.go @@ -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) } diff --git a/internal/ingestion/component/tokenizer_test.go b/internal/ingestion/component/tokenizer_test.go index ccd14a17ce..a1096f0f71 100644 --- a/internal/ingestion/component/tokenizer_test.go +++ b/internal/ingestion/component/tokenizer_test.go @@ -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) } diff --git a/internal/ingestion/task/pipeline_executor.go b/internal/ingestion/task/pipeline_executor.go index 2598e8a266..76c941812e 100644 --- a/internal/ingestion/task/pipeline_executor.go +++ b/internal/ingestion/task/pipeline_executor.go @@ -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 diff --git a/internal/tokenizer/tokenizer.go b/internal/tokenizer/tokenizer.go index 67a67e068e..87ccc17312 100644 --- a/internal/tokenizer/tokenizer.go +++ b/internal/tokenizer/tokenizer.go @@ -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 diff --git a/internal/tokenizer/tokenizer_concurrent_test.go b/internal/tokenizer/tokenizer_concurrent_test.go index 166bcce54f..4fb6ae50be 100644 --- a/internal/tokenizer/tokenizer_concurrent_test.go +++ b/internal/tokenizer/tokenizer_concurrent_test.go @@ -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{ diff --git a/internal/tokenizer/tokenizer_test.go b/internal/tokenizer/tokenizer_test.go index 48bbaddef2..f81f94a4d8 100644 --- a/internal/tokenizer/tokenizer_test.go +++ b/internal/tokenizer/tokenizer_test.go @@ -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 // ---------------------------------------------------------------------------