fix: 支持GSE_DATA_PATH环境变量加载外部数据文件

This commit is contained in:
2026-04-21 09:56:34 +08:00
parent afa9062170
commit f671096dbe

View File

@@ -2,6 +2,8 @@ package utils
import ( import (
"context" "context"
"os"
"path/filepath"
"sort" "sort"
"sync" "sync"
@@ -40,7 +42,47 @@ type gseTool struct {
func newGseTool() (tool *gseTool, err error) { func newGseTool() (tool *gseTool, err error) {
// 1. 初始化分词器 // 1. 初始化分词器
var seg gse.Segmenter var seg gse.Segmenter
// 内置词典(无外部文件)
// 获取GSE数据文件路径
gseDataPath := os.Getenv("GSE_DATA_PATH")
if gseDataPath != "" {
// 使用外部数据文件
dictPath := filepath.Join(gseDataPath, "dict", "zh")
idfPath := filepath.Join(gseDataPath, "dict", "zh", "idf.txt")
stopPath := filepath.Join(gseDataPath, "dict", "zh", "stop.txt")
// 加载词典
err = seg.LoadDict(filepath.Join(dictPath, "dict.txt"))
if err != nil {
return
}
// 加载停用词
err = seg.LoadStop(stopPath)
if err != nil {
glog.Warning(context.Background(), "加载停用词失败,继续:", err)
}
// 2. 初始化 TF-IDF 提取器
tfidf := &extracker.TagExtracter{}
tfidf.WithGse(seg)
err = tfidf.LoadIdf(idfPath)
if err != nil {
return
}
// 3. 初始化 TextRank 提取器
tr := &extracker.TextRanker{}
tr.WithGse(seg)
tool = &gseTool{
seg: seg,
tfidf: tfidf,
tr: tr,
}
} else {
// 使用内置embed数据
err = seg.LoadDictEmbed() err = seg.LoadDictEmbed()
if err != nil { if err != nil {
return return
@@ -68,6 +110,7 @@ func newGseTool() (tool *gseTool, err error) {
tfidf: tfidf, tfidf: tfidf,
tr: tr, tr: tr,
} }
}
return return
} }