Ask-ANRG / model_main.py
FloraJ's picture
update model
de3c2ee
raw
history blame
2.45 kB
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
# Get & embed user's question
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
"""