Embedding 模块

脱离笨重服务的本地向量计算引擎，为向量化服务提供了更简单、更易用的方案。

---

模块概述与核心组件

pkg/embedding 模块专为文本与图像的稠密向量化（Embedding）设计。其最大亮点是完全通过纯 Go 实现推理（基于 ONNX Runtime），彻底打破了对笨重的 Ollama 本地服务端或 Python (Transformers) 环境的依赖。只需编译一个二进制包，您的系统便拥有了高性能的离线向量化引擎。

模块核心类/接口：

• embedding.Provider: 表示具备文本嵌入能力的基础接口。
• embedding.MultimodalProvider: 继承自 Provider，具备处理图像和多模态的能力。
• embedding.BatchProcessor: 高并发处理工厂，提供针对大批量文本的哈希缓存和并发分批处理。
• embedding.Downloader: 内置极客工具，能够直接从 HuggingFace 的分发网络静默下载分片化的模型并在本地加载。

---

开箱即用的 Embedding Providers

GoChat 为业界主流的优质开源向量模型提供了极简的工厂加载方法。即使模型不在您的本地，框架也能一键拉取。

BGEProvider

简介：Baidu BGE (BAAI General Embedding)，百度通用嵌入模型，基于 Transformer 架构，支持中文、英文、日文、韩文、法文、西班牙文、德文、意大利文、俄文、葡萄牙文、越南文、印度文、阿拉伯文、韩文、泰文、西班牙文、丹麦文、芬兰文、瑞典文、挪威文、波兰文、捷克文、罗马尼亚文、希腊文、匈牙利文、保加利亚文、爱沙尼亚文、丹麦文、

• 特性与优势：BGE（BAAI General Embedding）系列是目前综合表现极优的中文及其它多语言编码器，深度优化了基于不对称语境（短 Query 搜长 Document）的 RAG 检索场景。
• 适用场景：国内业务、需要强大中文语境识别的知识库系统。
• 如何使用：

• // 直接加载 BGE-small 模型 (如果不存在，将自动下载至默认目录 ~/.embedding)
  provider, err := embedding.WithBEG("bge-small-zh-v1.5", "")
  if err != nil {
      panic(err)
  }
  
  vectors, err := provider.Embed(context.Background(), []string{"自然语言处理"})
  fmt.Printf("BGE 向量维度: %d, 提取特征数: %d\n", provider.Dimension(), len(vectors[0]))
  // 输出: BGE 向量维度: 512, 提取特征数: 512
  

SentenceBERTProvider

SentenceBERTProvider 是基于 sentence-transformers 的 Sentence-BERT 模型封装。

• 特性与优势：基于传统的 Sentence-BERT (如 all-MiniLM-L6-v2) 架构，主要解决对称语义相似度（Semantic Textual Similarity）匹配。它的模型体积极小（约 45MB），加载和计算速度快到极致。
• 适用场景：英文业务、基础词句余弦相似度匹配、对内存及性能有严苛限制的边缘设备。
• 如何使用：

  // 直接加载 Sentence-BERT 英文模型
  provider, err := embedding.WithBERT("all-MiniLM-L6-v2", "")
  
  vectors, _ := provider.Embed(context.Background(), []string{"Machine learning is fascinating."})
  fmt.Printf("SentenceBERT 向量维度: %d\n", provider.Dimension())
  // 输出: SentenceBERT 向量维度: 384
  

---

海量数据利器：批量处理与进度追踪

当您需要处理一个包含上万段文本（Chunks）的 PDF 知识库时，直接调用 provider.Embed 极易引发 OOM（内存溢出）并导致计算闲置。

BatchProcessor (批处理器) 正是为工业级海量数据设计的：

1. 自动并发切批：通过 MaxBatchSize 控制每次输送给 ONNX 引擎的数组大小，并使用 MaxConcurrent 拉满多核心 CPU 算力。
2. LRU Hash 缓存：自动忽略内容完全相同的文本片段，避免了昂贵的底层张量重算。
3. 带取消机制的原子进度条：提供回调供界面渲染进度。

完整示例：高并发知识库入库

package main

import (
    "context"
    "fmt"
    "github.com/DotNetAge/gochat/pkg/embedding"
)

func main() {
    ctx := context.Background()

    // 1. 初始化 Provider
    provider, _ := embedding.WithBEG("bge-small-zh-v1.5", "")

    // 2. 初始化 BatchProcessor
    batchProcessor := embedding.NewBatchProcessor(provider, embedding.BatchOptions{
        MaxBatchSize:  32, // 每个批次 32 段文本
        MaxConcurrent: 4,  // 最高 4 线程并发推理
    })

    // 模拟从数据库或 PDF 中拆分的上万段短句
    texts := []string{
        "Go 是一种高性能并发编程语言。",
        "什么是文本向量化（Embedding）？",
        // ... (此处省略 10000 条数据)
        "这是最后一条数据。",
    }

    // 3. 执行进度跟踪计算
    embeddings, err := batchProcessor.ProcessWithProgress(ctx, texts, func(current, total int, err error) bool {
        if err != nil {
            fmt.Printf("处理失败: %v\n", err)
            return false // 返回 false 可提前中断整个批处理任务
        }
        fmt.Printf("\r正在提取特征... 当前进度: %d / %d (%.1f%%)", current, total, float64(current)/float64(total)*100)
        return true
    })

    if err != nil {
        panic(err)
    }
    fmt.Printf("\n成功生成 %d 条 Embedding 数据！\n", len(embeddings))
}

