la04 commited on
Commit
812f60c
·
verified ·
1 Parent(s): c376788

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +58 -151
app.py CHANGED
@@ -1,100 +1,74 @@
1
  import gradio as gr
2
  import os
3
- api_token = os.getenv("HF_TOKEN")
4
-
5
-
6
  from langchain_community.vectorstores import FAISS
7
  from langchain_community.document_loaders import PyPDFLoader
8
  from langchain.text_splitter import RecursiveCharacterTextSplitter
9
- from langchain_community.vectorstores import Chroma
10
- from langchain.chains import ConversationalRetrievalChain
11
  from langchain_community.embeddings import HuggingFaceEmbeddings
12
- from langchain_community.llms import HuggingFacePipeline
13
- from langchain.chains import ConversationChain
14
- from langchain.memory import ConversationBufferMemory
15
  from langchain_community.llms import HuggingFaceEndpoint
16
- import torch
17
 
18
- list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]
 
19
  list_llm_simple = [os.path.basename(llm) for llm in list_llm]
20
 
21
- # Load and split PDF document
22
  def load_doc(list_file_path):
23
- # Processing for one document only
24
- # loader = PyPDFLoader(file_path)
25
- # pages = loader.load()
26
  loaders = [PyPDFLoader(x) for x in list_file_path]
27
  pages = []
28
  for loader in loaders:
29
  pages.extend(loader.load())
30
  text_splitter = RecursiveCharacterTextSplitter(
31
- chunk_size = 1024,
32
- chunk_overlap = 64
33
- )
34
  doc_splits = text_splitter.split_documents(pages)
35
  return doc_splits
36
 
37
- # Create vector database
38
  def create_db(splits):
39
  embeddings = HuggingFaceEmbeddings()
40
  vectordb = FAISS.from_documents(splits, embeddings)
41
  return vectordb
42
 
 
 
 
 
 
 
 
 
43
 
44
- # Initialize langchain LLM chain
45
- def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
46
- if llm_model == "meta-llama/Meta-Llama-3-8B-Instruct":
47
- llm = HuggingFaceEndpoint(
48
- repo_id=llm_model,
49
- huggingfacehub_api_token = api_token,
50
- temperature = temperature,
51
- max_new_tokens = max_tokens,
52
- top_k = top_k,
53
- )
54
- else:
55
- llm = HuggingFaceEndpoint(
56
- huggingfacehub_api_token = api_token,
57
- repo_id=llm_model,
58
- temperature = temperature,
59
- max_new_tokens = max_tokens,
60
- top_k = top_k,
61
- )
62
-
63
  memory = ConversationBufferMemory(
64
  memory_key="chat_history",
65
  output_key='answer',
66
  return_messages=True
67
  )
68
 
69
- retriever=vector_db.as_retriever()
70
  qa_chain = ConversationalRetrievalChain.from_llm(
71
  llm,
72
  retriever=retriever,
73
- chain_type="stuff",
74
  memory=memory,
75
  return_source_documents=True,
76
- verbose=False,
77
  )
78
  return qa_chain
79
 
80
- # Initialize database
81
- def initialize_database(list_file_obj, progress=gr.Progress()):
82
- # Create a list of documents (when valid)
83
  list_file_path = [x.name for x in list_file_obj if x is not None]
84
- # Load document and create splits
85
  doc_splits = load_doc(list_file_path)
86
- # Create or load vector database
87
  vector_db = create_db(doc_splits)
88
- return vector_db, "Database created!"
89
 
90
- # Initialize LLM
91
- def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db, progress=gr.Progress()):
92
- # print("llm_option",llm_option)
93
  llm_name = list_llm[llm_option]
94
- print("llm_name: ",llm_name)
95
- qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db, progress)
96
- return qa_chain, "QA chain initialized. Chatbot is ready!"
97
-
98
 
99
  def format_chat_history(message, chat_history):
100
  formatted_chat_history = []
@@ -102,115 +76,48 @@ def format_chat_history(message, chat_history):
102
  formatted_chat_history.append(f"User: {user_message}")
