Merge pull request #34 from TheDataGuild/refactor/measure-time-decorator

Refactor/measure time decorator

mind_palace/index.py

@@ -1,15 +1,21 @@
 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"):
+    embed_model = "local"
     service_context = li.ServiceContext.from_defaults(
         llm=OpenAI(model=model),
-        embed_model="local"
+        embed_model=embed_model
         # Q: how much does a better embedding model help?
     )
 
-    print(f"Creating an index with {len(nodes)} number of nodes using model {model}")
+    print(
+        f"Creating an index with {len(nodes)} number of nodes using embedding model {embed_model} and querying LLM {model}"
+    )
     index = li.VectorStoreIndex(nodes, service_context=service_context)
 
     return index

mind_palace/measure.py

@@ -2,8 +2,14 @@
 
 
 def time_function(func):
-    start_time = time.time()
-    result = func()
-    end_time = time.time()
-    elapsed_time = end_time - start_time
-    return result, elapsed_time
+    """Decorator to time and print timing of a single function call"""
+
+    def timed_func(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        elapsed_sec = end_time - start_time
+        print(f"Elapsed time calling {func.__name__}(): {elapsed_sec:.1f} seconds")
+        return result
+
+    return timed_func

tests/end_to_end/test_query.py

@@ -1,12 +1,18 @@
-from tests.context import extract, index, measure
+from tests.context import extract, index
 import llama_index as li
+import pytest
 
 
 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)
+
+
+@pytest.mark.skip(reason="This test takes a minute to run")
+def test_bootstrap_index():
+    xml_dir = "./resources/xmls/12-pdfs-from-steve-aug-22/"
+    vector_index = bootstrap_index(xml_dir)
+    assert vector_index is not None
 
 
 def persist_index(vector_index):