Introduce the ability to list datasets and models from the CLI

src/cli.py

@@ -3,12 +3,14 @@
 import matplotlib.pyplot as plt
 
 from src.common.data_models import EmbeddingDefinition
+from src.dataset.provider import dataset_provider
 from src.evaluation import (
     LinearProbeClassifier,
     WeightedKNNClassifier,
     ZeroShotClassifier,
 )
 from src.evaluation.evaluator import export_evaluation_to_csv, run_evaluation
+from src.models.CLIP_model import CLIPModel
 from src.plotting.animation import build_animation, save_animation_to_file
 from src.utils import read_all_cached_embeddings
 
@@ -70,12 +72,31 @@ def animate_embeddings(model_datasets: list[str]):
     plt.show()
 
 
+def list_command(args):
+    list_models_datasets(args.all)
+
+
+def list_models_datasets(all: bool = False):
+    if all:
+        datasets = dataset_provider.list_dataset_names()
+        models = CLIPModel.list_models()
+        print(f"Available datasets are: {', '.join(datasets)}")
+        print(f"Available models are: {', '.join(models)}")
+        return
+    else:
+        defns: list[EmbeddingDefinition] = [d for k, v in read_all_cached_embeddings().items() for d in v]
+        print(f"Available model_dataset pairs: {', '.join([str(defn) for defn in defns])}")
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Multiclip builder and Benchmarker")
     subparsers = parser.add_subparsers()
 
     parser_build = subparsers.add_parser("build", help="build embeddings")
-    parser_build.add_argument("model_dataset", type=str, help="model, dataset pair delimited by model/dataset")
+
+    parser_build.add_argument(
+        "model_dataset", type=str, help="model, dataset pair delimited by model/dataset", default=None
+    )
     parser_build.set_defaults(func=build_command)
 
     parser_evaluate = subparsers.add_parser("evaluate", help="evaluate embeddings")
@@ -99,5 +120,8 @@ def animate_embeddings(model_datasets: list[str]):
 
     parser_animation.set_defaults(func=animate_command)
 
+    parser_list = subparsers.add_parser("list", help="List available dataset, model pairs")
+    parser_list.add_argument("--all", action="store_true", help="List all available models and dataset")
+    parser_list.set_defaults(func=list_command)
     args = parser.parse_args()
     args.func(args)

src/dataset/provider.py

@@ -14,6 +14,9 @@ def get_dataset(self, title: str) -> Dataset:
         source, kwargs = self._datasets[title]
         return source(title, **kwargs)
 
+    def list_dataset_names(self) -> list[str]:
+        return list(self._datasets.keys())
+
 
 dataset_provider = DatasetProvider()
 dataset_provider.register_dataset("LungCancer4Types", HFDataset, title_in_source="Kabil007/LungCancer4Types")
@@ -22,7 +25,7 @@ def get_dataset(self, title: str) -> Dataset:
 dataset_provider.register_dataset("Alzheimer-MRI", HFDataset, title_in_source="Falah/Alzheimer_MRI")
 dataset_provider.register_dataset("Skin-cancer", HFDataset, title_in_source="marmal88/skin_cancer")
 dataset_provider.register_dataset(
-    "chest-xray-classification", HFDataset, title_in_source="tpakov/chest-xray-classification"
+    "chest-xray-classification", HFDataset, title_in_source="trpakov/chest-xray-classification"
 )
 dataset_provider.register_dataset(
     "sports-classification", HFDataset, title_in_source="HES-XPLAIN/SportsImageClassification"

src/models/CLIP_model.py

@@ -52,3 +52,7 @@ def process_fn(batch):
             return batch
 
         return process_fn
+
+    @classmethod
+    def list_models(self) -> list[str]:
+        return list(OPTIONS.keys())