LlamaIndex 目前的支持只有 OpenAI 和 Ollama,通过 Ollama 可以跑很多开源模型,在本地跑的 LlamaIndex Demo 中,使用了 Llama2 作为后端模型,效果非常一般。之前注册了零一万物的开发者平台,目前还有很多体验的 credits,所以想要将零一万物的模型作为 LlamaIndex 的后端模型,应该是非常简单,因为零一万物的开发者文档里面就是使用了 OpenAI 的 Python SDK,就是改了一下 api_base 和api_key。翻了一下 llama_index 中 llms 下 openai 的实现,把代码拷贝过来,将一些模型的枚举值干掉,改成零一万物的就行
代码片段:
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.embeddings import resolve_embed_model
from llm.lingyi import LingYi
documents = SimpleDirectoryReader('data').load_data()
Settings.embed_model = resolve_embed_model("local:BAAI/bge-small-en-v1.5")
Settings.llm = LingYi(model='yi-34b-chat-0205', api_key="xxxxxxxxx",
api_base="https://api.lingyiwanwu.com/v1")
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()
response = query_engine.query("What did the author do growing up?")
print(response)
返回内容是:
The author grew up writing and programming. They wrote short stories and tried writing programs on an IBM 1401 in 9th grade. They also built and programmed microcomputers, which they found more engaging than the batch processing machines they had previously used. The author was fascinated by the idea of intelligent computers and was influenced by science fiction, leading them to study AI in college.
一个简单的 RAG 程序就完成了