---

多模态支持：图文向量化 (CLIP)

CLIPProvider 是 GoChat Embedding 模块中最特殊也是最强大的成员。它不仅能处理文本，还能直接对图片的像素进行理解，并将图文特征投射到相同的向量空间中。这意味着您可以用文本向量去直接计算与图片向量的余弦距离！

• 特性与优势：一键打通文本域与视觉域。内嵌基于 Go 原生 Image 库的图形处理器（Resize, Center Crop, Normalize）。
• 适用场景：构建“以文搜图”图库、视觉零样本分类系统。

CLIP 快速使用示例

package main

import (
    "context"
    "fmt"
    "os"
    "github.com/DotNetAge/gochat/pkg/embedding"
)

func main() {
    // 拉取 CLIP 模型
    provider, err := embedding.WithCLIP("clip-vit-base-patch32", "")
    if err != nil {
        panic(err)
    }

    ctx := context.Background()

    // 1. 文本编码 (以查询词为例)
    textVectors, _ := provider.Embed(ctx, []string{"A photo of a cute cat"})
    fmt.Printf("文本向量生成成功，维度: %d\n", len(textVectors[0]))

    // 2. 图像编码 (以数据库图源为例)
    imgData, _ := os.ReadFile("cute_cat.jpg")
    imageVectors, err := provider.EmbedImages(ctx, [][]byte{imgData})
    if err != nil {
        panic(err)
    }
    fmt.Printf("图片向量生成成功，维度: %d\n", len(imageVectors[0]))

    // 后续可以直接计算 textVectors[0] 与 imageVectors[0] 的相似度评分
}

---

扩展开发指南：如何接入自定义本地模型

随着 AI 社区的蓬勃发展，如果您在网上下载到了一个新的或微调过的 ONNX Embedding 模型，GoChat 为您敞开了无缝集成的大门。

标准实现流程

1. 实现 EmbeddingModel 接口：提供模型加载与生命周期管理（Run() 和 Close()）。
2. 将模型注入基础 Provider：使用 embedding.New(Config{}) 将您的自定义模型包裹进提供并发锁与内置 Tokenizer 的标准 LocalProvider。

代码骨架与集成范例

本示例展示了如何创建自定义模型驱动并将其注册到系统的高性能批处理架构中：

package main

import (
    "context"
    "fmt"
    "log"
    "github.com/DotNetAge/gochat/pkg/embedding"
)

// 1. 实现底层的 EmbeddingModel 接口
type CustomONNXModel struct {
    dimension int
    // onnxSession *onnxruntime.Session // 实际项目中持有真实的推理会话
}

func NewCustomONNXModel(modelPath string, dim int) (*CustomONNXModel, error) {
    // 在这里加载您的 .onnx 权重文件...
    return &CustomONNXModel{dimension: dim}, nil
}

// 核心推理逻辑：接管从 Tokenizer 产出的 input_ids 并返回最后的特征张量
func (m *CustomONNXModel) Run(inputs map[string]interface{}) (map[string]interface{}, error) {
    inputIDs, ok := inputs["input_ids"].([][]int64)
    if !ok {
        return nil, fmt.Errorf("无效的输入张量")
    }

    batchSize := len(inputIDs)

    // 执行 ONNX Runtime 推理...
    // 假设推理输出的矩阵为 embeddings 
    embeddings := make([][]float32, batchSize)
    for i := 0; i < batchSize; i++ {
        embeddings[i] = make([]float32, m.dimension) // Mock
    }

    // 必须返回包含 "last_hidden_state" 或 "embeddings" 键的 map
    return map[string]interface{}{
        "last_hidden_state": embeddings,
    }, nil
}

func (m *CustomONNXModel) Close() error {
    // 释放 CGo 或系统内存
    return nil
}

func main() {
    ctx := context.Background()

    // 2. 实例化自定义模型引擎
    model, _ := NewCustomONNXModel("path/to/custom_finetuned.onnx", 768)
    defer model.Close()

    // 3. 将自定义引擎包装为通用的 Provider
    provider, err := embedding.New(embedding.Config{
        Model:        model,  // 注入您的引擎
        Dimension:    768,    // 声明输出维度
        MaxBatchSize: 32,     // 声明最大张量批次限制
    })
    if err != nil {
        log.Fatalf("生成 Provider 失败: %v", err)
    }

    // 4. 成功集成！您可以直接将它喂给 BatchProcessor 享受所有高级特性了
    processor := embedding.NewBatchProcessor(provider, embedding.BatchOptions{
        MaxBatchSize:  16,
        MaxConcurrent: 2,
    })

    results, _ := processor.Process(ctx, []string{"这是我的自定义微调模型的威力！"})
    fmt.Printf("成功处理，提取到了 %d 维特征。\n", len(results[0]))
}

利用这一扩展机制，您可以轻松应对未来各种奇形怪状的新模型协议。