llama-index 结合chatglm3-6B 利用RAG 基于文档智能问答

news2024/12/23 17:50:02

简介

llamaindex结合chatglm3使用

import os
import torch
from llama_index.core import VectorStoreIndex, ServiceContext
from llama_index.core.callbacks import CallbackManager
from llama_index.core.llms.callbacks import llm_completion_callback
from llama_index.core.prompts.prompts import SimpleInputPrompt
from llama_index.legacy.embeddings import HuggingFaceEmbedding
from llama_index.legacy.llms import HuggingFaceLLM
from llama_index.core import ServiceContext

os.environ["CUDA_VISIBLE_DEVICES"] = "0"


# from llama_index.core import VectorStoreIndex, SummaryIndex, SimpleDirectoryReader, ServiceContext, LLMPredictor
from llama_index.legacy.llms import (HuggingFaceLLM, CustomLLM, CompletionResponse, CompletionResponseGen, LLMMetadata)

from typing import Optional, List, Mapping, Any
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoModel
from transformers.generation import GenerationConfig

from modelscope import snapshot_download

m3e 向量编码模型

from llama_index.legacy.embeddings import HuggingFaceEmbedding

#模型下载
# load embedding
emb_path = snapshot_download('jieshenai/m3e-base')
embedding_model = HuggingFaceEmbedding(emb_path)

自定义模型

model_name = "chatglm3-6b"
model_path = snapshot_download('ZhipuAI/chatglm3-6b')

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModel.from_pretrained(model_path, trust_remote_code=True).half().cuda()
model = model.eval()

# set context window size
context_window = 2048
# set number of output tokens
num_output = 256


class ChatGML(CustomLLM):
    @property
    def metadata(self) -> LLMMetadata:
        """Get LLM metadata."""
        return LLMMetadata(
            context_window=context_window,
            num_output=num_output,
            model_name=model_name,
        )

    @llm_completion_callback()
    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
        prompt_length = len(prompt)

        # only return newly generated tokens
        text,_ = model.chat(tokenizer, prompt, history=[])
        return CompletionResponse(text=text)

    @llm_completion_callback()
    def stream_complete(
        self, prompt: str, **kwargs: Any
    ) -> CompletionResponseGen:
        raise NotImplementedError()

llm_model = ChatGML()

from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader('data').load_data()

service_context = ServiceContext.from_defaults(
						llm=llm_model, embed_model=embedding_model)

在data文件夹下，放txt文件
data/1.txt:

在遥远的东方，有一个名为“梦溪”的小村庄。这里山水相依，风景如画，村民们过着宁静而祥和的生活。
梦溪村有个少年，名叫小风。他聪明伶俐，心地善良，深受村民们的喜爱。小风有个特别的爱好，那就是喜欢在夜晚仰望星空，幻想着天上的故事。
有一天，小风在村边的小溪旁捡到一颗奇特的石头，它闪烁着幽幽的蓝光，仿佛蕴藏着无尽的秘密。小风好奇地将石头带回家，放在床头，每晚都伴着它的光芒入睡。
渐渐地，小风发现自己开始做起了一个奇怪的梦。在梦中，他来到了一个神秘的世界，那里有着奇异的生物和壮丽的景色。他遇到了一个美丽的少女，她告诉他，这颗石头其实是一个通往另一个世界的门户，而她就是那个世界的守护者。
少女告诉小风，她的世界正面临着巨大的危机，需要小风的帮助。小风虽然有些害怕，但想到自己或许能成为一个英雄，便鼓起勇气答应了她的请求。
从此，小风每晚都进入那个神秘的世界，与少女一起冒险，解决各种难题。他学会了如何与奇异的生物沟通，如何运用智慧战胜困难。他的勇气和善良感染了那个世界的居民们，他们纷纷加入小风的行列，共同抵抗邪恶势力。
经过一段艰苦的奋战，小风终于帮助少女解决了危机，恢复了那个世界的和平。少女感激不已，送给小风一件神奇的礼物——一把能够召唤风的力量的魔法扇。
当小风醒来时，他发现自己手中真的握着那把魔法扇。他走到村边的小溪旁，轻轻一挥扇子，顿时一阵清风吹过，带来了花香和鸟鸣。
从此，小风成了梦溪村的小英雄。他用魔法扇帮助村民们解决困难，保护村庄的安宁。而那颗奇特的石头，也永远留在了他的床头，成为他勇敢和善良的见证。

# create index
index = VectorStoreIndex.from_documents(documents, service_context=service_context)

# query engine
query_engine = index.as_query_engine()

# query
response = query_engine.query("少女感激不已，送给小风一件神奇的礼物是什么？")
print(response)