1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
from langchain.document_loaders import DirectoryLoader, unstructured
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma, faiss
from langchain.llms import ChatGLM
from langchain.chains import RetrievalQA
import gradio as gr
import time
import os
embedding_model_dict = {
"ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
"ernie-base": "nghuyong/ernie-3.0-base-zh",
"text2vec": "GanymedeNil/text2vec-large-chinese",
"text2vec2": "uer/sbert-base-chinese-nli",
"text2vec3": "shibing624/text2vec-base-chinese",
}
def load_documents(directory="books"):
"""
Load documents
"""
loader = DirectoryLoader(directory)
documents = loader.load()
# text: abcdefg, chunk_size:3, chunk_overlap: 1. result：abc cde efg
text_spliter = CharacterTextSplitter(chunk_size=256, chunk_overlap=0)
split_docs = text_spliter.split_documents(documents)
return split_docs
def load_embedding_model(model_name="ernie-tiny"):
"""
Load embedding model
"""
encode_kwargs = {"normalize_embeddings": False}
model_kwargs = {"device": "cuda:0"}
return HuggingFaceEmbeddings(
model_name=embedding_model_dict[model_name], # 这里如果不能在线下载，可以使用本地下载好的文件
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs
)
def store_chroma(docs, embeddings, persist_directory="VectorStore"):
db = Chroma.from_documents(docs, embeddings, persist_directory=persist_directory)
db.persist()
return db
# Chatbot demo with multimodal input (text, markdown, LaTeX, code blocks, image, audio, & video).
# Plus shows support for streaming text.
def print_like_dislike(x: gr.LikeData):
print(x.index, x.value, x.liked)
def add_text(history, text):
history = history + [(text, None)]
return history, gr.Textbox(value="", interactive=False)
def add_file(history, file):
# 取到文件夹（临时文件夹）
directory = os.path.dirname(file.name)
documents = load_documents(directory)
store_chroma(documents, embeddings)
history = history + [((file.name,), None)]
return history
def bot(history):
message = history[-1][0]
if isinstance(message, tuple):
response = "文件上传成功"
else:
response = qa.run(message)
history[-1][1] = ""
for character in response:
history[-1][1] += character
time.sleep(0.05)
yield history
if __name__ == "__main__":
embeddings = load_embedding_model("text2vec3")
if not os.path.exists("VectorStore"):
documents = load_documents()
db = store_chroma(documents, embeddings)
else:
db = Chroma(persist_directory="VectorStore", embedding_function=embeddings)
# 需要根据ChatGLM项目本地启动大模型服务在8000端口
llm = ChatGLM(
endpoint="http://127.0.0.1:8000",
max_token=80000,
top_p=0.9
)
retriever = db.as_retriever()
qa = RetrievalQA.from_chain_type(
llm=llm,
chain_type='stuff',
retriever=retriever
)
with gr.Blocks() as demo:
chatbot = gr.Chatbot(
[],
elem_id="chatbot",
bubble_full_width=False,
avatar_images=(None, (os.path.join(os.path.dirname(__file__), "avatar.png"))),
)
with gr.Row():
txt = gr.Textbox(
scale=4,
show_label=False,
placeholder="Enter text and press enter, or upload an image",
container=False,
)
btn = gr.UploadButton("📁", file_types=["txt"])
txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
bot, chatbot, chatbot, api_name="bot_response"
)
txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False).then(
bot, chatbot, chatbot
)
chatbot.like(print_like_dislike, None, None)
demo.queue()
demo.launch()
|
