feat(available_dataset): allow saving output to a file via print redirection

Author: javadr

rankify/dataset/dataset.py

@@ -657,7 +657,7 @@ def __getitem__(self,idx) -> Document:
 
 
     @staticmethod
-    def available_dataset() -> None:
+    def available_dataset(file=None) -> None:
         """
         Prints information about available datasets.
 
@@ -666,7 +666,7 @@ def available_dataset() -> None:
             Dataset.available_dataset()
             ```
         """
-        get_datasets_info()
+        get_datasets_info(file=file)
 
 
 

rankify/utils/dataset/utils.py

@@ -2,7 +2,7 @@
 import pandas as pd
 from prettytable import PrettyTable
 
-def get_datasets_info():
+def get_datasets_info(file=None):
     table = PrettyTable(['Retriever', 'Dataset', 'Original ext', 'Compressed','Desc','URL'])
     for retriever, datasets in HF_PRE_DEFIND_DATASET.items():
         for dataset_name, dataset_info in datasets.items():
@@ -17,4 +17,4 @@ def get_datasets_info():
             }
             table.add_row(flattened_entry.values())
 
-    print(table)
+    print(table, file=file)