From bd9e9f175723019d549a4ddf4499e5f4e931eeab Mon Sep 17 00:00:00 2001 From: Paul Lam Date: Sun, 24 Sep 2023 06:29:50 +0900 Subject: [PATCH] apply @time_function --- mind_palace/index.py | 3 +++ tests/end_to_end/test_query.py | 6 ++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/mind_palace/index.py b/mind_palace/index.py index a1d5b59..ea581cf 100644 --- a/mind_palace/index.py +++ b/mind_palace/index.py @@ -1,7 +1,10 @@ from llama_index.llms import OpenAI import llama_index as li +from measure import time_function + +@time_function def index_nodes(nodes, model="gpt-3.5-turbo"): service_context = li.ServiceContext.from_defaults( llm=OpenAI(model=model), diff --git a/tests/end_to_end/test_query.py b/tests/end_to_end/test_query.py index b56c81a..d877bf8 100644 --- a/tests/end_to_end/test_query.py +++ b/tests/end_to_end/test_query.py @@ -1,12 +1,10 @@ -from tests.context import extract, index, measure +from tests.context import extract, index import llama_index as li def bootstrap_index(xml_dir): nodes = extract.seed_nodes(xml_dir) - vector_index, elapsed_time = measure.time_function(lambda: index.index_nodes(nodes)) - print(f"Elapsed time {elapsed_time:.1f} seconds: Indexed {len(nodes)} nodes") - return vector_index + return index.index_nodes(nodes) def persist_index(vector_index):