common/gse/gse.go

package utils

import (
	"context"
	"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数据文件在可执行文件同级的gse/dict目录下
	gseDataPath := "gse"
	dictPath := filepath.Join(gseDataPath, "dict", "zh")
	idfPath := filepath.Join(dictPath, "idf.txt")
	stopPath := filepath.Join(dictPath, "stop.txt")

	// 加载词典
	err = seg.LoadDict(filepath.Join(dictPath, "dict.txt"))
	if err != nil {
		glog.Warning(context.Background(), "加载gse词典失败，尝试embed模式:", err)
		// 回退到embed模式
		err = seg.LoadDictEmbed()
		if err != nil {
			return
		}
		err = seg.LoadStopEmbed()
		if err != nil {
			return
		}
		tfidf := &extracker.TagExtracter{}
		tfidf.WithGse(seg)
		err = tfidf.LoadIdf()
		if err != nil {
			return
		}
		tr := &extracker.TextRanker{}
		tr.WithGse(seg)
		tool = &gseTool{
			seg:   seg,
			tfidf: tfidf,
			tr:    tr,
		}
		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,
	}
	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)
}