Spaces:
Sleeping
Sleeping
Commit
·
24680b6
1
Parent(s):
2d56ea0
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,116 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#这个版本有个问题,如果在运行状况下,增删文件,不会重新装载文件并构建向量数据库!
|
2 |
+
#最后只能够通过添加按钮的方式,手动解决这个问题!!!
|
3 |
+
#尝试设置一个判断变量file_uploaded
|
4 |
+
|
5 |
+
import streamlit as st
|
6 |
+
from llama_index import VectorStoreIndex, SimpleDirectoryReader
|
7 |
+
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
|
8 |
+
from llama_index import LangchainEmbedding, ServiceContext
|
9 |
+
from llama_index import StorageContext, load_index_from_storage
|
10 |
+
from llama_index import LLMPredictor
|
11 |
+
#from transformers import HuggingFaceHub
|
12 |
+
from langchain import HuggingFaceHub
|
13 |
+
from streamlit.components.v1 import html
|
14 |
+
from pathlib import Path
|
15 |
+
from time import sleep
|
16 |
+
import random
|
17 |
+
import string
|
18 |
+
|
19 |
+
import os
|
20 |
+
from dotenv import load_dotenv
|
21 |
+
load_dotenv()
|
22 |
+
|
23 |
+
import timeit
|
24 |
+
|
25 |
+
st.set_page_config(page_title="Open AI Doc-Chat Assistant", layout="wide")
|
26 |
+
st.subheader("Open AI Doc-Chat Assistant: Life Enhancing with AI!")
|
27 |
+
|
28 |
+
css_file = "main.css"
|
29 |
+
with open(css_file) as f:
|
30 |
+
st.markdown("<style>{}</style>".format(f.read()), unsafe_allow_html=True)
|
31 |
+
|
32 |
+
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
|
33 |
+
|
34 |
+
def generate_random_string(length):
|
35 |
+
letters = string.ascii_lowercase
|
36 |
+
return ''.join(random.choice(letters) for i in range(length))
|
37 |
+
|
38 |
+
if "query_engine" not in st.session_state:
|
39 |
+
st.session_state.query_engine = None
|
40 |
+
|
41 |
+
if "file_uploaded" not in st.session_state:
|
42 |
+
st.session_state.file_uploaded = False
|
43 |
+
|
44 |
+
with st.sidebar:
|
45 |
+
st.subheader("Upload your pdf Documents Here: ")
|
46 |
+
pdf_files = st.file_uploader("Choose your PDF Files and Press OK", type=['pdf'], accept_multiple_files=True)
|
47 |
+
if not pdf_files:
|
48 |
+
st.warning("请上传文档文件")
|
49 |
+
st.stop()
|
50 |
+
else:
|
51 |
+
st.session_state.file_uploaded = True
|
52 |
+
|
53 |
+
if st.session_state.file_uploaded:
|
54 |
+
uploadedfile_path=generate_random_string(20)
|
55 |
+
os.makedirs(uploadedfile_path)
|
56 |
+
for pdf_file in pdf_files:
|
57 |
+
file_path = os.path.join(uploadedfile_path, pdf_file.name)
|
58 |
+
with open(file_path, 'wb') as f:
|
59 |
+
f.write(pdf_file.read())
|
60 |
+
st.success(f"File '{pdf_file.name}' saved successfully.")
|
61 |
+
# if st.button('Process for QA'):
|
62 |
+
try:
|
63 |
+
start_1 = timeit.default_timer() # Start timer
|
64 |
+
st.write(f"QA文档加载开始:{start_1}")
|
65 |
+
documents = SimpleDirectoryReader(uploadedfile_path).load_data()
|
66 |
+
end_1 = timeit.default_timer() # Start timer
|
67 |
+
st.write(f"QA文档加载结束:{end_1}")
|
68 |
+
st.write(f"QA文档加载耗时:{end_1 - start_1}")
|
69 |
+
except Exception as e:
|
70 |
+
print("文档加载出现问题/Waiting for path creation.")
|
71 |
+
st.warning("文档加载出现问题/Waiting for path creation.")
|
72 |
+
st.stop()
|
73 |
+
start_2 = timeit.default_timer() # Start timer
|
74 |
+
st.write(f"向量模型加载开始:{start_2}")
|
75 |
+
if "embed_model" not in st.session_state:
|
76 |
+
st.session_state.embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2'))
|
77 |
+
end_2 = timeit.default_timer() # Start timer
|
78 |
+
st.write(f"向量模型加载结束:{end_2}")
|
79 |
+
st.write(f"向量模型加载耗时:{end_2 - start_2}")
|
80 |
+
if "llm_predictor" not in st.session_state:
|
81 |
+
st.session_state.llm_predictor = LLMPredictor(HuggingFaceHub(repo_id="HuggingFaceH4/starchat-beta", model_kwargs={"min_length":100, "max_new_tokens":1024, "do_sample":True, "temperature":0.2,"top_k":50, "top_p":0.95, "eos_token_id":49155}))
|
82 |
+
if "service_context" not in st.session_state:
|
83 |
+
st.session_state.service_context = ServiceContext.from_defaults(llm_predictor=st.session_state.llm_predictor, embed_model=st.session_state.embed_model)
|
84 |
+
start_3 = timeit.default_timer() # Start timer
|
85 |
+
st.write(f"向量库构建开始:{start_3}")
|
86 |
+
new_index = VectorStoreIndex.from_documents(documents, service_context=st.session_state.service_context)
|
87 |
+
end_3 = timeit.default_timer() # Start timer
|
88 |
+
st.write(f"向量库构建结束:{end_3}")
|
89 |
+
st.write(f"向量库构建耗时:{end_3 - start_3}")
|
90 |
+
directory_path = generate_random_string(20)
|
91 |
+
os.makedirs(directory_path)
|
92 |
+
new_index.storage_context.persist("directory_path")
|
93 |
+
storage_context = StorageContext.from_defaults(persist_dir="directory_path")
|
94 |
+
start_4 = timeit.default_timer() # Start timer
|
95 |
+
st.write(f"向量库装载开始:{start_4}")
|
96 |
+
loadedindex = load_index_from_storage(storage_context=storage_context, service_context=st.session_state.service_context)
|
97 |
+
end_4 = timeit.default_timer() # Start timer
|
98 |
+
st.write(f"向量库装载结束:{end_4}")
|
99 |
+
st.write(f"向量库装载耗时:{end_4 - start_4}")
|
100 |
+
query_engine = loadedindex.as_query_engine()
|
101 |
+
st.session_state.query_engine = query_engine
|
102 |
+
|
103 |
+
user_question = st.text_input("Enter your query:")
|
104 |
+
if user_question != "" and not user_question.strip().isspace() and not user_question == "" and not user_question.strip() == "" and not user_question.isspace() and st.session_state.query_engine is not None:
|
105 |
+
print("user question: "+user_question)
|
106 |
+
with st.spinner("AI Thinking...Please wait a while to Cheers!"):
|
107 |
+
start_5 = timeit.default_timer() # Start timer
|
108 |
+
st.write(f"Query Engine - AI QA开始:{start_5}")
|
109 |
+
initial_response = st.session_state.query_engine.query(user_question)
|
110 |
+
end_5 = timeit.default_timer() # Start timer
|
111 |
+
st.write(f"Query Engine - AI QA结束:{end_5}")
|
112 |
+
st.write(f"Query Engine - AI QA耗时:{end_5 - start_5}")
|
113 |
+
temp_ai_response=str(initial_response)
|
114 |
+
final_ai_response=temp_ai_response.partition('<|end|>')[0]
|
115 |
+
print("AI Response:\n"+final_ai_response)
|
116 |
+
st.write("AI Response:\n\n"+final_ai_response)
|