go get github.com/cloudwego/eino-ext/components/embedding/ollama

{
  "embedding_api_url": "http://localhost:11434",
  "embedding_model": "nomic-embed-text"
}

package rag

import (
	"context"
	"eino_study/config"
	"fmt"
	"log"
	"math"

	"github.com/cloudwego/eino-ext/components/embedding/ollama"
)

var Embeddder *ollama.Embedder

func NewEmbedder() *ollama.Embedder {
	ctx := context.Background()
	var err error
	config := config.NewConfig("../config/config.json")
	Embeddder, err = ollama.NewEmbedder(ctx, &ollama.EmbeddingConfig{
		BaseURL: config.EmbeddingAPIURL,
		Model:   config.EmbeddingModel,
	})
	if err != nil {
		log.Fatalf("NewEmbedder failed, err=%v", err)
	}
	return Embeddder
}

func Embedding(ctx context.Context, text []string) ([][]float64, error) {
	embedder := NewEmbedder()
	embeddings, err := embedder.EmbedStrings(ctx, text)
	if err != nil {
		log.Fatalf("EmbedStrings failed, err=%v", err)
	}
	// 归一化处理
	for i, v := range embeddings {
		embeddings[i] = NormalizeVector(v)
	}
	return embeddings, nil
}

func NormalizeVector(v []float64) []float64 {
	if len(v) == 0 {
		return nil
	}
	// 计算向量的模长
	a := 0.
	for _, x := range v {
		a += x * x
	}
	a = math.Sqrt(a)
	// 归一化向量
	for i := range v {
		v[i] /= a
	}
	return v
}

// 计算向量内积
// 如果向量已经归一化,那么内积的结果就是向量之间的夹角余弦值
func DotProduct(v1, v2 []float64) float64 {
	if len(v1) != len(v2) {
		return 0.
	}
	a := 0.
	for i := range v1 {
		a += v1[i] * v2[i]
	}
	return a
}

// 多个向量按位求平均
func AvgOfVector(vectors [][]float64) ([]float64, error) {
	if len(vectors) == 0 {
		return nil, nil
	}
	if len(vectors) == 1 {
		return NormalizeVector(vectors[0]), nil
	}
	sum := make([]float64, len(vectors[0]))
	for i := 0; i < len(vectors); i++ {
		if len(vectors[i]) != len(sum) {
			return nil, fmt.Errorf("%dth vector dim not equal to first vector", i+1)
		}
		for index, v := range vectors[i] {
			sum[index] += v
		}
	}
	for i := range sum {
		sum[i] /= float64(len(vectors))
	}
	return NormalizeVector(sum), nil
}

func WordSim() {
	// 定义待比较的词语
	words := []string{"大连理工大学", "中国海洋大学", "大连", "青岛"}
	ctx := context.Background()

	// 获取向量
	embeddings, err := Embedding(ctx, words)
	if err != nil {
		log.Fatal(err)
	}

	// 计算两两之间的相似度
	for i := 0; i < len(words); i++ {
		for j := i + 1; j < len(words); j++ {
			dot := DotProduct(embeddings[i], embeddings[j])
			fmt.Printf("sim of %s and %s is %.4f\n", words[i], words[j], dot)
		}
	}

	// 文档向量化示例
	// doc1: 大学相关文档
	doc1, _ := AvgOfVector([][]float64{embeddings[0], embeddings[1]})
	// doc2: 城市相关文档
	doc2, _ := AvgOfVector([][]float64{embeddings[2], embeddings[3]})

	// 测试新词与文档的相似度
	words = []string{"中国计量大学", "杭州"}
	vectors, _ := Embedding(ctx, words)
	for i, vector := range vectors {
		dot := DotProduct(vector, doc1)
		fmt.Printf("sim of word %s and doc1 is %.4f\n", words[i], dot)
		dot = DotProduct(vector, doc2)
		fmt.Printf("sim of word %s and doc2 is %.4f\n", words[i], dot)
	}
}

=== RUN   TestEmbedding
sim of 大连理工大学 and 中国海洋大学 is 0.8183
sim of 大连理工大学 and 大连 is 0.8241
sim of 大连理工大学 and 青岛 is 0.5443
sim of 中国海洋大学 and 大连 is 0.6569
sim of 中国海洋大学 and 青岛 is 0.5597
sim of 大连 and 青岛 is 0.5946

sim of word 中国计量大学 and doc1 is 0.9411
sim of word 中国计量大学 and doc2 is 0.6970
sim of word 杭州 and doc1 is 0.5964
sim of word 杭州 and doc2 is 0.6367
--- PASS: TestEmbedding (0.67s)
PASS