从零学习大模型——使用GLM-4-9B-Chat + BGE-M3 + langchain + chroma建立的本地RAG应用(三)——将BGE-M3接入langchain

BGE-M3 是第一个具有多功能、多语言和多粒度特性的文本检索模型。

  • 多功能:可以同时执行三种检索功能:单向量检索、多向量检索和稀疏检索。
  • 多语言:支持100多种工作语言。
  • 多粒度:它能够处理不同粒度的输入,从短句子到长达8192个词汇的长文档。

为了构建RAG应用,我们需要用到向量数据库和embedding模型进行数据集的向量化存储和后续用到的相似度检索。

本文介绍如何将BGE-M3接入LangChain。

下载BGE-M3

运行以下命令:

git lfs install

git clone https://www.modelscope.cn/Xorbits/bge-m3.git

注意:安装时为了方便后续统一管理,执行克隆命令时要先将目录切换到/root/autodl-tmp

embedding模型——BGE-M3的搭建(以算力云平台为例)

安装lfs时若遇到问题可参考上方链接

与LLM模型一样,Embedding模型也需要继承langchain.embeddings.base的一个子类并实现相关函数。

/root/autodl-tmp目录下新建textEmbeddings.py并输入以下语句:

from transformers import AutoTokenizer, AutoModel
import torch
from langchain.embeddings.base import Embeddings

class BGEM3Embeddings(Embeddings):
    def __init__(self, model, tokenizer):
        self.model = model
        self.tokenizer = tokenizer

    def embed_text(self, text: str) -> list:
        inputs = self.tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = self.model(**inputs)
        embeddings = outputs.last_hidden_state.mean(dim=1).squeeze().tolist()
        return embeddings

    def embed_documents(self, texts: list) -> list:
        embeddings = [self.embed_text(text) for text in texts]
        return embeddings

    def embed_query(self, text: str) -> list:
        return self.embed_text(text)

同样我们再新建一个测试文件测试是否接入成功

from transformers import AutoTokenizer, AutoModel
import torch
from testEmbeddings import BGEM3Embeddings

model_name = "bge-m3" 
save_directory = "/root/autodl-tmp/bge-m3"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
embedding = BGEM3Embeddings(model, tokenizer)

# 生成文本嵌入
text = "这是一个示例文本"
embedding = embedding.embed_text(text)
print(embedding)

# 生成多个文档的嵌入
documents = ["文档1", "文档2", "文档3"]
document_embeddings = embedding.embed_documents(documents)
print(document_embeddings)

# 生成查询的嵌入
query = "这是一个查询文本"
query_embedding = embedding.embed_query(query)
print(query_embedding)

输出了很多向量构成的数组说明运行正常。

posted @ 2024-07-10 14:38  YTARO  阅读(368)  评论(0编辑  收藏  举报