103
  formatted_chat_history.append(f"Assistant: {bot_message}")
104
  return formatted_chat_history
105
-
106
 
 
107
  def conversation(qa_chain, message, history):
108
  formatted_chat_history = format_chat_history(message, history)
109
- # Generate response using QA chain
110
  response = qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
111
  response_answer = response["answer"]
112
- if response_answer.find("Helpful Answer:") != -1:
113
  response_answer = response_answer.split("Helpful Answer:")[-1]
114
- response_sources = response["source_documents"]
115
- response_source1 = response_sources[0].page_content.strip()
116
- response_source2 = response_sources[1].page_content.strip()
117
- response_source3 = response_sources[2].page_content.strip()
118
- # Langchain sources are zero-based
119
- response_source1_page = response_sources[0].metadata["page"] + 1
120
- response_source2_page = response_sources[1].metadata["page"] + 1
121
- response_source3_page = response_sources[2].metadata["page"] + 1
122
- # Append user message and response to chat history
123
  new_history = history + [(message, response_answer)]
124
- return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
125
-
126
-
127
- def upload_file(file_obj):
128
- list_file_path = []
129
- for idx, file in enumerate(file_obj):
130
- file_path = file_obj.name
131
- list_file_path.append(file_path)
132
- return list_file_path
133
-
134
 
 
135
  def demo():
136
- # with gr.Blocks(theme=gr.themes.Default(primary_hue="sky")) as demo:
137
- with gr.Blocks(theme=gr.themes.Default(primary_hue="red", secondary_hue="pink", neutral_hue = "sky")) as demo:
138
  vector_db = gr.State()
139
  qa_chain = gr.State()
140
- gr.HTML("<center><h1>RAG PDF chatbot</h1><center>")
141
- gr.Markdown("""<b>Query your PDF documents!</b> This AI agent is designed to perform retrieval augmented generation (RAG) on PDF documents. The app is hosted on Hugging Face Hub for the sole purpose of demonstration. \
142
- <b>Please do not upload confidential documents.</b>
143
- """)
144
  with gr.Row():
145
- with gr.Column(scale = 86):
146
- gr.Markdown("<b>Step 1 - Upload PDF documents and Initialize RAG pipeline</b>")
147
- with gr.Row():
148
- document = gr.Files(height=300, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload PDF documents")
149
- with gr.Row():
150
- db_btn = gr.Button("Create vector database")
151
- with gr.Row():
152
- db_progress = gr.Textbox(value="Not initialized", show_label=False) # label="Vector database status",
153
- gr.Markdown("<style>body { font-size: 16px; }</style><b>Select Large Language Model (LLM) and input parameters</b>")
154
- with gr.Row():
155
- llm_btn = gr.Radio(list_llm_simple, label="Available LLMs", value = list_llm_simple[0], type="index") # info="Select LLM", show_label=False
156
- with gr.Row():
157
- with gr.Accordion("LLM input parameters", open=False):
158
- with gr.Row():
159
- slider_temperature = gr.Slider(minimum = 0.01, maximum = 1.0, value=0.5, step=0.1, label="Temperature", info="Controls randomness in token generation", interactive=True)
160
- with gr.Row():
161
- slider_maxtokens = gr.Slider(minimum = 128, maximum = 9192, value=4096, step=128, label="Max New Tokens", info="Maximum number of tokens to be generated",interactive=True)
162
- with gr.Row():
163
- slider_topk = gr.Slider(minimum = 1, maximum = 10, value=3, step=1, label="top-k", info="Number of tokens to select the next token from", interactive=True)
164
- with gr.Row():
165
- qachain_btn = gr.Button("Initialize Question Answering Chatbot")
166
- with gr.Row():
167
- llm_progress = gr.Textbox(value="Not initialized", show_label=False) # label="Chatbot status",
168
-
169
- with gr.Column(scale = 200):
170
- gr.Markdown("<b>Step 2 - Chat with your Document</b>")
171
- chatbot = gr.Chatbot(height=505)
172
- with gr.Accordion("Relevent context from the source document", open=False):
173
- with gr.Row():
174
- doc_source1 = gr.Textbox(label="Reference 1", lines=2, container=True, scale=20)
175
- source1_page = gr.Number(label="Page", scale=1)
176
- with gr.Row():
177
- doc_source2 = gr.Textbox(label="Reference 2", lines=2, container=True, scale=20)
178
- source2_page = gr.Number(label="Page", scale=1)
179
- with gr.Row():
180
- doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)
181
- source3_page = gr.Number(label="Page", scale=1)
182
- with gr.Row():
183
- msg = gr.Textbox(placeholder="Ask a question", container=True)
184
- with gr.Row():
185
- submit_btn = gr.Button("Submit")
186
- clear_btn = gr.ClearButton([msg, chatbot], value="Clear")
187
-
188
- # Preprocessing events
189
- db_btn.click(initialize_database, \
190
- inputs=[document], \
191
- outputs=[vector_db, db_progress])
192
- qachain_btn.click(initialize_LLM, \
193
- inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], \
194
- outputs=[qa_chain, llm_progress]).then(lambda:[None,"",0,"",0,"",0], \
195
- inputs=None, \
196
- outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
197
- queue=False)
198
-
199
- # Chatbot events
200
- msg.submit(conversation, \
201
- inputs=[qa_chain, msg, chatbot], \
202
- outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
203
- queue=False)
204
- submit_btn.click(conversation, \
205
- inputs=[qa_chain, msg, chatbot], \
206
- outputs=[qa_chain, msg, chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
207
- queue=False)
208
- clear_btn.click(lambda:[None,"",0,"",0,"",0], \
209
- inputs=None, \
210
- outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], \
211
- queue=False)
212
- demo.queue().launch(debug=True)
213
-
214
 
