diff --git a/utils/gse.go b/utils/gse.go index 64539a6..b75c6ff 100644 --- a/utils/gse.go +++ b/utils/gse.go @@ -2,6 +2,8 @@ package utils import ( "context" + "os" + "path/filepath" "sort" "sync" @@ -40,33 +42,74 @@ type gseTool struct { func newGseTool() (tool *gseTool, err error) { // 1. 初始化分词器 var seg gse.Segmenter - // 内置词典(无外部文件) - err = seg.LoadDictEmbed() - if err != nil { - return - } - // 内置停用词(v1.0.2 标准) - err = seg.LoadStopEmbed() - if err != nil { - return - } - // 2. 初始化 TF-IDF 提取器 - tfidf := &extracker.TagExtracter{} - tfidf.WithGse(seg) - err = tfidf.LoadIdf() - if err != nil { - return - } + // 获取GSE数据文件路径 + gseDataPath := os.Getenv("GSE_DATA_PATH") - // 3. 初始化 TextRank 提取器 - tr := &extracker.TextRanker{} - tr.WithGse(seg) + if gseDataPath != "" { + // 使用外部数据文件 + dictPath := filepath.Join(gseDataPath, "dict", "zh") + idfPath := filepath.Join(gseDataPath, "dict", "zh", "idf.txt") + stopPath := filepath.Join(gseDataPath, "dict", "zh", "stop.txt") - tool = &gseTool{ - seg: seg, - tfidf: tfidf, - tr: tr, + // 加载词典 + 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() + if err != nil { + return + } + // 内置停用词(v1.0.2 标准) + err = seg.LoadStopEmbed() + if err != nil { + return + } + + // 2. 初始化 TF-IDF 提取器 + tfidf := &extracker.TagExtracter{} + tfidf.WithGse(seg) + err = tfidf.LoadIdf() + if err != nil { + return + } + + // 3. 初始化 TextRank 提取器 + tr := &extracker.TextRanker{} + tr.WithGse(seg) + + tool = &gseTool{ + seg: seg, + tfidf: tfidf, + tr: tr, + } } return }