add some logging to welcome fns

TheDataGuild · Sep 28, 2023 · e89ce63 · e89ce63
1 parent 6533a84
commit e89ce63
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/mind_palace/welcome.py b/mind_palace/welcome.py
@@ -2,6 +2,9 @@
 
 from docs import Section
 from llama_index.llms import ChatMessage, OpenAI
+from measure import time_function
+
+from mind_palace.measure import time_function
 
 
 def parse_abstracts(nodes) -> List[str]:
@@ -29,7 +32,9 @@ def _summarize_prompt(abstracts: List[str]):
     }
 
 
+@time_function
 def summarize(gpt_model, texts: List[str]) -> str:
+    print(f"Summarizing {len(texts)} pieces of texts with {gpt_model} model.")
     prompt = _summarize_prompt(texts)
     messages = [
         ChatMessage(role="system", content=prompt["system"]),
@@ -51,7 +56,11 @@ def _extract_keywords_output(message):
     return [keyword.split(". ")[1] for keyword in keywords]
 
 
+@time_function
 def extract_keywords(gpt_model, texts: List[str]) -> List[str]:
+    print(
+        f"Extracting keywords from {len(texts)} pieces of texts with {gpt_model} model."
+    )
     prompt = _extract_keywords_prompt(texts)
     messages = [
         ChatMessage(role="system", content=prompt["system"]),