215
  if __name__ == "__main__":
216
- demo()
 
1
  import gradio as gr
2
  import os
 
 
 
3
  from langchain_community.vectorstores import FAISS
4
  from langchain_community.document_loaders import PyPDFLoader
5
  from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 
6
  from langchain_community.embeddings import HuggingFaceEmbeddings
7
+ from langchain.chains import ConversationalRetrievalChain
 
 
8
  from langchain_community.llms import HuggingFaceEndpoint
9
+ from langchain.memory import ConversationBufferMemory
10
 
11
+ # Liste der Modelle
12
+ list_llm = ["google/flan-t5-small", "distilbert-base-uncased"] # Leichtere, CPU-freundliche Modelle
13
  list_llm_simple = [os.path.basename(llm) for llm in list_llm]
14
 
15
+ # PDF-Dokument laden und aufteilen
16
  def load_doc(list_file_path):
 
 
 
17
  loaders = [PyPDFLoader(x) for x in list_file_path]
18
  pages = []
19
  for loader in loaders:
20
  pages.extend(loader.load())
21
  text_splitter = RecursiveCharacterTextSplitter(
22
+ chunk_size=512, # Kleinere Chunks für schnellere Verarbeitung
23
+ chunk_overlap=32
24
+ )
25
  doc_splits = text_splitter.split_documents(pages)
26
  return doc_splits
27
 
28
+ # Erstellen der Vektordatenbank
29
  def create_db(splits):
30
  embeddings = HuggingFaceEmbeddings()
31
  vectordb = FAISS.from_documents(splits, embeddings)
32
  return vectordb
33
 
34
+ # Initialisierung des LLM Chains
35
+ def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
36
+ llm = HuggingFaceEndpoint(
37
+ repo_id=llm_model,
38
+ temperature=temperature,
39
+ max_new_tokens=max_tokens,
40
+ top_k=top_k
41
+ )
42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
  memory = ConversationBufferMemory(
44
  memory_key="chat_history",
45
  output_key='answer',
46
  return_messages=True
47
  )
48
 
49
+ retriever = vector_db.as_retriever()
50
  qa_chain = ConversationalRetrievalChain.from_llm(
51
  llm,
52
  retriever=retriever,
53
+ chain_type="stuff",
54
  memory=memory,
55
  return_source_documents=True,
56
+ verbose=False
57
  )
