devbox
喵喵爱编程
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

langchain自定义vllm chat model实现知识库RAG快速推理_langchain vllm

作者：喵喵爱编程 | 2024-08-23 07:20:25

赞

踩

langchain vllm

原本方法：继承langchain.llms.base import LLM类 实现一个CustomLLM

BaseLLM官方文档

我使用的是Qwen2模型，原本改写的模型：

  1. class QwenLM(LLM):
  2.     # 基于本地 Qwen 自定义 LLM 类
  3.     tokenizer : AutoTokenizer = None
  4.     model: AutoModelForCausalLM = None
  5.  
  6.     def __init__(self, model_path :str):
  7.         # model_path: Qwen 模型路径
  8.         # 从本地初始化模型
  9.         super().__init__()
  10.         print("正在从本地加载模型...")
  11.         model_dir = ''
  12.         self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
  13.         self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True).eval()
  14.         # Specify hyperparameters for generation
  15.         self.model.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True) # 可指定不同的生成长度、top_p等相关超参
  16.         print("完成本地模型的加载")
  17.  
  18.     def history_to_messages(self,history):
  19.         messages = [{'role': 'system', 'content': 'xxx'}]
  20.         for h in history:
  21.             messages.append({'role': 'user', 'content': h[0]})
  22.             messages.append({'role': 'assistant', 'content': h[1]})
  23.         return messages
  24.  
  25.     def model_chat(self,model, tokenizer, prompt, history=None):
  26.  
  27.         model_name_or_path = 'Qwen2'
  28.         
  29.         # Assume other models have chat method
  30.         if "Qwen2" not in model_name_or_path:
  31.             return model.chat(tokenizer, prompt, history=history)
  32.         else:
  33.             if history is None:
  34.                 history = []
  35.             else:
  36.                 history = copy.deepcopy(history)
  37.                     
  38.             messages = self.history_to_messages(history)
  39.             messages.append({'role': 'user', 'content': prompt})
  40.             text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
  41.             model_inputs = tokenizer([text], return_tensors="pt").to("cuda")
  42.             generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
  43.             generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
  44.             response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
  45.             history.append((prompt, response))
  46.             return response, history
  47.  
  48.     def _call(self, prompt : str, stop: Optional[List[str]] = None,
  49.                 run_manager: Optional[CallbackManagerForLLMRun] = None,
  50.                 **kwargs: Any):
  51.         # 重写调用函数
  52.         # response, history = self.model.chat(self.tokenizer, prompt , history=[])
  53.  
  54.         response,history = self.model_chat(self.model,self.tokenizer,prompt)
  55.         return response
  56.    
  57.     @property
  58.     def _llm_type(self) -> str:
  59.         return "QwenLM"

但是发现用这种情况去推理的时候每次都要load model和tokenizer，非常耗时，所以想到先把模型用vllm部署然后推理，能大幅提升模型推理的效率，故我改写并定义VLLMChatModel模型如下：

  1. class VLLMChatModel(LLM):
  2.     # 基于vLLM的自定义LLM类
  3.     api_client: OpenAI = None
  4.  
  5.     def __init__(self, api_base_url: str):
  6.         super().__init__()
  7.         self.api_client = OpenAI(api_key="EMPTY", base_url=api_base_url)
  8.  
  9.     def history_to_messages(self, history):
  10.         messages = [{'role': 'system', 'content': ''}] #system prompt自行补充
  11.         for h in history:
  12.             messages.append({'role': 'user', 'content': h[0]})
  13.             messages.append({'role': 'assistant', 'content': h[1]})
  14.         return messages
  15.  
  16.     def model_chat(self, prompt, history=None):
  17.         if history is None:
  18.             history = []
  19.  
  20.         messages = self.history_to_messages(history)
  21.         messages.append({'role': 'user', 'content': prompt})
  22.  
  23.         # 调用vLLM API
  24.         response = self.api_client.chat.completions.create(
  25.             model="Qwen2-7B-Instruct_2",
  26.             messages=messages,
  27.         )
  28.  
  29.         generated_text = response.choices[0].message.content
  30.         history.append((prompt, generated_text))
  31.         return generated_text, history
  32.  
  33.     def _call(self, prompt: str, stop: Optional[List[str]] = None,
  34.               run_manager: Optional[CallbackManagerForLLMRun] = None,
  35.               **kwargs: Any):
  36.         # 重写调用函数
  37.         response, history = self.model_chat(prompt)
  38.         return response
  39.     
  40.     @property
  41.     def _llm_type(self) -> str:
  42.         return "VLLMChatModel"

