|
import argparse |
|
|
|
from utils import get_embeddings, search_document_annoy, \ |
|
answer_with_gpt3_with_function_calls, transform_user_question, debug_print |
|
|
|
def main(args): |
|
""" |
|
- Get & embed user's question |
|
- Find top choices of documents based on embedding search |
|
- Generate response |
|
""" |
|
top_k = args.top_k |
|
annoy_metric = args.annoy_metric |
|
|
|
|
|
while True: |
|
print("Hi! What question do you have for ANGR? press 0 to exit") |
|
user_input = input() |
|
if user_input == '0': |
|
break |
|
|
|
assert top_k > 0, 'k must be a integer greater than 0' |
|
|
|
if args.user_query_preprocess: |
|
chatgpt_question = transform_user_question(user_input, args.model) |
|
else: |
|
chatgpt_question = user_input |
|
debug_print("chatgpt_question: ", chatgpt_question) |
|
|
|
try: |
|
user_q_embedding = get_embeddings(chatgpt_question) |
|
document = search_document_annoy(user_q_embedding, top_k=top_k, metric=annoy_metric) |
|
reply = answer_with_gpt3_with_function_calls(document, user_input, args.model) |
|
print(reply) |
|
except Exception as e: |
|
print(e) |
|
print("Error when trying to get embedding for the user query. Please try with a shorter question.") |
|
|
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--top_k', default=3, type=int, help="Number of documents to retrieve") |
|
parser.add_argument('--annoy_metric', default='angular', |
|
choices=['angular', 'euclidean', 'manhattan', 'hamming', 'dot'], |
|
help="metric for annoy algorithm") |
|
parser.add_argument('--model', default="gpt-3.5-turbo") |
|
parser.add_argument('--user_query_preprocess', default=False, help="Whether to ask ChatGPT to generate an additional query for itself based on user's question") |
|
parser.add_argument('--debug', default=False, help="define whether to print the debug statement") |
|
args = parser.parse_args() |
|
|
|
main(args) |
|
|
|
""" |
|
python main.py --top_k 3 --annoy_metric dot --user_query_preprocess True |
|
|
|
Questions: |
|
- list math-related publications |
|
- list all publications |
|
- Give me a publication written by J Coleman |
|
- Could you please provide me with a publication authored by J Coleman? |
|
- Can you find a publication authored by J Coleman? |
|
- Give me the link to Jared Coleman's homepage |
|
|
|
""" |
|
|