mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-24 17:36:47 +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:
@@ -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