vllm启动的代码如下 ：

  1. export CUDA_VISIBLE_DEVICES=4,5,6,7
  2.  
  3. python -m vllm.entrypoints.openai.api_server \
  4.     --served-model-name Qwen2-7B-Instruct_2 \ #这个对应定义函数中的model=""
  5.     --model '' \  #Qwen2模型的local path
  6.     --gpu-memory-utilization 0.9 \
  7.     --tensor-parallel-size 4 \ #并行的卡数
  8.     --port 8002 \
  9.     --dtype float

调用定义好的VLLMChatModel:

  1. llm = VLLMChatModel(api_base_url="http://localhost:8002/v1")
  2.  
  3. # llm = QwenLM(model_path = "") #之前定义的QwenLM其实也可以用

这样之后就可以实现在langchain框架下用自定义的vllm chat model来实现高效率的推理了

完整项目：我实现的是一个用langchain框架结合知识库RAG然后用Qwen2-7b-Instruct模型推理的项目。

langchain框架结合知识库RAG框架代码参考：

https://blog.csdn.net/FL1623863129/article/details/137900683?spm=1001.2014.3001.5506icon-default.png?t=N7T8https://blog.csdn.net/FL1623863129/article/details/137900683?spm=1001.2014.3001.5506完整代码：
 

  1. import os
  2. os.environ['CUDA_VISIBLE_DEVICES'] = '4,5,6,7'
  3.  
  4. from langchain.document_loaders import UnstructuredFileLoader
  5. from langchain.document_loaders import UnstructuredMarkdownLoader
  6. from langchain_community.document_loaders.csv_loader import CSVLoader
  7. from langchain.text_splitter import RecursiveCharacterTextSplitter
  8. from langchain.vectorstores import Chroma
  9. from langchain.embeddings.huggingface import HuggingFaceEmbeddings
  10. from langchain.prompts import PromptTemplate
  11. from langchain.chains import RetrievalQA
  12. from tqdm import tqdm
  13. import os
  14. from langchain.llms.base import LLM
  15. from typing import Any, List, Optional
  16. from langchain.callbacks.manager import CallbackManagerForLLMRun
  17. from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
  18.  
  19. from openai import OpenAI
  20.  
  21. import copy
  22.  
  23.  
  24. class VLLMChatModel(LLM):
  25.     # 基于vLLM的自定义LLM类
  26.     api_client: OpenAI = None
  27.  
  28.     def __init__(self, api_base_url: str):
  29.         super().__init__()
  30.         self.api_client = OpenAI(api_key="EMPTY", base_url=api_base_url)
  31.  
  32.     def history_to_messages(self, history):
  33.         messages = [{'role': 'system', 'content': ''}] #system prompt自行补充
  34.         for h in history:
  35.             messages.append({'role': 'user', 'content': h[0]})
  36.             messages.append({'role': 'assistant', 'content': h[1]})
  37.         return messages
  38.  
  39.     def model_chat(self, prompt, history=None):
  40.         if history is None:
  41.             history = []
  42.  
  43.         messages = self.history_to_messages(history)
  44.         messages.append({'role': 'user', 'content': prompt})
  45.  
  46.         # 调用vLLM API
  47.         response = self.api_client.chat.completions.create(
  48.             model="Qwen2-7B-Instruct_2",
  49.             messages=messages,
  50.         )
  51.  
  52.         generated_text = response.choices[0].message.content
  53.         history.append((prompt, generated_text))
  54.         return generated_text, history
  55.  
  56.     def _call(self, prompt: str, stop: Optional[List[str]] = None,
  57.               run_manager: Optional[CallbackManagerForLLMRun] = None,
  58.               **kwargs: Any):
  59.         # 重写调用函数
  60.         response, history = self.model_chat(prompt)
  61.         return response
  62.     
  63.     @property
  64.     def _llm_type(self) -> str:
  65.         return "VLLMChatModel"
  66.     
  67.  
  68. class QwenLM(LLM):
  69.     # 基于本地 Qwen 自定义 LLM 类
  70.     tokenizer : AutoTokenizer = None
  71.     model: AutoModelForCausalLM = None
  72.  
  73.     def __init__(self, model_path :str):
  74.         # model_path: Qwen 模型路径
  75.         # 从本地初始化模型
  76.         super().__init__()
  77.         print("正在从本地加载模型...")
  78.         model_dir = '' #Qwen2模型的local path
  79.         self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
  80.         self.model = AutoModelForCausalLM.from_pretrained(model_dir, device_map="auto", trust_remote_code=True).eval()
  81.         # Specify hyperparameters for generation
  82.         self.model.generation_config = GenerationConfig.from_pretrained(model_dir, trust_remote_code=True) # 可指定不同的生成长度、top_p等相关超参
  83.         print("完成本地模型的加载")
  84.  
  85.     def history_to_messages(self,history):
  86.         messages = [{'role': 'system', 'content': ''}] #system prompt自行补充
  87.         for h in history:
  88.             messages.append({'role': 'user', 'content': h[0]})
  89.             messages.append({'role': 'assistant', 'content': h[1]})
  90.         return messages
  91.  
  92.     def model_chat(self,model, tokenizer, prompt, history=None):
  93.  
  94.         model_name_or_path = 'Qwen2'
  95.         
  96.         # Assume other models have chat method
  97.         if "Qwen2" not in model_name_or_path:
  98.             return model.chat(tokenizer, prompt, history=history)
  99.         else:
  100.             if history is None:
  101.                 history = []
  102.             else:
  103.                 history = copy.deepcopy(history)
  104.                     
  105.             messages = self.history_to_messages(history)
  106.             messages.append({'role': 'user', 'content': prompt})
  107.             text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
  108.             model_inputs = tokenizer([text], return_tensors="pt").to("cuda")
  109.             generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
  110.             generated_ids = [output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
  111.             response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
  112.             history.append((prompt, response))
  113.             return response, history
  114.  
  115.     def _call(self, prompt : str, stop: Optional[List[str]] = None,
  116.                 run_manager: Optional[CallbackManagerForLLMRun] = None,
  117.                 **kwargs: Any):
  118.         # 重写调用函数
  119.         # response, history = self.model.chat(self.tokenizer, prompt , history=[])
  120.  
  121.         response,history = self.model_chat(self.model,self.tokenizer,prompt)
  122.         return response
  123.     
  124.         
  125.     @property
  126.     def _llm_type(self) -> str:
  127.         return "QwenLM"
  128.  
  129.  
  130. # 获取文件路径函数
  131. def get_files(dir_path):
  132.     # args：dir_path，目标文件夹路径
  133.     file_list = []
  134.     for filepath, dirnames, filenames in os.walk(dir_path):
  135.         # os.walk 函数将递归遍历指定文件夹
  136.         for filename in filenames:
  137.             # 通过后缀名判断文件类型是否满足要求
  138.             if filename.endswith(".md"):
  139.                 # 如果满足要求，将其绝对路径加入到结果列表
  140.                 file_list.append(os.path.join(filepath, filename))
  141.             elif filename.endswith(".txt"):
  142.                 file_list.append(os.path.join(filepath, filename))
  143.             elif filename.endswith(".csv"):
  144.                 file_list.append(os.path.join(filepath, filename))
  145.     return file_list
  146.  
  147. # 加载文件函数
  148. def get_text(dir_path):
  149.     # args：dir_path，目标文件夹路径
  150.     # 首先调用上文定义的函数得到目标文件路径列表
  151.     file_lst = get_files(dir_path)
  152.     # docs 存放加载之后的纯文本对象
  153.     docs = []
  154.     # 遍历所有目标文件
  155.     for one_file in tqdm(file_lst):
  156.         file_type = one_file.split('.')[-1]
  157.         if file_type == 'md':
  158.             loader = UnstructuredMarkdownLoader(one_file)
  159.         elif file_type == 'txt':
  160.             loader = UnstructuredFileLoader(one_file)
  161.         elif file_type == 'csv':
  162.             loader = CSVLoader(one_file)
  163.         else:
  164.             # 如果是不符合条件的文件，直接跳过
  165.             continue
  166.         docs.extend(loader.load())
  167.     return docs
  168.  
  169. # 目标文件夹
  170. # tar_dir = [
  171. #     "/root/autodl-tmp/qwen",
  172. #     "/root/autodl-tmp/Qwen",
  173. # ]
  174.  
  175.  
  176. #构建向量数据库并持久化到磁盘上
  177. def generate_vectordb():
  178.     tar_dir = ["/raid/home/specter/enhance_llm/qa_bot/hbv_data"]
  179.  
  180.     # 加载目标文件
  181.     docs = []
  182.     for dir_path in tar_dir:
  183.         docs.extend(get_text(dir_path))
  184.  
  185.     # 对文本进行分块
  186.     text_splitter = RecursiveCharacterTextSplitter(
  187.         chunk_size=500, chunk_overlap=150)
  188.     split_docs = text_splitter.split_documents(docs)
  189.  
  190.     # 加载开源词向量模型
  191.     embeddings = HuggingFaceEmbeddings(model_name="/raid/home/specter/enhance_llm/qa_bot/embedding_model")
  192.  
  193.     # 构建向量数据库
  194.     # 定义持久化路径
  195.     persist_directory = '/raid/home/specter/enhance_llm/qa_bot/data_base/vector_db/chroma'
  196.     # 加载数据库
  197.     vectordb = Chroma.from_documents(
  198.         documents=split_docs,
  199.         embedding=embeddings,
  200.         persist_directory=persist_directory  # 允许我们将persist_directory目录保存到磁盘上
  201.     )
  202.     # 将加载的向量数据库持久化到磁盘上
  203.     vectordb.persist()
  204.  
  205.  
  206.  
  207. def decode_2(user_description):
  208.  
  209.     #构建向量数据库并持久化到磁盘上
  210.     generate_vectordb()
  211.         
  212.     # 加载开源词向量模型
  213.     embeddings = HuggingFaceEmbeddings(model_name="/raid/home/specter/enhance_llm/qa_bot/embedding_model")
  214.  
  215.     # 向量数据库持久化路径
  216.     persist_directory = '/raid/home/specter/enhance_llm/qa_bot/data_base/vector_db/chroma'
  217.  
  218.     # 加载数据库
  219.     vectordb = Chroma(
  220.         persist_directory=persist_directory, 
  221.         embedding_function=embeddings
  222.     )
  223.  
  224.     llm = VLLMChatModel(api_base_url="http://localhost:8002/v1")
  225.  
  226.     # llm = QwenLM(model_path = "") #loacl path自行补充
  227.  
  228.     # 我们所构造的 Prompt 模板
  229.     template = """使用以下上下文来回答最后的问题。如果你不知道答案，就说你不知道，不要试图编造答案。尽量使答案简明扼要。总是在回答的最后说“谢谢你的提问！”。
  230.     {context}
  231.     问题: {question}
  232.     有用的回答:"""
  233.  
  234.     # 调用 LangChain 的方法来实例化一个 Template 对象，该对象包含了 context 和 question 两个变量，在实际调用时，这两个变量会被检索到的文档片段和用户提问填充
  235.     QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context","question"],template=template)
  236.  
  237.     qa_chain = RetrievalQA.from_chain_type(llm,retriever=vectordb.as_retriever(),return_source_documents=True,chain_type_kwargs={"prompt":QA_CHAIN_PROMPT})
  238.  
  239.     question = user_description
  240.  
  241.     result = qa_chain({"query": question})
  242.     print("检索问答链回答 question 的结果：")
  243.     print(result["result"])
  244.  
  245.     return result["result"]
  246.  
  247.     # # 仅 LLM 回答效果
  248.     # result_2 = llm(question)
  249.     # print("大模型回答 question 的结果：")
  250.     # print(result_2)
  251.  
  252.  
  253. if __name__ =='__main__':
  254.     decode_2('') #补充user_description内容
声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop博客】
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号