mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-09-08 10:14:35 +08:00
86c520a336
### Summary Closes #18414. `rag/res/term.freq` is not shipped, and both term-weight implementations therefore assigned the same `300` fallback frequency to every lowercase Latin token. With no tokenizer frequency, NER, or POS signal, function words and content words received identical lexical boosts. This PR adds the same bounded out-of-vocabulary prior to Python and Go: - Use it only when the explicit DF dictionary or tokenizer has no frequency. - Count Latin, Greek, and Cyrillic letters, including uppercase and accented forms. - Keep the existing frequency of `300` for words up to three letters, halve it every two additional letters, and clamp it at `10`. - Reject digits, underscores, and logographic terms so Chinese and other existing fine-grained-tokenizer paths are unchanged. - Treat an absent optional `term.freq` as the supported fallback path without a startup warning, while still logging inaccessible or malformed dictionaries. A corpus-derived table was intentionally not added: that would require provenance/licensing decisions, language detection, and handling cross-language homographs. The bounded prior is deterministic, dependency-free, and fixes the equal-weight degradation for whitespace-delimited alphabetic languages without claiming corpus-specific precision. Python and Go consume one shared fixture covering ASCII, uppercase, accented Latin, Greek, Cyrillic, separators, invalid mixed tokens, and a CJK non-match. Both sides also verify the issue's ordering (`was < largest < supplier < equipment`) and that an explicit dictionary entry still takes precedence. Co-authored-by: Loong <184861530+yzl0ng@users.noreply.github.com>
896 lines
22 KiB
Go
896 lines
22 KiB
Go
// Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
|
|
package nlp
|
|
|
|
import (
|
|
"encoding/json"
|
|
"os"
|
|
"path/filepath"
|
|
"reflect"
|
|
"slices"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
type oovFrequencyFixture struct {
|
|
Term string `json:"term"`
|
|
Frequency *float64 `json:"frequency"`
|
|
}
|
|
|
|
func repositoryRoot(t *testing.T) string {
|
|
t.Helper()
|
|
dir, err := os.Getwd()
|
|
if err != nil {
|
|
t.Fatalf("get working directory: %v", err)
|
|
}
|
|
for {
|
|
if _, err := os.Stat(filepath.Join(dir, "go.mod")); err == nil {
|
|
return dir
|
|
}
|
|
parent := filepath.Dir(dir)
|
|
if parent == dir {
|
|
t.Fatal("repository root not found")
|
|
}
|
|
dir = parent
|
|
}
|
|
}
|
|
|
|
func TestAlphabeticOOVFrequencyMatchesSharedFixture(t *testing.T) {
|
|
fixturePath := filepath.Join(repositoryRoot(t), "test", "fixtures", "term_weight_oov.json")
|
|
data, err := os.ReadFile(fixturePath)
|
|
if err != nil {
|
|
t.Fatalf("read OOV fixture: %v", err)
|
|
}
|
|
|
|
var cases []oovFrequencyFixture
|
|
if err := json.Unmarshal(data, &cases); err != nil {
|
|
t.Fatalf("decode OOV fixture: %v", err)
|
|
}
|
|
for _, testCase := range cases {
|
|
got, ok := alphabeticOOVFrequency(testCase.Term)
|
|
if testCase.Frequency == nil {
|
|
if ok {
|
|
t.Errorf("alphabeticOOVFrequency(%q) = %v, true; want no fallback", testCase.Term, got)
|
|
}
|
|
continue
|
|
}
|
|
if !ok || got != *testCase.Frequency {
|
|
t.Errorf("alphabeticOOVFrequency(%q) = %v, %v; want %v, true", testCase.Term, got, ok, *testCase.Frequency)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestOOVWeightsFavorLongerContentTerms(t *testing.T) {
|
|
d := NewTermWeightDealer(t.TempDir())
|
|
terms := []string{"was", "largest", "supplier", "equipment"}
|
|
weights := d.Weights(terms, false)
|
|
if len(weights) != len(terms) {
|
|
t.Fatalf("Weights returned %d terms; want %d", len(weights), len(terms))
|
|
}
|
|
for i := 1; i < len(weights); i++ {
|
|
if weights[i-1].Weight >= weights[i].Weight {
|
|
t.Fatalf("weights are not increasing with OOV term length: %v", weights)
|
|
}
|
|
}
|
|
|
|
d.df["equipment"] = 1_000_000
|
|
weights = d.Weights([]string{"supplier", "equipment"}, false)
|
|
if weights[1].Weight >= weights[0].Weight {
|
|
t.Fatalf("dictionary frequency should take precedence over the OOV prior: %v", weights)
|
|
}
|
|
}
|
|
|
|
// TestNewTermWeightDealer tests the constructor
|
|
func TestNewTermWeightDealer(t *testing.T) {
|
|
// Test with empty resPath
|
|
d := NewTermWeightDealer("")
|
|
if d == nil {
|
|
t.Fatal("NewTermWeightDealer returned nil")
|
|
}
|
|
|
|
// Check stop words are initialized
|
|
if len(d.stopWords) == 0 {
|
|
t.Error("Stop words not initialized")
|
|
}
|
|
|
|
// Check stop word exists
|
|
if _, ok := d.stopWords["请问"]; !ok {
|
|
t.Error("Expected stop word '请问' not found")
|
|
}
|
|
|
|
// Test with non-existent resPath (should not panic)
|
|
d2 := NewTermWeightDealer("/nonexistent/path")
|
|
if d2 == nil {
|
|
t.Fatal("NewTermWeightDealer returned nil for non-existent path")
|
|
}
|
|
}
|
|
|
|
// TestNewTermWeightDealerWithMockFiles tests with mock dictionary files
|
|
func TestNewTermWeightDealerWithMockFiles(t *testing.T) {
|
|
// Create temporary directory with mock files
|
|
tmpDir := t.TempDir()
|
|
|
|
// Create mock ner.json
|
|
nerData := `{
|
|
"北京": "loca",
|
|
"腾讯": "corp",
|
|
"func": "func",
|
|
"toxic": "toxic"
|
|
}`
|
|
if err := os.WriteFile(filepath.Join(tmpDir, "ner.json"), []byte(nerData), 0644); err != nil {
|
|
t.Fatalf("Failed to create mock ner.json: %v", err)
|
|
}
|
|
|
|
// Create mock term.freq
|
|
freqData := "hello\t100\nworld\t200\ntest\t50\n"
|
|
if err := os.WriteFile(filepath.Join(tmpDir, "term.freq"), []byte(freqData), 0644); err != nil {
|
|
t.Fatalf("Failed to create mock term.freq: %v", err)
|
|
}
|
|
|
|
d := NewTermWeightDealer(tmpDir)
|
|
|
|
// Check NE dictionary
|
|
if ne := d.Ner("北京"); ne != "loca" {
|
|
t.Errorf("Expected NE 'loca' for '北京', got '%s'", ne)
|
|
}
|
|
if ne := d.Ner("腾讯"); ne != "corp" {
|
|
t.Errorf("Expected NE 'corp' for '腾讯', got '%s'", ne)
|
|
}
|
|
|
|
// Check DF dictionary
|
|
if df := d.GetDF(); len(df) != 3 {
|
|
t.Errorf("Expected 3 entries in DF, got %d", len(df))
|
|
}
|
|
}
|
|
|
|
// TestPretoken tests the pretokenization function
|
|
func TestPretoken(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
tests := []struct {
|
|
name string
|
|
txt string
|
|
num bool
|
|
stpwd bool
|
|
expected []string
|
|
}{
|
|
{
|
|
name: "simple text",
|
|
txt: "hello world",
|
|
num: false,
|
|
stpwd: true,
|
|
expected: []string{}, // May vary based on tokenizer
|
|
},
|
|
{
|
|
name: "with stop words",
|
|
txt: "请问你好吗",
|
|
num: false,
|
|
stpwd: true,
|
|
expected: []string{}, // Stop words should be removed
|
|
},
|
|
{
|
|
name: "with numbers (num=true)",
|
|
txt: "123",
|
|
num: true,
|
|
stpwd: true,
|
|
expected: []string{}, // Single digit may be filtered
|
|
},
|
|
{
|
|
name: "empty text",
|
|
txt: "",
|
|
num: false,
|
|
stpwd: true,
|
|
expected: []string{},
|
|
},
|
|
{
|
|
name: "only punctuation",
|
|
txt: ",。!?",
|
|
num: false,
|
|
stpwd: true,
|
|
expected: []string{},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := d.Pretoken(tt.txt, tt.num, tt.stpwd)
|
|
// Just check it doesn't panic and returns a slice
|
|
if result == nil {
|
|
t.Error("Pretoken returned nil")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestTokenMerge tests token merging
|
|
func TestTokenMerge(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
tests := []struct {
|
|
name string
|
|
tks []string
|
|
expected []string
|
|
}{
|
|
{
|
|
name: "empty input",
|
|
tks: []string{},
|
|
expected: []string{},
|
|
},
|
|
{
|
|
name: "single token",
|
|
tks: []string{"hello"},
|
|
expected: []string{"hello"},
|
|
},
|
|
{
|
|
name: "consecutive short tokens",
|
|
tks: []string{"a", "b", "c"},
|
|
expected: []string{"a b c"}, // Should merge
|
|
},
|
|
{
|
|
name: "mixed tokens",
|
|
tks: []string{"a", "hello", "b"},
|
|
expected: []string{"a", "hello", "b"},
|
|
},
|
|
{
|
|
name: "first term single char followed by multi-char",
|
|
tks: []string{"多", "工位"},
|
|
expected: []string{"多 工位"}, // Special case
|
|
},
|
|
{
|
|
name: "too many short tokens (>=5)",
|
|
tks: []string{"a", "b", "c", "d", "e", "f"},
|
|
expected: []string{"a b", "c d", "e f"}, // Merge in pairs
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := d.TokenMerge(tt.tks)
|
|
if !reflect.DeepEqual(result, tt.expected) {
|
|
// Debug: print detailed comparison
|
|
t.Errorf("TokenMerge(%v) = %v (len=%d), expected %v (len=%d)",
|
|
tt.tks, result, len(result), tt.expected, len(tt.expected))
|
|
for i, r := range result {
|
|
t.Errorf(" result[%d] = %q (len=%d)", i, r, len(r))
|
|
}
|
|
for i, e := range tt.expected {
|
|
t.Errorf(" expected[%d] = %q (len=%d)", i, e, len(e))
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestNer tests named entity recognition
|
|
func TestNer(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
|
|
// Create mock ner.json
|
|
nerData := `{
|
|
"北京": "loca",
|
|
"腾讯": "corp",
|
|
"阿里巴巴": "corp"
|
|
}`
|
|
if err := os.WriteFile(filepath.Join(tmpDir, "ner.json"), []byte(nerData), 0644); err != nil {
|
|
t.Fatalf("Failed to create mock ner.json: %v", err)
|
|
}
|
|
|
|
d := NewTermWeightDealer(tmpDir)
|
|
|
|
tests := []struct {
|
|
term string
|
|
expected string
|
|
}{
|
|
{"北京", "loca"},
|
|
{"腾讯", "corp"},
|
|
{"阿里巴巴", "corp"},
|
|
{"不存在", ""},
|
|
{"", ""},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.term, func(t *testing.T) {
|
|
result := d.Ner(tt.term)
|
|
if result != tt.expected {
|
|
t.Errorf("Ner('%s') = '%s', expected '%s'", tt.term, result, tt.expected)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestSplit tests text splitting
|
|
func TestSplit(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
tests := []struct {
|
|
name string
|
|
txt string
|
|
expected []string
|
|
}{
|
|
{
|
|
name: "simple split",
|
|
txt: "hello world test",
|
|
// Consecutive English words ending with letters are merged
|
|
expected: []string{"hello world test"},
|
|
},
|
|
{
|
|
name: "consecutive English words",
|
|
txt: "machine learning algorithm",
|
|
expected: []string{"machine learning algorithm"}, // Should merge
|
|
},
|
|
{
|
|
name: "mixed Chinese and English",
|
|
txt: "hello 世界 world",
|
|
// "hello" ends with letter, "世界" doesn't start with letter but doesn't end with letter either
|
|
expected: []string{"hello", "世界", "world"},
|
|
},
|
|
{
|
|
name: "empty string",
|
|
txt: "",
|
|
expected: []string{""},
|
|
},
|
|
{
|
|
name: "multiple spaces",
|
|
txt: "hello world",
|
|
// Multiple spaces are normalized, then merged if both end with letters
|
|
expected: []string{"hello world"},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
result := d.Split(tt.txt)
|
|
if !reflect.DeepEqual(result, tt.expected) {
|
|
t.Errorf("Split('%s') = %v (len=%d), expected %v (len=%d)",
|
|
tt.txt, result, len(result), tt.expected, len(tt.expected))
|
|
for i, r := range result {
|
|
t.Errorf(" result[%d] = %q", i, r)
|
|
}
|
|
for i, e := range tt.expected {
|
|
t.Errorf(" expected[%d] = %q", i, e)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestWeights tests weight calculation
|
|
func TestWeights(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
|
|
// Create mock ner.json
|
|
nerData := `{
|
|
"toxic": "toxic",
|
|
"func": "func",
|
|
"corp": "corp",
|
|
"loca": "loca"
|
|
}`
|
|
if err := os.WriteFile(filepath.Join(tmpDir, "ner.json"), []byte(nerData), 0644); err != nil {
|
|
t.Fatalf("Failed to create mock ner.json: %v", err)
|
|
}
|
|
|
|
// Create mock term.freq
|
|
freqData := "hello\t100\nworld\t200\n"
|
|
if err := os.WriteFile(filepath.Join(tmpDir, "term.freq"), []byte(freqData), 0644); err != nil {
|
|
t.Fatalf("Failed to create mock term.freq: %v", err)
|
|
}
|
|
|
|
d := NewTermWeightDealer(tmpDir)
|
|
|
|
t.Run("without preprocess", func(t *testing.T) {
|
|
tks := []string{"hello", "world", "123"}
|
|
weights := d.Weights(tks, false)
|
|
|
|
if len(weights) != len(tks) {
|
|
t.Errorf("Expected %d weights, got %d", len(tks), len(weights))
|
|
}
|
|
|
|
// Check weights sum to 1 (normalized)
|
|
sum := 0.0
|
|
for _, tw := range weights {
|
|
sum += tw.Weight
|
|
}
|
|
if sum < 0.99 || sum > 1.01 {
|
|
t.Errorf("Weights should sum to ~1, got %f", sum)
|
|
}
|
|
})
|
|
|
|
t.Run("with preprocess", func(t *testing.T) {
|
|
tks := []string{"hello world", "test"}
|
|
weights := d.Weights(tks, true)
|
|
|
|
// Check it doesn't panic and returns results
|
|
if weights == nil {
|
|
t.Error("Weights returned nil")
|
|
}
|
|
})
|
|
|
|
t.Run("empty input", func(t *testing.T) {
|
|
weights := d.Weights([]string{}, false)
|
|
if len(weights) != 0 {
|
|
t.Errorf("Expected empty weights for empty input, got %d", len(weights))
|
|
}
|
|
})
|
|
|
|
t.Run("ner weight effect", func(t *testing.T) {
|
|
tmpDir2 := t.TempDir()
|
|
nerData := `{"toxicterm": "toxic"}`
|
|
os.WriteFile(filepath.Join(tmpDir2, "ner.json"), []byte(nerData), 0644)
|
|
d2 := NewTermWeightDealer(tmpDir2)
|
|
|
|
tks := []string{"toxicterm", "normal"}
|
|
weights := d2.Weights(tks, false)
|
|
|
|
if len(weights) != 2 {
|
|
t.Fatalf("Expected 2 weights, got %d", len(weights))
|
|
}
|
|
|
|
// toxicterm should have higher weight (nerWeight=2)
|
|
if weights[0].Weight <= weights[1].Weight {
|
|
t.Error("Expected toxicterm to have higher weight than normal term")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestWeightsWithNER tests NER type weight effects
|
|
func TestWeightsWithNER(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
|
|
// Create mock ner.json with all types
|
|
nerData := `{
|
|
"toxic_word": "toxic",
|
|
"func_word": "func",
|
|
"corp_name": "corp",
|
|
"location": "loca",
|
|
"school": "sch",
|
|
"stock": "stock",
|
|
"firstname": "firstnm"
|
|
}`
|
|
if err := os.WriteFile(filepath.Join(tmpDir, "ner.json"), []byte(nerData), 0644); err != nil {
|
|
t.Fatalf("Failed to create mock ner.json: %v", err)
|
|
}
|
|
|
|
d := NewTermWeightDealer(tmpDir)
|
|
|
|
tests := []struct {
|
|
term string
|
|
expectedType string
|
|
}{
|
|
{"toxic_word", "toxic"},
|
|
{"func_word", "func"},
|
|
{"corp_name", "corp"},
|
|
{"location", "loca"},
|
|
{"school", "sch"},
|
|
{"stock", "stock"},
|
|
{"firstname", "firstnm"},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.term, func(t *testing.T) {
|
|
ne := d.Ner(tt.term)
|
|
if ne != tt.expectedType {
|
|
t.Errorf("Ner('%s') = '%s', expected '%s'", tt.term, ne, tt.expectedType)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestGetters tests the getter methods
|
|
func TestGetters(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
|
|
// Create mock files
|
|
nerData := `{"test": "type"}`
|
|
os.WriteFile(filepath.Join(tmpDir, "ner.json"), []byte(nerData), 0644)
|
|
os.WriteFile(filepath.Join(tmpDir, "term.freq"), []byte("word\t10\n"), 0644)
|
|
|
|
d := NewTermWeightDealer(tmpDir)
|
|
|
|
t.Run("GetStopWords", func(t *testing.T) {
|
|
sw := d.GetStopWords()
|
|
if len(sw) == 0 {
|
|
t.Error("GetStopWords returned empty map")
|
|
}
|
|
if _, ok := sw["请问"]; !ok {
|
|
t.Error("Expected stop word '请问' not in map")
|
|
}
|
|
})
|
|
|
|
t.Run("GetNE", func(t *testing.T) {
|
|
ne := d.GetNE()
|
|
if len(ne) != 1 {
|
|
t.Errorf("Expected 1 NE entry, got %d", len(ne))
|
|
}
|
|
if ne["test"] != "type" {
|
|
t.Error("NE dictionary content incorrect")
|
|
}
|
|
})
|
|
|
|
t.Run("GetDF", func(t *testing.T) {
|
|
df := d.GetDF()
|
|
if len(df) != 1 {
|
|
t.Errorf("Expected 1 DF entry, got %d", len(df))
|
|
}
|
|
if df["word"] != 10 {
|
|
t.Error("DF dictionary content incorrect")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestLoadDict tests dictionary loading
|
|
func TestLoadDict(t *testing.T) {
|
|
t.Run("load with frequency", func(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
content := "word1\t100\nword2\t200\nword3\t300\n"
|
|
fn := filepath.Join(tmpDir, "test.freq")
|
|
os.WriteFile(fn, []byte(content), 0644)
|
|
|
|
dict := loadDict(fn)
|
|
if len(dict) != 3 {
|
|
t.Errorf("Expected 3 entries, got %d", len(dict))
|
|
}
|
|
if dict["word1"] != 100 {
|
|
t.Errorf("Expected word1=100, got %d", dict["word1"])
|
|
}
|
|
})
|
|
|
|
t.Run("load without frequency (set mode)", func(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
content := "word1\nword2\nword3\n"
|
|
fn := filepath.Join(tmpDir, "test.freq")
|
|
os.WriteFile(fn, []byte(content), 0644)
|
|
|
|
dict := loadDict(fn)
|
|
if len(dict) != 3 {
|
|
t.Errorf("Expected 3 entries, got %d", len(dict))
|
|
}
|
|
// All values should be 0 in set mode
|
|
for k, v := range dict {
|
|
if v != 0 {
|
|
t.Errorf("Expected %s=0 in set mode, got %d", k, v)
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("load non-existent file", func(t *testing.T) {
|
|
dict := loadDict("/nonexistent/file.txt")
|
|
if dict == nil {
|
|
t.Error("loadDict should return empty map, not nil")
|
|
}
|
|
if len(dict) != 0 {
|
|
t.Error("loadDict should return empty map for non-existent file")
|
|
}
|
|
})
|
|
|
|
t.Run("load with malformed lines", func(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
content := "word1\t100\n\n\nword2\tnotanumber\nword3"
|
|
fn := filepath.Join(tmpDir, "test.freq")
|
|
os.WriteFile(fn, []byte(content), 0644)
|
|
|
|
dict := loadDict(fn)
|
|
// Should handle empty lines and invalid numbers gracefully
|
|
if len(dict) < 1 {
|
|
t.Error("Should handle malformed lines gracefully")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestWeightsNormalization tests weight normalization
|
|
func TestWeightsNormalization(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
tests := []struct {
|
|
name string
|
|
tks []string
|
|
}{
|
|
{
|
|
name: "single token",
|
|
tks: []string{"hello"},
|
|
},
|
|
{
|
|
name: "multiple tokens",
|
|
tks: []string{"hello", "world", "test"},
|
|
},
|
|
{
|
|
name: "many tokens",
|
|
tks: []string{"a", "b", "c", "d", "e"},
|
|
},
|
|
}
|
|
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
weights := d.Weights(tt.tks, false)
|
|
|
|
if len(weights) != len(tt.tks) {
|
|
t.Fatalf("Expected %d weights, got %d", len(tt.tks), len(weights))
|
|
}
|
|
|
|
// Sum should be approximately 1
|
|
sum := 0.0
|
|
for _, tw := range weights {
|
|
sum += tw.Weight
|
|
// Individual weights should be non-negative
|
|
if tw.Weight < 0 {
|
|
t.Errorf("Weight for '%s' is negative: %f", tw.Term, tw.Weight)
|
|
}
|
|
}
|
|
|
|
if sum < 0.99 || sum > 1.01 {
|
|
t.Errorf("Weights sum to %f, expected ~1.0", sum)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// TestSplitWithNER tests Split with NER considerations
|
|
func TestSplitWithNER(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
|
|
// Create mock ner.json
|
|
nerData := `{
|
|
"function": "func"
|
|
}`
|
|
os.WriteFile(filepath.Join(tmpDir, "ner.json"), []byte(nerData), 0644)
|
|
|
|
d := NewTermWeightDealer(tmpDir)
|
|
|
|
t.Run("func type should not merge", func(t *testing.T) {
|
|
// If one of the words has NE type "func", they should not merge
|
|
result := d.Split("hello function")
|
|
// "hello" and "function" should not merge because function has type "func"
|
|
if len(result) != 2 {
|
|
t.Logf("Result: %v", result)
|
|
}
|
|
})
|
|
}
|
|
|
|
// BenchmarkWeights benchmarks the Weights function
|
|
func BenchmarkWeights(b *testing.B) {
|
|
d := NewTermWeightDealer("")
|
|
tks := []string{"hello", "world", "this", "is", "a", "test", "of", "term", "weights", "calculation"}
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
d.Weights(tks, false)
|
|
}
|
|
}
|
|
|
|
// BenchmarkTokenMerge benchmarks the TokenMerge function
|
|
func BenchmarkTokenMerge(b *testing.B) {
|
|
d := NewTermWeightDealer("")
|
|
tks := []string{"a", "b", "c", "d", "e", "hello", "world", "x", "y", "z"}
|
|
|
|
b.ResetTimer()
|
|
for i := 0; i < b.N; i++ {
|
|
d.TokenMerge(tks)
|
|
}
|
|
}
|
|
|
|
// TestTermWeightStructure tests the TermWeight struct
|
|
func TestTermWeightStructure(t *testing.T) {
|
|
tw := TermWeight{
|
|
Term: "test",
|
|
Weight: 0.5,
|
|
}
|
|
|
|
if tw.Term != "test" {
|
|
t.Error("Term field incorrect")
|
|
}
|
|
if tw.Weight != 0.5 {
|
|
t.Error("Weight field incorrect")
|
|
}
|
|
}
|
|
|
|
// TestIntegration tests an integrated workflow
|
|
func TestIntegration(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
|
|
// Create mock dictionaries
|
|
nerData := `{
|
|
"北京": "loca",
|
|
"腾讯": "corp"
|
|
}`
|
|
os.WriteFile(filepath.Join(tmpDir, "ner.json"), []byte(nerData), 0644)
|
|
os.WriteFile(filepath.Join(tmpDir, "term.freq"), []byte("北京\t1000\n腾讯\t500\n"), 0644)
|
|
|
|
d := NewTermWeightDealer(tmpDir)
|
|
|
|
// Full workflow: text -> split -> pretoken -> token_merge -> weights
|
|
text := "北京 腾讯 公司"
|
|
|
|
// Step 1: Split
|
|
splitted := d.Split(text)
|
|
if len(splitted) == 0 {
|
|
t.Fatal("Split returned empty result")
|
|
}
|
|
|
|
// Step 2: Pretoken
|
|
var allTokens []string
|
|
for _, s := range splitted {
|
|
tokens := d.Pretoken(s, true, true)
|
|
allTokens = append(allTokens, tokens...)
|
|
}
|
|
|
|
// Step 3: Token merge
|
|
merged := d.TokenMerge(allTokens)
|
|
|
|
// Step 4: Calculate weights
|
|
weights := d.Weights(merged, false)
|
|
|
|
// Verify results
|
|
if len(weights) == 0 && len(merged) > 0 {
|
|
t.Error("Weights calculation failed")
|
|
}
|
|
|
|
// Check weights sum to 1
|
|
sum := 0.0
|
|
for _, w := range weights {
|
|
sum += w.Weight
|
|
}
|
|
if sum < 0.99 || sum > 1.01 {
|
|
t.Errorf("Final weights sum to %f, expected ~1.0", sum)
|
|
}
|
|
}
|
|
|
|
// TestWeightsEdgeCases tests edge cases for weight calculation
|
|
func TestWeightsEdgeCases(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
t.Run("numbers pattern", func(t *testing.T) {
|
|
tks := []string{"123,45", "abc"}
|
|
weights := d.Weights(tks, false)
|
|
if len(weights) != 2 {
|
|
t.Fatalf("Expected 2 weights, got %d", len(weights))
|
|
}
|
|
// Numbers should get nerWeight=2
|
|
})
|
|
|
|
t.Run("short letters pattern", func(t *testing.T) {
|
|
tks := []string{"ab", "abc"}
|
|
weights := d.Weights(tks, false)
|
|
if len(weights) != 2 {
|
|
t.Fatalf("Expected 2 weights, got %d", len(weights))
|
|
}
|
|
})
|
|
|
|
t.Run("letter pattern with spaces", func(t *testing.T) {
|
|
tks := []string{"hello world test"}
|
|
weights := d.Weights(tks, true)
|
|
// Should not panic
|
|
if weights == nil {
|
|
t.Error("Weights returned nil for letter pattern")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestPretokenWithNumbers tests pretoken with num parameter
|
|
func TestPretokenWithNumbers(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
t.Run("num=false filters single digits", func(t *testing.T) {
|
|
result := d.Pretoken("5", false, true)
|
|
// Single digit should be filtered when num=false
|
|
if slices.Contains(result, "5") {
|
|
t.Error("Single digit should be filtered when num=false")
|
|
}
|
|
})
|
|
|
|
t.Run("num=true keeps single digits", func(t *testing.T) {
|
|
result := d.Pretoken("5 123", true, true)
|
|
// Check at least something is returned
|
|
if len(result) == 0 {
|
|
t.Log("Single digit may still be filtered by other rules")
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestPretokenStopWords tests pretoken with stpwd parameter
|
|
func TestPretokenStopWords(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
t.Run("stpwd=true removes stop words", func(t *testing.T) {
|
|
result := d.Pretoken("请问", true, true)
|
|
// "请问" is a stop word
|
|
for _, r := range result {
|
|
if r == "请问" {
|
|
t.Error("Stop word should be removed when stpwd=true")
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("stpwd=false keeps stop words", func(t *testing.T) {
|
|
result := d.Pretoken("请问", true, false)
|
|
// With tokenizer, this might still filter it
|
|
_ = result
|
|
})
|
|
}
|
|
|
|
// TestTokenMergeEdgeCases tests edge cases for token merging
|
|
func TestTokenMergeEdgeCases(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
t.Run("nil input", func(t *testing.T) {
|
|
result := d.TokenMerge(nil)
|
|
if len(result) != 0 {
|
|
t.Error("TokenMerge(nil) should return empty slice")
|
|
}
|
|
})
|
|
|
|
t.Run("empty strings in input", func(t *testing.T) {
|
|
result := d.TokenMerge([]string{"", "a", "", "b", ""})
|
|
// Empty strings should be filtered
|
|
for _, r := range result {
|
|
if r == "" {
|
|
t.Error("Empty strings should be filtered")
|
|
}
|
|
}
|
|
})
|
|
|
|
t.Run("exactly 4 short tokens", func(t *testing.T) {
|
|
// 4 short tokens should be merged as one group (not split into pairs)
|
|
result := d.TokenMerge([]string{"a", "b", "c", "d"})
|
|
expected := []string{"a b c d"}
|
|
if !reflect.DeepEqual(result, expected) {
|
|
t.Errorf("Expected %v, got %v", expected, result)
|
|
}
|
|
})
|
|
|
|
t.Run("exactly 5 short tokens", func(t *testing.T) {
|
|
// 5 short tokens should be split into pairs
|
|
result := d.TokenMerge([]string{"a", "b", "c", "d", "e"})
|
|
// Should be: a b, c d (e is left? depends on implementation)
|
|
if len(result) < 2 {
|
|
t.Errorf("Expected at least 2 groups for 5 tokens, got %d: %v", len(result), result)
|
|
}
|
|
})
|
|
}
|
|
|
|
// TestSplitEdgeCases tests edge cases for splitting
|
|
func TestSplitEdgeCases(t *testing.T) {
|
|
d := NewTermWeightDealer("")
|
|
|
|
t.Run("tabs and spaces", func(t *testing.T) {
|
|
result := d.Split("hello\tworld\t\ttest")
|
|
// Tabs should be normalized to single space
|
|
hasTab := false
|
|
for _, r := range result {
|
|
if strings.Contains(r, "\t") {
|
|
hasTab = true
|
|
break
|
|
}
|
|
}
|
|
if hasTab {
|
|
t.Error("Tabs should be normalized")
|
|
}
|
|
})
|
|
|
|
t.Run("consecutive English with different NE types", func(t *testing.T) {
|
|
tmpDir := t.TempDir()
|
|
nerData := `{
|
|
"hello": "func",
|
|
"world": "corp"
|
|
}`
|
|
os.WriteFile(filepath.Join(tmpDir, "ner.json"), []byte(nerData), 0644)
|
|
d2 := NewTermWeightDealer(tmpDir)
|
|
|
|
result := d2.Split("hello world")
|
|
// Both have NE types, so they should NOT merge
|
|
if len(result) != 2 {
|
|
t.Errorf("Expected 2 tokens when both have NE types, got %d: %v", len(result), result)
|
|
}
|
|
})
|
|
}
|