LlamaIndex中的CustomLLM（在线加载模型）

一.使用 Flask 将模型封装为 REST 接口

主要是将 complete()和 stream_complete()方法封装为 REST 接口，如下所示：

from flask import Flask, request, jsonify
from transformers import AutoTokenizer, AutoModelForCausalLM
app = Flask(__name__)class QwenModel:def __init__(self, pretrained_model_name_or_path):# self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, device_map="cpu", trust_remote_code=True)  # CPU方式加载模型# self.model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, device_map="cpu", trust_remote_code=True)  # CPU方式加载模型self.tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, device_map="cuda", trust_remote_code=True)  # GPU方式加载模型self.model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, device_map="cuda", trust_remote_code=True)  # GPU方式加载模型self.model = self.model.float()def generate_completion(self, prompt):# inputs = self.tokenizer.encode(prompt, return_tensors="pt")  # CPU方式加载模型inputs = self.tokenizer.encode(prompt, return_tensors="pt").cuda()  # GPU方式加载模型outputs = self.model.generate(inputs, max_length=128)response = self.tokenizer.decode(outputs[0])return responsepretrained_model_name_or_path = r'L:\20230713_HuggingFaceModel\20230925_Qwen\Qwen-1_8B'
qwen_model = QwenModel(pretrained_model_name_or_path)@app.route('/complete', methods=['POST'])
def complete():data = request.get_json()prompt = data.get('prompt', '')result = qwen_model.generate_completion(prompt)return jsonify({'text': result})@app.route('/stream_complete', methods=['POST'])
def stream_complete():data = request.get_json()prompt = data.get('prompt', '')result = list(qwen_model.generate_completion(prompt))return jsonify(result)if __name__ == "__main__":app.run(debug=False, port=5050, host='0.0.0.0')

二.通过 requests.post 请求方式调用接口

主要是通过 requests.post 请求方式来实现 complete()和 stream_complete()函数，如下所示：

from typing import Anyimport requests
from llama_index import ServiceContext, SimpleDirectoryReader, SummaryIndex
from llama_index.llms import CustomLLM, CompletionResponse, CompletionResponseGen, LLMMetadata
from llama_index.llms.base import llm_completion_callbackclass QwenCustomLLM(CustomLLM):context_window: int = 8192num_output: int = 128model_name: str = "Qwen-1_8B"base_url: str = "http://127.0.0.1:5050"tokenizer: object = Nonemodel: object = Nonedef __init__(self):super().__init__()@propertydef metadata(self) -> LLMMetadata:return LLMMetadata(context_window=self.context_window,num_output=self.num_output,model_name=self.model_name,)@llm_completion_callback()def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:data = {'prompt': prompt}response = requests.post(f'{self.base_url}/complete', json=data)result = response.json()return CompletionResponse(text=result['text'])@llm_completion_callback()def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:data = {'prompt': prompt}response = requests.post(f'{self.base_url}/stream_complete', json=data)result = response.json()for token in result:yield CompletionResponse(text=token, delta=token)if __name__ == "__main__":llm = QwenCustomLLM()# 方式1：本地加载模型方式进行调用service_context = ServiceContext.from_defaults(llm=llm, embed_model="local:L:/20230713_HuggingFaceModel/BAAI_bge-large-zh")# TODO 方式2：调整embed_model为在线模型# 思路：继承BaseEmbedding类，将加载Embedding模型部分封装为REST接口。可重点参考下OpenAIEmbedding类的实现。# service_context = ServiceContext.from_defaults(llm=llm, embed_model=BgeLargeZhEmbedding())documents = SimpleDirectoryReader("./data").load_data()index = SummaryIndex.from_documents(documents, service_context=service_context)query_engine = index.as_query_engine()# 你能够像以前一样使用llm.complete和llm.stream_completeresponse_complete = llm.complete("您好")print(response_complete)response_stream_complete = list(llm.stream_complete("您好"))print(response_stream_complete)response = query_engine.query("花未眠")print(response)

上述代码在加载 Embedding 模型的时候还是从本地加载的，这部分也是可通过 REST 接口方式调用的。基本思路：继承 BaseEmbedding 类，将加载 Embedding 模型部分封装为 REST 接口。可重点参考下 OpenAIEmbedding 类的实现。