RAG实战6-如何在LlamaIndex中使用自己搭建的API
RAG实战6-如何在LlamaIndex使用自己搭建的大模型API
在搭建一个大模型API服务中，我们介绍了如何使用SWIFT框架搭建一个大模型API服务。在RAG实战1-5中，我们一直使用的是本地加载大模型的方式来调用大模型，本文将介绍如何在LlamaIndex中使用自己搭建的大模型API。
LlamaIndex支持部分厂商的API配置，如OpenAI，但我们想使用的是自己在服务器上搭建的API服务，这个时候需要我们定制一个LLM类，代码如下：
from typing import Any
from llama_index.core import PromptTemplate, Settings, StorageContext, load_index_from_storage
from llama_index.core.base.llms.types import LLMMetadata, CompletionResponse, CompletionResponseGen
from llama_index.core.llms import CustomLLM
from llama_index.core.llms.callbacks import llm_completion_callback
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from swift.llm import get_model_list_client, XRequestConfig, inference_client
# API
model_list = get_model_list_client()
model_type = model_list.data[0].id
print(f'API model_type: {model_type}')
request_config = XRequestConfig(seed=42)
# 定制自己的LLM类
class BianCangLLM(CustomLLM):
context_window: int = 4096
num_output: int = 2048
model_name: str = "BianCang"
@property
def metadata(self) -> LLMMetadata:
"""Get LLM metadata."""
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:
resp = inference_client(model_type, prompt, [], request_config=request_config)
return CompletionResponse(text=resp.choices[0].message.content)
@llm_completion_callback()
def stream_complete(
self, prompt: str, **kwargs: Any
) -> CompletionResponseGen:
resp = inference_client(model_type, prompt, [], request_config=request_config)
response = ""
for token in resp.choices[0].message.content:
response += token
yield CompletionResponse(text=response, delta=token)
# 定义system prompt
SYSTEM_PROMPT = """你是一个医疗人工智能助手。"""
query_wrapper_prompt = PromptTemplate(
"[INST]<<SYS>>\n" + SYSTEM_PROMPT + "<</SYS>>\n\n{query_str}[/INST] "
)
# 定义qa prompt
qa_prompt_tmpl_str = (
"上下文信息如下。\n"
"---------------------\n"
"{context_str}\n"
"---------------------\n"
"请根据上下文信息而不是先验知识来回答以下的查询。"
"作为一个医疗人工智能助手，你的回答要尽可能严谨。\n"
"Query: {query_str}\n"
"Answer: "
)
qa_prompt_tmpl = PromptTemplate(qa_prompt_tmpl_str)
# 使用自定义的LLM API
Settings.llm = BianCangLLM()
# 使用llama-index-embeddings-huggingface构建本地embedding模型
Settings.embed_model = HuggingFaceEmbedding(
model_name="E:\\LLMs\\bge-base-zh-v1.5"
)
# 从存储文件中读取embedding向量和向量索引
storage_context = StorageContext.from_defaults(persist_dir="doc_emb")
index = load_index_from_storage(storage_context)
# 构建查询引擎
query_engine = index.as_query_engine(similarity_top_k=5)
# 更新查询引擎中的prompt template
query_engine.update_prompts(
{"response_synthesizer:text_qa_template": qa_prompt_tmpl}
)
# 查询获得答案
response = query_engine.query("不耐疲劳，口燥、咽干可能是哪些证候？")
print(response)
代码的核心是实现BianCangLLM类，该类继承自LlamaIndex的CustomLLM类。我们需要重写父类中的def metadata(self) -> LLMMetadata、def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse、def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:。其中，metadata负责定义大模型的一些参数属性；complete负责调用大模型API服务并直接返回响应；stream_complete负责调用大模型API服务并以流式输出的形式返回响应。
运行代码，同样可以得到类似于之前的效果：
根据提供的上下文信息，口燥、咽干的症状可能与以下几个中医证候相关：
1. 津液不足证（4.6.1.1）：由于津液生成不足或体内燥热，可能导致口眼喉鼻干燥。
2. 津亏热结证（4.6.3.2）：津液亏虚加上热邪内结，也会出现口燥咽干的表现。
3. 津液亏涸证（4.6.1.2）：津液亏损严重时，口干、唇裂、鼻燥、舌燥是其特征，可能包括咽干。
4. 燥干清窍证（3.6.3.2）：气候干燥导致的津液耗损，会引起口鼻咽喉干燥。
5. 津伤化燥证（6.3.1）：燥热内蕴或内热化燥可能引起口干舌燥，伴有多尿、消瘦等症状。
因此，这些证候都可能与不耐疲劳和口燥咽干的临床表现相关，但具体诊断需要结合其他症状和中医辨证原则。建议患者就诊中医师以获取专业诊断。
好处是，我们不需要每次启动RAG应用时都加载一遍大模型权重了。