58
  return qa_chain
59
 
60
+ # Initialisierung der Datenbank
61
+ def initialize_database(list_file_obj):
 
62
  list_file_path = [x.name for x in list_file_obj if x is not None]
 
63
  doc_splits = load_doc(list_file_path)
 
64
  vector_db = create_db(doc_splits)
65
+ return vector_db, "Datenbank erfolgreich erstellt!"
66
 
67
+ # Initialisierung des LLMs
68
+ def initialize_LLM(llm_option, llm_temperature, max_tokens, top_k, vector_db):
 
69
  llm_name = list_llm[llm_option]
70
+ qa_chain = initialize_llmchain(llm_name, llm_temperature, max_tokens, top_k, vector_db)
71
+ return qa_chain, "LLM erfolgreich initialisiert! Chatbot ist bereit."
 
 
72
 
73
  def format_chat_history(message, chat_history):
74
  formatted_chat_history = []
 
76
  formatted_chat_history.append(f"User: {user_message}")
77
  formatted_chat_history.append(f"Assistant: {bot_message}")
78
  return formatted_chat_history
 
79
 
80
+ # Chat-Funktion
81
  def conversation(qa_chain, message, history):
82
  formatted_chat_history = format_chat_history(message, history)
 
83
  response = qa_chain.invoke({"question": message, "chat_history": formatted_chat_history})
84
  response_answer = response["answer"]
85
+ if "Helpful Answer:" in response_answer:
86
  response_answer = response_answer.split("Helpful Answer:")[-1]
 
 
 
 
 
 
 
 
 
87
  new_history = history + [(message, response_answer)]
88
+ return qa_chain, gr.update(value=""), new_history
 
 
 
 
 
 
 
 
 
89
 
90
+ # Gradio App erstellen
91
  def demo():
92
+ with gr.Blocks() as demo:
 
93
  vector_db = gr.State()
94
  qa_chain = gr.State()
95
+ gr.HTML("<center><h1>RAG PDF Chatbot</h1></center>")
 
 
 
96
  with gr.Row():
97
+ with gr.Column():
98
+ gr.Markdown("### Schritt 1: Lade PDF-Dokument hoch")
99
+ document = gr.Files(height=300, file_count="multiple", file_types=[".pdf"], interactive=True)
100
+ db_btn = gr.Button("Erstelle Vektordatenbank")
101
+ db_progress = gr.Textbox(value="Nicht initialisiert", show_label=False)
102
+ gr.Markdown("### Schritt 2: Wähle LLM und Einstellungen")
103
+ llm_btn = gr.Radio(list_llm_simple, label="Verfügbare Modelle", value=list_llm_simple[0], type="index")
104
+ slider_temperature = gr.Slider(0.01, 1.0, value=0.5, step=0.1, label="Temperature")
105
+ slider_maxtokens = gr.Slider(64, 512, value=256, step=64, label="Max Tokens")
106
+ slider_topk = gr.Slider(1, 10, value=3, step=1, label="Top-k")
107
+ qachain_btn = gr.Button("Initialisiere QA-Chatbot")
108
+ llm_progress = gr.Textbox(value="Nicht initialisiert", show_label=False)
109
+
110
+ with gr.Column():
111
+ gr.Markdown("### Schritt 3: Stelle Fragen an dein Dokument")
112
+ chatbot = gr.Chatbot(height=400, type="messages")
113
+ msg = gr.Textbox(placeholder="Frage stellen...")
114
+ submit_btn = gr.Button("Absenden")
115
+
116
+ db_btn.click(initialize_database, [document], [vector_db, db_progress])
117
+ qachain_btn.click(initialize_LLM, [llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], [qa_chain, llm_progress])
118
+ msg.submit(conversation, [qa_chain, msg, chatbot], [qa_chain, msg, chatbot])
119
+ submit_btn.click(conversation, [qa_chain, msg, chatbot], [qa_chain, msg, chatbot])
120
+ demo.launch(debug=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
121
 
122
  if __name__ == "__main__":
123
+ demo()