从零学习大模型——使用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)
输出了很多向量构成的数组说明运行正常。