apply @time_function

TheDataGuild · Sep 23, 2023 · bd9e9f1 · bd9e9f1
1 parent c864aae
commit bd9e9f1
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git 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
@@ -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):