第二章:langchain文本向量化(embed)搭建与详细教程-本地服务方式(下)
文章目录
- 前言
- 一、本地构建模型与服务端搭建代码
- 1、完整代码
- 2、结果示例
- 二、基于flask服务端构建langchain调用的向量方法代码
- 1、完整代码
- 2、结果示例
- 总结
前言
在上篇文章介绍了langchain源码embed方法与基于api key方式构建向量模型。然而,某些情况,我们调用的模型不存在支持框架能将其转为openai api key形式或不被支持。基于此,我们使用服务端与客户端的方法来构建一个能替换上一章内容的方法。我们构建本地构建一个向量模型,通过服务方式发送,再被langchain所接受形式。当然,这个方法也适用大语言模型实现形式,只需要发送服务接口与langchain调用保持一致即可。
一、本地构建模型与服务端搭建代码
我这边本地使用transformer方法来构建一个text embed模型,而服务使用flask方法来实现,我已在代码内做了解读,将不再介绍。
1、完整代码
from flask import Flask, request, jsonify # 导入Flask框架及其处理请求和响应的模块
from transformers import AutoTokenizer, AutoModel # 导入Hugging Face的AutoTokenizer和AutoModel用于加载预训练模型
import torch # PyTorch库,用于深度学习相关的操作
# 设置设备为GPU(如果有)或CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
def init_embed_model_huggingface(model_path, device=device):tokenizer = AutoTokenizer.from_pretrained(model_path) # 从指定路径加载预训练的分词器model = AutoModel.from_pretrained(model_path) # 从指定路径加载预训练的模型model.to(device) # 将模型移动到选定的设备(GPU/CPU)model.eval() # 设置模型为评估模式,关闭dropout等训练时才需要的功能return tokenizer, model
model_path = '/langchain/embed_model/bge-small-zh-v1.5' # 模型文件的本地路径
tokenizer, model = init_embed_model_huggingface(model_path) # 初始化分词器和模型
app = Flask(__name__) # 创建一个Flask应用实例
@app.route('/', methods=['POST']) # 定义路由和允许的方法,这里只接受POST请求
def get_embedding():# 检查请求头中的Authorization字段是否匹配特定的Bearer Tokenif request.headers.get('Authorization') != 'Bearer sk-JaJE4fXaktHAI2MB295FAd58004f7eBcE255B63CdD6F2':return jsonify({"error": "Unauthorized"}), 403 # 如果不匹配,则返回未授权错误data = request.json # 从请求中获取JSON格式的数据sentences = data.get('input') # 获取"data"键对应的值,即要嵌入的句子列表if not sentences:return jsonify({"error": "No input provided"}), 400 # 如果没有提供输入,则返回错误信息encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt') # 对输入文本进行编码# 将输入张量移动到与模型相同的设备上encoded_input = {key: tensor.to(device) for key, tensor in encoded_input.items()}with torch.no_grad(): # 禁用梯度计算,减少内存占用和加速计算model_output = model(**encoded_input) # 使用模型进行前向计算sentence_embeddings = model_output[0][:, 0] # 提取每个句子的第一个token的隐藏状态作为句子嵌入# 归一化并转换为列表格式sentence_embeddings = torch.nn.functional.normalize(sentence_embeddings, p=2, dim=1).tolist()# 构建响应内容,包含嵌入结果、使用的模型名称及对象类型response = {"data": [{"embedding": emb, "index": idx} for idx, emb in enumerate(sentence_embeddings)],"model": "BAAI/bge-small-zh-v1.5","object": "list"}return jsonify(response) # 返回JSON格式的响应
if __name__ == '__main__':app.run(host='0.0.0.0', port=55000) # 运行Flask应用,监听所有IP地址的55000端口
2、结果示例
启动成功后结果如下:
当然,上面构建模型可以使用transformer方式,也可以使用flagembed框架。无论使用什么方法来构建模型,只要能进行发送,被客户端提供的内容能给出结果即可。而想被langchain所调用,只需构建满足langchain接口即可,该接口或格式已在上一章进行了说明。所以总结,客户端的代码满足langchain格式即可,而服务端与客户端如何接口,就看你个人如何定义了。
二、基于flask服务端构建langchain调用的向量方法代码
在上面启动了服务端后,我们需要构建符合langchain可调用的text embed模型,而该模型构建需要什么接口,我已在上一章给出了解读,就是embed_documents与embed_query的2个方法,但也要记住异构的这2个方法哈(因langchain的Emebdding存在这个内容,这里不构建了)。
1、完整代码
import requests
from langchain_core.embeddings import Embeddingsclass EmbeddingsModel_flask(Embeddings):# 初始化方法,用于设置模型参数和API密钥def __init__(self, api_key=None, base_url=None, model="BAAI/bge-small-zh-v1.5", **kwargs):# 设置模型名称,默认为"BAAI/bge-small-zh-v1.5"self.model = model# 设置API的基础URL,如果未提供则从kwargs中获取self.base_url = base_url or kwargs.get("base_url")# 设置API密钥self.api_key = api_key# 获取嵌入向量的私有方法def _get_embeddings(self, texts):# 设置请求头,包括授权信息和内容类型headers = {'Authorization': f'Bearer {self.api_key}','Content-Type': 'application/json'}# 设置请求体,包含输入文本data = {"input": texts}# 发送POST请求到API,获取嵌入向量response = requests.post(self.base_url, json=data, headers=headers)# 检查响应状态码,如果是200则表示成功if response.status_code == 200:# 从响应中提取嵌入向量return [item['embedding'] for item in response.json()['data']]else:# 如果响应状态码不是200,则抛出异常raise Exception(f"Error: {response.text}")# 嵌入文档的方法def embed_documents(self, texts, chunk_size=None):# 调用私有方法获取嵌入向量return self._get_embeddings(texts)# 嵌入查询的方法def embed_query(self, text):# 调用私有方法获取嵌入向量embeddings = self._get_embeddings([text])# 返回第一个嵌入向量,如果没有则返回Nonereturn embeddings[0] if embeddings else None# 写一个本地加载的Demo
def embed_model_flask():api_url = "http://132.142.995.20:55000/" # 修改这里的URL以匹配服务器端点api_key_embeddings = "sk-JaJE4fXaktHAI2MB295FAd5804f7eBcE5B63CdD6F2"embeddings_model = EmbeddingsModel_flask(api_key=api_key_embeddings, base_url=api_url)return embeddings_modelif __name__ == "__main__":api_url = "http://132.142.995.20:55000/" # 修改这里的URL以匹配服务器端点api_key_embeddings = "sk-JaJEfXaktHAI2MB295FAd5004f7eBcE255B63CdD6F2"client = EmbeddingsModel_flask(api_key=api_key_embeddings, base_url=api_url)data = ["今天天气很好", "我们很开心,和大家一起出来玩"]embedding = client.embed_query(data[0]) # 单个查询print("Single query embedding:", embedding)embeddings = client.embed_documents(data) # 多文档查询print("Multiple documents embeddings:", embeddings)
2、结果示例
总结
当然,使用服务是一个可行方法,但至于用什么来做,可根据你的想法来。这里,你需要构建满足langchain接口即可。另外,该模式不仅使用text embed模型,也适用任何其它模型,如大语言模型。