LlamaIndex中的CustomLLM（本地加载模型）
输入“/”快速插入内容
LlamaIndex中的CustomLLM（本地加载模型）
LlamaIndex中的CustomLLM（本地加载模型）
飞书用户44435月24日修改
LlamaIndex中接口基本上调用的是OpenAI，如果想想调用自定义模型可以吗？答案当然是可以的。经过查找找到了自定义大语言模型的简单抽象基类class CustomLLM(LLM)。
一.CustomLLM(LLM)源码
只要子类必须实现 `__init__`、`_complete`、`_stream_complete` 和 `metadata` 方法即可。
代码块
class CustomLLM(LLM):
"""Simple abstract base class for custom LLMs.
Subclasses must implement the `__init__`, `_complete`,
`_stream_complete`, and `metadata` methods.
"""
@llm_chat_callback()
def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
chat_fn = completion_to_chat_decorator(self.complete)
return chat_fn(messages, **kwargs)
@llm_chat_callback()
def stream_chat(
self, messages: Sequence[ChatMessage], **kwargs: Any
) -> ChatResponseGen:
stream_chat_fn = stream_completion_to_chat_decorator(self.stream_complete)
return stream_chat_fn(messages, **kwargs)
@llm_chat_callback()
async def achat(
self,
messages: Sequence[ChatMessage],
**kwargs: Any,
) -> ChatResponse:
return self.chat(messages, **kwargs)
@llm_chat_callback()
async def astream_chat(
self,
messages: Sequence[ChatMessage],
**kwargs: Any,
) -> ChatResponseAsyncGen:
async def gen() -> ChatResponseAsyncGen:
for message in self.stream_chat(messages, **kwargs):
yield message
# NOTE: convert generator to async generator
return gen()
@llm_completion_callback()
async def acomplete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
return self.complete(prompt, **kwargs)
@llm_completion_callback()
async def astream_complete(
self, prompt: str, **kwargs: Any
) -> CompletionResponseAsyncGen:
async def gen() -> CompletionResponseAsyncGen:
for message in self.stream_complete(prompt, **kwargs):
yield message
# NOTE: convert generator to async generator
return gen()
@classmethod
def class_name(cls) -> str:
return "custom_llm"