-
Notifications
You must be signed in to change notification settings - Fork 5
/
main.py
76 lines (60 loc) · 1.8 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import argparse
from config import Config
from loader.reader import DataReader
from loader.index import Index
def load_index_and_run_engine(config):
index = Index(
storage_dir=config.storage_dir,
model_type=config.model_type
).load()
chat_engine = index.as_chat_engine(similarity_top_k=3, chat_mode="context")
while True:
query_str = input("\nQ: ")
if query_str == "exit":
print("Cya!")
break
if query_str.strip() == "":
continue
streaming_response = chat_engine.stream_chat(query_str)
for token in streaming_response.response_gen:
print(token, end="")
def load_data_and_store_index(config):
data_loader = DataReader(data_dir=config.data_dir)
docs = data_loader.load()
_ = Index(
storage_dir=config.storage_dir,
documents=docs,
model_type=config.model_type
).persist()
def get_arg_parser():
parser = argparse.ArgumentParser(
prog='llm',
description='LLM prompt augmentation to chat with your documents',
usage='%(prog)s [options]'
)
parser.add_argument(
'-L', '--load-data',
help='Ingest your documents and create index',
action='store_true',
required=False
)
parser.add_argument(
'-Q', '--query-data',
help='Chat with your documents. Type "exit" to quit',
action='store_true',
required=False
)
return parser
def main():
parser = get_arg_parser()
args = parser.parse_args()
config = Config()
if args.load_data:
load_data_and_store_index(config)
elif args.query_data:
load_index_and_run_engine(config)
else:
print("Invalid arguments")
parser.print_help()
if __name__ == "__main__":
main()