当前位置：首页 > article >正文

LlamaIndex中使用本地LLM和Embedding

article 2025/2/23 16:26:42

LlamaIndex默认会调用OpenAI的text-davinci-002模型对应的API，用于获得大模型输出，这种方式在很多情况下对国内用户不太方便，如果本地有大模型可以部署，可以按照以下方式在LlamaIndex中使用本地的LLM和Embedding（这里LLM使用chatglm2-6b，Embedding使用m3e-base）：

import torch
from transformers import AutoModel, AutoTokenizer
from llama_index.llms import HuggingFaceLLM
from llama_index import VectorStoreIndex, ServiceContext
from llama_index import LangchainEmbedding, ServiceContext
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index import Prompt, PromptHelper
from llama_index.node_parser import SimpleNodeParser
from llama_index.langchain_helpers.text_splitter import TextSplitter, TokenTextSplitter
from llama_index import set_global_service_context

# 需要使用GPU才能运行
device = 'cuda'

# 自定义输入大模型的prompt
TEMPLATE_STR = """我们在下面提供了上下文信息

{context_str}
根据此信息，请回答问题：{query_str}
"""
QA_TEMPLATE = Prompt(TEMPLATE_STR)

# 加载本地LLM，需提供本地LLM模型文件的路径
llm_tokenizer = AutoTokenizer.from_pretrained('/models/chatglm2-6b/', trust_remote_code=True, device=device)
llm_model = AutoModel.from_pretrained('/models/chatglm2-6b/', trust_remote_code=True, device=device)
chatglm2 = HuggingFaceLLM(model=llm_model, tokenizer=llm_tokenizer)

# 加载本地Embedding，需提供本地Embedding模型文件的路径
embed_tokenizer = AutoTokenizer.from_pretrained('/models/moka-ai/m3e-base/', trust_remote_code=True, device=device)
embed_model = LangchainEmbedding(
                HuggingFaceEmbeddings(model_name='/models/moka-ai/m3e-base/'), 
                tokenizer=embed_tokenizer)

node_parser = SimpleNodeParser(text_splitter=TokenTextSplitter(tokenizer=embed_tokenizer))
prompt_helper = PromptHelper(tokenizer=llm_tokenizer)
service_context = ServiceContext.from_defaults(
                llm=chatglm2, 
                prompt_helper=prompt_helper, 
                embed_model=embed_model, 
                node_parser=node_parser)
set_global_service_context(service_context)

documents = SimpleDirectoryReader('/path/to/your/files').load_data()
index = VectorStoreIndex.from_documents(documents, service_context=service_context)

 # 查询引擎
query_engine = index.as_query_engine(text_qa_template=QA_TEMPLATE)
# 聊天引擎
chat_eigine = index.as_chat_engine()

response = query_engine.query("your question")
print(response)