binqiangliu commited on
Commit
24680b6
·
1 Parent(s): 2d56ea0

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +116 -0
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)