logging -> loguru: final stage

crazy_functions/rag_fns/llama_index_worker.py

@@ -1,6 +1,7 @@
 import llama_index
 import os
 import atexit
+from loguru import logger
 from typing import List
 from llama_index.core import Document
 from llama_index.core.schema import TextNode
@@ -41,14 +42,14 @@ class SaveLoad():
         return True
 
     def save_to_checkpoint(self, checkpoint_dir=None):
-        print(f'saving vector store to: {checkpoint_dir}')
+        logger.info(f'saving vector store to: {checkpoint_dir}')
         if checkpoint_dir is None: checkpoint_dir = self.checkpoint_dir
         self.vs_index.storage_context.persist(persist_dir=checkpoint_dir)
 
     def load_from_checkpoint(self, checkpoint_dir=None):
         if checkpoint_dir is None: checkpoint_dir = self.checkpoint_dir
         if self.does_checkpoint_exist(checkpoint_dir=checkpoint_dir):
-            print('loading checkpoint from disk')
+            logger.info('loading checkpoint from disk')
             from llama_index.core import StorageContext, load_index_from_storage
             storage_context = StorageContext.from_defaults(persist_dir=checkpoint_dir)
             self.vs_index = load_index_from_storage(storage_context, embed_model=self.embed_model)
@@ -85,9 +86,9 @@ class LlamaIndexRagWorker(SaveLoad):
         self.vs_index.storage_context.index_store.to_dict()
         docstore = self.vs_index.storage_context.docstore.docs
         vector_store_preview = "\n".join([ f"{_id} | {tn.text}" for _id, tn in docstore.items() ])
-        print('\n++ --------inspect_vector_store begin--------')
-        print(vector_store_preview)
-        print('oo --------inspect_vector_store end--------')
+        logger.info('\n++ --------inspect_vector_store begin--------')
+        logger.info(vector_store_preview)
+        logger.info('oo --------inspect_vector_store end--------')
         return vector_store_preview
 
     def add_documents_to_vector_store(self, document_list):
@@ -125,5 +126,5 @@ class LlamaIndexRagWorker(SaveLoad):
         
     def generate_node_array_preview(self, nodes):
         buf = "\n".join(([f"(No.{i+1} | score {n.score:.3f}): {n.text}" for i, n in enumerate(nodes)]))
-        if self.debug_mode: print(buf)
+        if self.debug_mode: logger.info(buf)
         return buf

crazy_functions/rag_fns/milvus_worker.py

@@ -2,6 +2,7 @@ import llama_index
 import os
 import atexit
 from typing import List
+from loguru import logger
 from llama_index.core import Document
 from llama_index.core.schema import TextNode
 from request_llms.embed_models.openai_embed import OpenAiEmbeddingModel
@@ -44,14 +45,14 @@ class MilvusSaveLoad():
         return True
 
     def save_to_checkpoint(self, checkpoint_dir=None):
-        print(f'saving vector store to: {checkpoint_dir}')
+        logger.info(f'saving vector store to: {checkpoint_dir}')
         # if checkpoint_dir is None: checkpoint_dir = self.checkpoint_dir
         # self.vs_index.storage_context.persist(persist_dir=checkpoint_dir)
 
     def load_from_checkpoint(self, checkpoint_dir=None):
         if checkpoint_dir is None: checkpoint_dir = self.checkpoint_dir
         if self.does_checkpoint_exist(checkpoint_dir=checkpoint_dir):
-            print('loading checkpoint from disk')
+            logger.info('loading checkpoint from disk')
             from llama_index.core import StorageContext, load_index_from_storage
             storage_context = StorageContext.from_defaults(persist_dir=checkpoint_dir)
             try:
@@ -101,7 +102,7 @@ class MilvusRagWorker(MilvusSaveLoad, LlamaIndexRagWorker):
             vector_store_preview = "\n".join(
                 [f"{node.id_} | {node.text}" for node in dummy_retrieve_res]
             )
-        print('\n++ --------inspect_vector_store begin--------')
-        print(vector_store_preview)
-        print('oo --------inspect_vector_store end--------')
+        logger.info('\n++ --------inspect_vector_store begin--------')
+        logger.info(vector_store_preview)
+        logger.info('oo --------inspect_vector_store end--------')
         return vector_store_preview