awacke1 commited on
Commit
3130701
1 Parent(s): dc50884

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +243 -0
app.py ADDED
@@ -0,0 +1,243 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import transformers
3
+ import gradio as gr
4
+ from ragatouille import RAGPretrainedModel
5
+ from huggingface_hub import InferenceClient
6
+ import re
7
+ from datetime import datetime
8
+ import json
9
+ import os
10
+
11
+ import arxiv
12
+ from utils import get_md_text_abstract, search_cleaner, get_arxiv_live_search
13
+
14
+ retrieve_results = 10
15
+ show_examples = False
16
+ llm_models_to_choose = ['mistralai/Mixtral-8x7B-Instruct-v0.1','mistralai/Mistral-7B-Instruct-v0.2', 'google/gemma-7b-it', 'None']
17
+
18
+ generate_kwargs = dict(
19
+ temperature = None,
20
+ max_new_tokens = 512,
21
+ top_p = None,
22
+ do_sample = False,
23
+ )
24
+
25
+ RAG = RAGPretrainedModel.from_index("colbert/indexes/arxiv_colbert")
26
+
27
+ try:
28
+ gr.Info("Setting up retriever, please wait...")
29
+ rag_initial_output = RAG.search("what is Mistral?", k = 1)
30
+ gr.Info("Retriever working successfully!")
31
+
32
+ except:
33
+ gr.Warning("Retriever not working!")
34
+
35
+ mark_text = '# 🩺🔍 Search Results\n'
36
+ header_text = "## Arxiv Paper Summary With QA Retrieval Augmented Generation \n"
37
+
38
+ try:
39
+ with open("README.md", "r") as f:
40
+ mdfile = f.read()
41
+ date_pattern = r'Index Last Updated : \d{4}-\d{2}-\d{2}'
42
+ match = re.search(date_pattern, mdfile)
43
+ date = match.group().split(': ')[1]
44
+ formatted_date = datetime.strptime(date, '%Y-%m-%d').strftime('%d %b %Y')
45
+ header_text += f'Index Last Updated: {formatted_date}\n'
46
+ index_info = f"Semantic Search - up to {formatted_date}"
47
+ except:
48
+ index_info = "Semantic Search"
49
+
50
+ database_choices = [index_info,'Arxiv Search - Latest - (EXPERIMENTAL)']
51
+
52
+ arx_client = arxiv.Client()
53
+ is_arxiv_available = True
54
+ check_arxiv_result = get_arxiv_live_search("What is Self Rewarding AI and how can it be used in Multi-Agent Systems?", arx_client, retrieve_results)
55
+ if len(check_arxiv_result) == 0:
56
+ is_arxiv_available = False
57
+ print("Arxiv search not working, switching to default search ...")
58
+ database_choices = [index_info]
59
+
60
+ if show_examples:
61
+ with open("sample_outputs.json", "r") as f:
62
+ sample_outputs = json.load(f)
63
+ output_placeholder = sample_outputs['output_placeholder']
64
+ md_text_initial = sample_outputs['search_placeholder']
65
+
66
+ else:
67
+ output_placeholder = None
68
+ md_text_initial = ''
69
+
70
+ def rag_cleaner(inp):
71
+ rank = inp['rank']
72
+ title = inp['document_metadata']['title']
73
+ content = inp['content']
74
+ date = inp['document_metadata']['_time']
75
+ return f"{rank}. <b> {title} </b> \n Date : {date} \n Abstract: {content}"
76
+
77
+ def get_prompt_text(question, context, formatted = True, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
78
+ if formatted:
79
+ sys_instruction = f"Context:\n {context} \n Given the following scientific paper abstracts, take a deep breath and lets think step by step to answer the question. Cite the titles of your sources when answering, do not cite links or dates."
80
+ message = f"Question: {question}"
81
+
82
+ if 'mistralai' in llm_model_picked:
83
+ return f"<s>" + f"[INST] {sys_instruction}" + f" {message}[/INST]"
84
+
85
+ elif 'gemma' in llm_model_picked:
86
+ return f"<bos><start_of_turn>user\n{sys_instruction}" + f" {message}<end_of_turn>\n"
87
+
88
+ return f"Context:\n {context} \n Given the following info, take a deep breath and lets think step by step to answer the question: {question}. Cite the titles of your sources when answering.\n\n"
89
+
90
+ def get_references(question, retriever, k = retrieve_results):
91
+ rag_out = retriever.search(query=question, k=k)
92
+ return rag_out
93
+
94
+ def get_rag(message):
95
+ return get_references(message, RAG)
96
+
97
+ def SaveResponseAndRead(result):
98
+ documentHTML5='''
99
+ <!DOCTYPE html>
100
+ <html>
101
+ <head>
102
+ <title>Read It Aloud</title>
103
+ <script type="text/javascript">
104
+ function readAloud() {
105
+ const text = document.getElementById("textArea").value;
106
+ const speech = new SpeechSynthesisUtterance(text);
107
+ window.speechSynthesis.speak(speech);
108
+ }
109
+ </script>
110
+ </head>
111
+ <body>
112
+ <h1>🔊 Read It Aloud</h1>
113
+ <textarea id="textArea" rows="10" cols="80">
114
+ '''
115
+ documentHTML5 = documentHTML5 + result
116
+ documentHTML5 = documentHTML5 + '''
117
+ </textarea>
118
+ <br>
119
+ <button onclick="readAloud()">🔊 Read Aloud</button>
120
+ </body>
121
+ </html>
122
+ '''
123
+ gr.HTML(documentHTML5)
124
+
125
+ def save_search_results(prompt, results, response):
126
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
127
+ filename = f"{timestamp}_{re.sub(r'[^\w\-_\. ]', '_', prompt)}.txt"
128
+ with open(filename, "w") as f:
129
+ f.write(f"# {prompt}\n\n")
130
+ f.write(f"## Search Results\n\n{results}\n\n")
131
+ f.write(f"## LLM Response\n\n{response}\n")
132
+ return filename
133
+
134
+ def get_past_searches():
135
+ txt_files = [f for f in os.listdir(".") if f.endswith(".txt") and f != "requirements.txt"]
136
+ return txt_files
137
+
138
+ with gr.Blocks(theme = gr.themes.Soft()) as demo:
139
+ header = gr.Markdown(header_text)
140
+
141
+ with gr.Row():
142
+ with gr.Column():
143
+ msg = gr.Textbox(label = 'Search', placeholder = 'What is Mistral?')
144
+
145
+ with gr.Accordion("Advanced Settings", open=False):
146
+ with gr.Row(equal_height = True):
147
+ llm_model = gr.Dropdown(choices = llm_models_to_choose, value = 'mistralai/Mistral-7B-Instruct-v0.2', label = 'LLM Model')
148
+ llm_results = gr.Slider(minimum=4, maximum=10, value=5, step=1, interactive=True, label="Top n results as context")
149
+ database_src = gr.Dropdown(choices = database_choices, value = index_info, label = 'Search Source')
150
+ stream_results = gr.Checkbox(value = True, label = "Stream output", visible = False)
151
+
152
+ output_text = gr.Textbox(show_label = True, container = True, label = 'LLM Answer', visible = True, placeholder = output_placeholder)
153
+ input = gr.Textbox(show_label = False, visible = False)
154
+ gr_md = gr.Markdown(mark_text + md_text_initial)
155
+
156
+ with gr.Column():
157
+ past_searches = gr.Dropdown(choices=get_past_searches(), label="Past Searches")
158
+ past_search_content = gr.Textbox(label="Past Search Content", visible=False)
159
+
160
+ def update_past_search_content(past_search):
161
+ if past_search:
162
+ with open(past_search, "r") as f:
163
+ content = f.read()
164
+ return gr.Textbox.update(value=content, visible=True)
165
+ else:
166
+ return gr.Textbox.update(visible=False)
167
+
168
+ past_searches.change(update_past_search_content, past_searches, past_search_content)
169
+
170
+ def update_with_rag_md(message, llm_results_use = 5, database_choice = index_info, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2'):
171
+ prompt_text_from_data = ""
172
+ database_to_use = database_choice
173
+ if database_choice == index_info:
174
+ rag_out = get_rag(message)
175
+ else:
176
+ arxiv_search_success = True
177
+ try:
178
+ rag_out = get_arxiv_live_search(message, arx_client, retrieve_results)
179
+ if len(rag_out) == 0:
180
+ arxiv_search_success = False
181
+ except:
182
+ arxiv_search_success = False
183
+
184
+ if not arxiv_search_success:
185
+ gr.Warning("Arxiv Search not working, switching to semantic search ...")
186
+ rag_out = get_rag(message)
187
+ database_to_use = index_info
188
+
189
+ md_text_updated = mark_text
190
+ for i in range(retrieve_results):
191
+ rag_answer = rag_out[i]
192
+ if i < llm_results_use:
193
+ md_text_paper, prompt_text = get_md_text_abstract(rag_answer, source = database_to_use, return_prompt_formatting = True)
194
+ prompt_text_from_data += f"{i+1}. {prompt_text}"
195
+ else:
196
+ md_text_paper = get_md_text_abstract(rag_answer, source = database_to_use)
197
+ md_text_updated += md_text_paper
198
+ prompt = get_prompt_text(message, prompt_text_from_data, llm_model_picked = llm_model_picked)
199
+
200
+ filename = save_search_results(message, md_text_updated, "")
201
+
202
+ with open(filename, "r") as f:
203
+ md_content = f.read()
204
+
205
+ return md_content, prompt, get_past_searches()
206
+
207
+ def ask_llm(prompt, llm_model_picked = 'mistralai/Mistral-7B-Instruct-v0.2', stream_outputs = False):
208
+ model_disabled_text = "LLM Model is disabled"
209
+ output = ""
210
+
211
+ if llm_model_picked == 'None':
212
+ if stream_outputs:
213
+ for out in model_disabled_text:
214
+ output += out
215
+ yield output
216
+ return output
217
+ else:
218
+ return model_disabled_text
219
+
220
+ client = InferenceClient(llm_model_picked)
221
+ try:
222
+ stream = client.text_generation(prompt, **generate_kwargs, stream=stream_outputs, details=False, return_full_text=False)
223
+
224
+ except:
225
+ gr.Warning("LLM Inference rate limit reached, try again later!")
226
+ return ""
227
+
228
+ if stream_outputs:
229
+ for response in stream:
230
+ output += response
231
+ SaveResponseAndRead(response)
232
+ yield output
233
+ return output
234
+ else:
235
+ return stream
236
+
237
+ msg.submit(update_with_rag_md, [msg, llm_results, database_src, llm_model], [gr_md, input, past_searches]).success(ask_llm, [input, llm_model, stream_results], output_text).then(
238
+ lambda response: save_search_results(msg.value, gr_md.value, response),
239
+ [msg, gr_md, output_text],
240
+ None
241
+ )
242
+
243
+ demo.queue().launch()