166 lines
3.6 KiB
Go
166 lines
3.6 KiB
Go
package utils
|
||
|
||
import (
|
||
"context"
|
||
"os"
|
||
"path/filepath"
|
||
"sort"
|
||
"sync"
|
||
|
||
"github.com/go-ego/gse"
|
||
"github.com/go-ego/gse/hmm/extracker"
|
||
"github.com/go-ego/gse/hmm/segment"
|
||
"github.com/gogf/gf/v2/os/glog"
|
||
)
|
||
|
||
// 全局工具实例（不再自动初始化）
|
||
var (
|
||
GseTool *gseTool
|
||
once sync.Once // 保证只初始化一次，线程安全
|
||
)
|
||
|
||
func InitGseTool(ctx context.Context) error {
|
||
var err error
|
||
once.Do(func() {
|
||
// 只执行一次初始化
|
||
GseTool, err = newGseTool()
|
||
})
|
||
if err != nil {
|
||
glog.Error(ctx, "gse 分词工具初始化失败:", err)
|
||
}
|
||
return err
|
||
}
|
||
|
||
// gseTool 关键词提取工具（gse v1.0.2 标准）
|
||
type gseTool struct {
|
||
seg gse.Segmenter
|
||
tfidf *extracker.TagExtracter
|
||
tr *extracker.TextRanker
|
||
}
|
||
|
||
// newGseTool 初始化工具（内置词典 + 停用词）
|
||
func newGseTool() (tool *gseTool, err error) {
|
||
// 1. 初始化分词器
|
||
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()
|
||
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
|
||
}
|
||
|
||
// Cut 分词（关键词提取唯一正确模式：精确模式 + HMM）
|
||
func (k *gseTool) Cut(text string) []string {
|
||
return k.seg.Cut(text, true)
|
||
}
|
||
|
||
// Keyword 最终输出：关键词 + 权重
|
||
type Keyword struct {
|
||
Word string `json:"word"`
|
||
Score float64 `json:"score"`
|
||
}
|
||
|
||
func (k *gseTool) Extract(text string, topN int) []Keyword {
|
||
// 1. 提取 TF-IDF
|
||
tfTags := k.extractTFIDF(text, topN)
|
||
|
||
// 2. 提取 TextRank
|
||
trTags := k.extractTextRank(text, topN)
|
||
|
||
// 3. 合并成最终关键词（业务最常用）
|
||
scoreMap := make(map[string]float64)
|
||
for _, tag := range tfTags {
|
||
scoreMap[tag.Text] = tag.Weight
|
||
}
|
||
for _, tag := range trTags {
|
||
scoreMap[tag.Text] = tag.Weight
|
||
}
|
||
|
||
// 转成切片并排序（高分在前）
|
||
res := make([]Keyword, 0, len(scoreMap))
|
||
for word, score := range scoreMap {
|
||
res = append(res, Keyword{Word: word, Score: score})
|
||
}
|
||
|
||
sort.Slice(res, func(i, j int) bool {
|
||
return res[i].Score > res[j].Score
|
||
})
|
||
|
||
return res
|
||
}
|
||
|
||
// ExtractTFIDF TF-IDF 关键词（带权重）90% 业务：文章标签、搜索、关键词
|
||
func (k *gseTool) extractTFIDF(text string, topN int) segment.Segments {
|
||
return k.tfidf.ExtractTags(text, topN)
|
||
}
|
||
|
||
// ExtractTextRank TextRank 关键词（带权重）长文本、摘要、语义理解
|
||
func (k *gseTool) extractTextRank(text string, topN int) segment.Segments {
|
||
return k.tr.TextRank(text, topN)
|
||
}
|