RAG實戰6-如何在LlamaIndex中使用自己搭建的API

一蓑烟雨度平生發表於2024-03-14
AIIndexAPI

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應用時都載入一遍大模型權重了。

相關文章