Sourab Mangrulkar commited on
Commit
f01bff8
1 Parent(s): a7bdbc0
Files changed (5) hide show
  1. agent.py +65 -0
  2. app.py +442 -0
  3. chunked_data.parquet +3 -0
  4. requirements.txt +9 -0
  5. search_index.bin +3 -0
agent.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from threading import Thread
2
+ from typing import Iterator
3
+
4
+ import torch
5
+ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
6
+
7
+ model_id = "meta-llama/Llama-2-7b-chat-hf"
8
+
9
+ if torch.cuda.is_available():
10
+ model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
11
+ else:
12
+ model = None
13
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
14
+
15
+
16
+ def get_prompt(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> str:
17
+ texts = [f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"]
18
+ # The first user input is _not_ stripped
19
+ do_strip = False
20
+ for user_input, response in chat_history:
21
+ user_input = user_input.strip() if do_strip else user_input
22
+ do_strip = True
23
+ texts.append(f"{user_input} [/INST] {response.strip()} </s><s>[INST] ")
24
+ message = message.strip() if do_strip else message
25
+ texts.append(f"{message} [/INST]")
26
+ return "".join(texts)
27
+
28
+
29
+ def get_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> int:
30
+ prompt = get_prompt(message, chat_history, system_prompt)
31
+ input_ids = tokenizer([prompt], return_tensors="np", add_special_tokens=False)["input_ids"]
32
+ return input_ids.shape[-1]
33
+
34
+
35
+ def run(
36
+ message: str,
37
+ chat_history: list[tuple[str, str]],
38
+ system_prompt: str,
39
+ max_new_tokens: int = 1024,
40
+ temperature: float = 0.8,
41
+ top_p: float = 0.95,
42
+ top_k: int = 50,
43
+ ) -> Iterator[str]:
44
+ prompt = get_prompt(message, chat_history, system_prompt)
45
+ inputs = tokenizer([prompt], return_tensors="pt", add_special_tokens=False).to("cuda")
46
+
47
+ streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
48
+ generate_kwargs = dict(
49
+ inputs,
50
+ streamer=streamer,
51
+ max_new_tokens=max_new_tokens,
52
+ do_sample=True,
53
+ top_p=top_p,
54
+ top_k=top_k,
55
+ temperature=temperature,
56
+ num_beams=1,
57
+ eos_token_id=tokenizer.eos_token_id,
58
+ )
59
+ t = Thread(target=model.generate, kwargs=generate_kwargs)
60
+ t.start()
61
+
62
+ outputs = []
63
+ for text in streamer:
64
+ outputs.append(text)
65
+ yield "".join(outputs)
app.py ADDED
@@ -0,0 +1,442 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import argparse
2
+ import os
3
+ import json
4
+ import re
5
+ from sentence_transformers import SentenceTransformer, CrossEncoder
6
+ import hnswlib
7
+ from typing import Iterator
8
+
9
+ import gradio as gr
10
+ import pandas as pd
11
+ import torch
12
+
13
+ from easyllm.clients import huggingface
14
+
15
+ from agent import get_input_token_length
16
+
17
+ huggingface.prompt_builder = "llama2"
18
+ huggingface.api_key = os.environ["HUGGINGFACE_TOKEN"]
19
+ MAX_MAX_NEW_TOKENS = 2048
20
+ DEFAULT_MAX_NEW_TOKENS = 1024
21
+ MAX_INPUT_TOKEN_LENGTH = 4000
22
+ EMBED_DIM = 1024
23
+ K = 10
24
+ EF = 100
25
+ SEARCH_INDEX = "search_index.bin"
26
+ DOCUMENT_DATASET = "chunked_data.parquet"
27
+ COSINE_THRESHOLD = 0.7
28
+
29
+ torch_device = "cuda" if torch.cuda.is_available() else "cpu"
30
+ print("Running on device:", torch_device)
31
+ print("CPU threads:", torch.get_num_threads())
32
+
33
+ biencoder = SentenceTransformer("intfloat/e5-large-v2", device=torch_device)
34
+ cross_encoder = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2", max_length=512, device=torch_device)
35
+
36
+
37
+ def create_qa_prompt(query, relevant_chunks):
38
+ stuffed_context = " ".join(relevant_chunks)
39
+ return f"""\
40
+ Use the following pieces of context given in to answer the question at the end. \
41
+ If you don't know the answer, just say that you don't know, don't try to make up an answer. \
42
+ Keep the answer short and succinct.
43
+
44
+ Context: {stuffed_context}
45
+ Question: {query}
46
+ Helpful Answer: \
47
+ """
48
+
49
+
50
+ def create_condense_question_prompt(question, chat_history):
51
+ return f"""\
52
+ Given the following conversation and a follow up question, \
53
+ rephrase the follow up question to be a standalone question in its original language. \
54
+ Output the json object with single field `question` and value being the rephrased standalone question.
55
+ Only output json object and nothing else.
56
+
57
+ Chat History:
58
+ {chat_history}
59
+ Follow Up Input: {question}
60
+ """
61
+
62
+
63
+ # https://www.philschmid.de/llama-2#how-to-prompt-llama-2-chat
64
+ def get_completion(
65
+ prompt,
66
+ system_prompt=None,
67
+ model="meta-llama/Llama-2-70b-chat-hf",
68
+ max_new_tokens=1024,
69
+ temperature=0.2,
70
+ top_p=0.95,
71
+ top_k=50,
72
+ stream=False,
73
+ debug=False,
74
+ ):
75
+ if temperature < 1e-2:
76
+ temperature = 1e-2
77
+ messages = []
78
+ if system_prompt is not None:
79
+ messages.append({"role": "system", "content": system_prompt})
80
+ messages.append({"role": "user", "content": prompt})
81
+ response = huggingface.ChatCompletion.create(
82
+ model=model,
83
+ messages=messages,
84
+ temperature=temperature, # this is the degree of randomness of the model's output
85
+ max_tokens=max_new_tokens, # this is the number of new tokens being generated
86
+ top_p=top_p,
87
+ top_k=top_k,
88
+ stream=stream,
89
+ debug=debug,
90
+ )
91
+ return response["choices"][0]["message"]["content"] if not stream else response
92
+
93
+
94
+ # load the index for the PEFT docs
95
+ def load_hnsw_index(index_file):
96
+ # Load the HNSW index from the specified file
97
+ index = hnswlib.Index(space="ip", dim=EMBED_DIM)
98
+ index.load_index(index_file)
99
+ return index
100
+
101
+
102
+ def create_query_embedding(query):
103
+ # Encode the query to get its embedding
104
+ embedding = biencoder.encode([query], normalize_embeddings=True)[0]
105
+ return embedding
106
+
107
+
108
+ def find_nearest_neighbors(query_embedding):
109
+ search_index.set_ef(EF)
110
+ # Find the k-nearest neighbors for the query embedding
111
+ labels, distances = search_index.knn_query(query_embedding, k=K)
112
+ labels = [label for label, distance in zip(labels[0], distances[0]) if (1 - distance) >= COSINE_THRESHOLD]
113
+ relevant_chunks = data_df.iloc[labels]["chunk_content"].tolist()
114
+ return relevant_chunks
115
+
116
+
117
+ def rerank_chunks_with_cross_encoder(query, chunks):
118
+ # Create a list of tuples, each containing a query-chunk pair
119
+ pairs = [(query, chunk) for chunk in chunks]
120
+
121
+ # Get scores for each query-chunk pair using the cross encoder
122
+ scores = cross_encoder.predict(pairs)
123
+
124
+ # Sort the chunks based on their scores in descending order
125
+ sorted_chunks = [chunk for _, chunk in sorted(zip(scores, chunks), reverse=True)]
126
+
127
+ return sorted_chunks
128
+
129
+
130
+ def generate_condensed_query(query, history):
131
+ chat_history = ""
132
+ for turn in history:
133
+ chat_history += f"Human: {turn[0]}\n"
134
+ chat_history += f"Assistant: {turn[1]}\n"
135
+
136
+ condense_question_prompt = create_condense_question_prompt(query, chat_history)
137
+ condensed_question = json.loads(get_completion(condense_question_prompt, max_new_tokens=64, temperature=0))
138
+ return condensed_question["question"]
139
+
140
+
141
+ DEFAULT_SYSTEM_PROMPT = """\
142
+ You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
143
+
144
+ If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\
145
+ """
146
+ MAX_MAX_NEW_TOKENS = 2048
147
+ DEFAULT_MAX_NEW_TOKENS = 1024
148
+ MAX_INPUT_TOKEN_LENGTH = 4000
149
+
150
+ DESCRIPTION = """
151
+ # PEFT Docs QA Chatbot 🤗
152
+ """
153
+
154
+ LICENSE = """
155
+ <p/>
156
+
157
+ ---
158
+ As a derivate work of [Llama-2-70b-chat](https://huggingface.co/meta-llama/Llama-2-70b-chat) by Meta,
159
+ this demo is governed by the original [license](https://huggingface.co/spaces/huggingface-projects/llama-2-70b-chat/blob/main/LICENSE.txt) and [acceptable use policy](https://huggingface.co/spaces/huggingface-projects/llama-2-70b-chat/blob/main/USE_POLICY.md).
160
+ """
161
+
162
+ if not torch.cuda.is_available():
163
+ DESCRIPTION += "\n<p>Running on CPU 🥶.</p>"
164
+
165
+
166
+ def clear_and_save_textbox(message: str) -> tuple[str, str]:
167
+ return "", message
168
+
169
+
170
+ def display_input(message: str, history: list[tuple[str, str]]) -> list[tuple[str, str]]:
171
+ history.append((message, ""))
172
+ return history
173
+
174
+
175
+ def delete_prev_fn(history: list[tuple[str, str]]) -> tuple[list[tuple[str, str]], str]:
176
+ try:
177
+ message, _ = history.pop()
178
+ except IndexError:
179
+ message = ""
180
+ return history, message or ""
181
+
182
+
183
+ def wrap_html_code(text):
184
+ pattern = r"<.*?>"
185
+ matches = re.findall(pattern, text)
186
+ if len(matches) > 0:
187
+ return f"```{text}```"
188
+ else:
189
+ return text
190
+
191
+
192
+ def generate(
193
+ message: str,
194
+ history_with_input: list[tuple[str, str]],
195
+ system_prompt: str,
196
+ max_new_tokens: int,
197
+ temperature: float,
198
+ top_p: float,
199
+ top_k: int,
200
+ ) -> Iterator[list[tuple[str, str]]]:
201
+ if max_new_tokens > MAX_MAX_NEW_TOKENS:
202
+ raise ValueError
203
+ history = history_with_input[:-1]
204
+ if len(history) > 0:
205
+ condensed_query = generate_condensed_query(message, history)
206
+ print(f"{condensed_query=}")
207
+ else:
208
+ condensed_query = message
209
+ query_embedding = create_query_embedding(condensed_query)
210
+ relevant_chunks = find_nearest_neighbors(query_embedding)
211
+ reranked_relevant_chunks = rerank_chunks_with_cross_encoder(condensed_query, relevant_chunks)
212
+ qa_prompt = create_qa_prompt(condensed_query, reranked_relevant_chunks)
213
+ print(f"{qa_prompt=}")
214
+ generator = get_completion(
215
+ qa_prompt,
216
+ system_prompt=system_prompt,
217
+ stream=True,
218
+ max_new_tokens=max_new_tokens,
219
+ temperature=temperature,
220
+ top_k=top_k,
221
+ top_p=top_p,
222
+ )
223
+
224
+ output = ""
225
+ for idx, response in enumerate(generator):
226
+ token = response["choices"][0]["delta"].get("content", "") or ""
227
+ output += token
228
+ if idx == 0:
229
+ history.append((message, output))
230
+ else:
231
+ history[-1] = (message, output)
232
+
233
+ history = [
234
+ (wrap_html_code(history[i][0].strip()), wrap_html_code(history[i][1].strip()))
235
+ for i in range(0, len(history))
236
+ ]
237
+ yield history
238
+
239
+ return history
240
+
241
+
242
+ def process_example(message: str) -> tuple[str, list[tuple[str, str]]]:
243
+ generator = generate(message, [], DEFAULT_SYSTEM_PROMPT, 1024, 0.2, 0.95, 50)
244
+ for x in generator:
245
+ pass
246
+ return "", x
247
+
248
+
249
+ def check_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> None:
250
+ input_token_length = get_input_token_length(message, chat_history, system_prompt)
251
+ if input_token_length > MAX_INPUT_TOKEN_LENGTH:
252
+ raise gr.Error(
253
+ f"The accumulated input is too long ({input_token_length} > {MAX_INPUT_TOKEN_LENGTH}). Clear your chat history and try again."
254
+ )
255
+
256
+
257
+ search_index = load_hnsw_index(SEARCH_INDEX)
258
+ data_df = pd.read_parquet(DOCUMENT_DATASET).reset_index()
259
+ with gr.Blocks(css="style.css") as demo:
260
+ gr.Markdown(DESCRIPTION)
261
+
262
+ with gr.Group():
263
+ chatbot = gr.Chatbot(label="Chatbot")
264
+ with gr.Row():
265
+ textbox = gr.Textbox(
266
+ container=False,
267
+ show_label=False,
268
+ placeholder="Type a message...",
269
+ scale=10,
270
+ )
271
+ submit_button = gr.Button("Submit", variant="primary", scale=1, min_width=0)
272
+ with gr.Row():
273
+ retry_button = gr.Button("🔄 Retry", variant="secondary")
274
+ undo_button = gr.Button("↩️ Undo", variant="secondary")
275
+ clear_button = gr.Button("🗑️ Clear", variant="secondary")
276
+
277
+ saved_input = gr.State()
278
+
279
+ with gr.Accordion(label="Advanced options", open=False):
280
+ system_prompt = gr.Textbox(label="System prompt", value=DEFAULT_SYSTEM_PROMPT, lines=6)
281
+ max_new_tokens = gr.Slider(
282
+ label="Max new tokens",
283
+ minimum=1,
284
+ maximum=MAX_MAX_NEW_TOKENS,
285
+ step=1,
286
+ value=DEFAULT_MAX_NEW_TOKENS,
287
+ )
288
+ temperature = gr.Slider(
289
+ label="Temperature",
290
+ minimum=0.1,
291
+ maximum=4.0,
292
+ step=0.1,
293
+ value=0.2,
294
+ )
295
+ top_p = gr.Slider(
296
+ label="Top-p (nucleus sampling)",
297
+ minimum=0.05,
298
+ maximum=1.0,
299
+ step=0.05,
300
+ value=0.95,
301
+ )
302
+ top_k = gr.Slider(
303
+ label="Top-k",
304
+ minimum=1,
305
+ maximum=1000,
306
+ step=1,
307
+ value=50,
308
+ )
309
+
310
+ gr.Examples(
311
+ examples=[
312
+ "What is 🤗 PEFT?",
313
+ "How do I create a LoraConfig?",
314
+ "What are the different tuners supported?",
315
+ "How do I use LoRA with custom models?",
316
+ "What are the different real-world applications that I can use PEFT for?",
317
+ ],
318
+ inputs=textbox,
319
+ outputs=[textbox, chatbot],
320
+ # fn=process_example,
321
+ cache_examples=False,
322
+ )
323
+
324
+ gr.Markdown(LICENSE)
325
+
326
+ textbox.submit(
327
+ fn=clear_and_save_textbox,
328
+ inputs=textbox,
329
+ outputs=[textbox, saved_input],
330
+ api_name=False,
331
+ queue=False,
332
+ ).then(fn=display_input, inputs=[saved_input, chatbot], outputs=chatbot, api_name=False, queue=False,).then(
333
+ fn=check_input_token_length,
334
+ inputs=[saved_input, chatbot, system_prompt],
335
+ api_name=False,
336
+ queue=False,
337
+ ).success(
338
+ fn=generate,
339
+ inputs=[
340
+ saved_input,
341
+ chatbot,
342
+ system_prompt,
343
+ max_new_tokens,
344
+ temperature,
345
+ top_p,
346
+ top_k,
347
+ ],
348
+ outputs=chatbot,
349
+ api_name=False,
350
+ )
351
+
352
+ button_event_preprocess = (
353
+ submit_button.click(
354
+ fn=clear_and_save_textbox,
355
+ inputs=textbox,
356
+ outputs=[textbox, saved_input],
357
+ api_name=False,
358
+ queue=False,
359
+ )
360
+ .then(
361
+ fn=display_input,
362
+ inputs=[saved_input, chatbot],
363
+ outputs=chatbot,
364
+ api_name=False,
365
+ queue=False,
366
+ )
367
+ .then(
368
+ fn=check_input_token_length,
369
+ inputs=[saved_input, chatbot, system_prompt],
370
+ api_name=False,
371
+ queue=False,
372
+ )
373
+ .success(
374
+ fn=generate,
375
+ inputs=[
376
+ saved_input,
377
+ chatbot,
378
+ system_prompt,
379
+ max_new_tokens,
380
+ temperature,
381
+ top_p,
382
+ top_k,
383
+ ],
384
+ outputs=chatbot,
385
+ api_name=False,
386
+ )
387
+ )
388
+
389
+ retry_button.click(
390
+ fn=delete_prev_fn,
391
+ inputs=chatbot,
392
+ outputs=[chatbot, saved_input],
393
+ api_name=False,
394
+ queue=False,
395
+ ).then(fn=display_input, inputs=[saved_input, chatbot], outputs=chatbot, api_name=False, queue=False,).then(
396
+ fn=generate,
397
+ inputs=[
398
+ saved_input,
399
+ chatbot,
400
+ system_prompt,
401
+ max_new_tokens,
402
+ temperature,
403
+ top_p,
404
+ top_k,
405
+ ],
406
+ outputs=chatbot,
407
+ api_name=False,
408
+ )
409
+
410
+ undo_button.click(
411
+ fn=delete_prev_fn,
412
+ inputs=chatbot,
413
+ outputs=[chatbot, saved_input],
414
+ api_name=False,
415
+ queue=False,
416
+ ).then(
417
+ fn=lambda x: x,
418
+ inputs=[saved_input],
419
+ outputs=textbox,
420
+ api_name=False,
421
+ queue=False,
422
+ )
423
+
424
+ clear_button.click(
425
+ fn=lambda: ([], ""),
426
+ outputs=[chatbot, saved_input],
427
+ queue=False,
428
+ api_name=False,
429
+ )
430
+
431
+ demo.queue(max_size=20).launch(debug=True, share=True)
432
+
433
+
434
+ # if __name__ == "__main__":
435
+ # parser = argparse.ArgumentParser(description="Script to create and use an HNSW index for similarity search.")
436
+ # parser.add_argument("--input_file", help="Input file containing text chunks in a Parquet format")
437
+ # parser.add_argument("--index_file", help="HNSW index file with .bin extension")
438
+ # args = parser.parse_args()
439
+
440
+ # data_df = pd.read_parquet(args.input_file).reset_index()
441
+ # search_index = load_hnsw_index(args.index_file)
442
+ # main()
chunked_data.parquet ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d9bd56497444e3ae38e46f271d20dff959a0b5de7b0a64355dbac51f2bef660e
3
+ size 109712
requirements.txt ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ sentence_transformers
2
+ hnswlib
3
+ huggingface_hub
4
+ scipy
5
+ numpy
6
+ pandas
7
+ git+https://github.com/philschmid/easyllm
8
+ pydantic==1.10.12
9
+ transformers
search_index.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:14e38e3cb1c2b2e64977ca2ca5ded4ebff397412e228d6777304626448da8680
3
+ size 4911056