Refactor Catalog command

ersilia/cli/commands/catalog.py

@@ -29,25 +29,15 @@ def catalog_cmd():
     # Example usage: ersilia catalog
     @ersilia_cli.command(help="List a catalog of models")
     @click.option(
-        "-l",
-        "--local",
-        is_flag=True,
-        default=False,
-        help="Show catalog of models available in the local computer",
-    )
-    @click.option(
-        "-h",
-        "--hub",
+        "-h/-l",
+        "--hub/--local",
         is_flag=True,
         default=False,
         help="Show catalog of models available in the model hub",
     )
     @click.option(
         "--file_name", "-f", default=None, type=click.STRING, help="Catalog file name"
     )
-    @click.option(
-        "--browser", is_flag=True, default=False, help="Show catalog in the browser"
-    )
     @click.option(
         "--more/--less",
         is_flag=True,
@@ -66,28 +56,28 @@ def catalog_cmd():
         required=False,
     )
     @click.option(
-        "--as-table",
+        "-j/-t",
+        "--as-json/--as-table",
         is_flag=True,
         default=False,
         help="Show catalog in table format",
     )
     def catalog(
-        local=False,
         hub=False,
         file_name=None,
         browser=False,
         more=False,
         card=False,
         model=None,
-        as_table=False,
+        as_json=False,
     ):
         if card and not model:
             click.echo(
                 click.style("Error: --card option requires a model ID", fg="red"),
                 err=True,
             )
             return
-        if card and model:
+        elif card and model:
             try:
                 mc = ModelCard()
                 model_metadata = mc.get(model, as_json=True)
@@ -104,52 +94,29 @@ def catalog(
             except Exception as e:
                 click.echo(click.style(f"Error fetching model metadata: {e}", fg="red"))
             return
+        else:
+            mc = ModelCatalog(less=not more)
 
-        # The idea here is to deter the user from running ersilia catalog --local --hub
-        if local and hub:
-            click.echo(
-                click.style(
-                    "Error: Cannot show local and hub models together", fg="red"
-                ),
-                err=True,
-            )
-            return
-
-        mc = ModelCatalog()
-        mc.only_identifier = False if more else True
-
-        if hub:
-            if browser:
-                mc.airtable()
-                return
-
-            catalog_table = mc.hub()
-
-        else:  # This will work even if the user doesn't explicitly specify the --local flag
-            if browser:
-                click.echo(
-                    click.style(
-                        "Error: Cannot show local models in the browser.\nPlease use the --hub option to see models in the browser.",
-                        fg="red",
-                    )
-                )
-                return
-            catalog_table = mc.local()
-            if not catalog_table.data:
-                click.echo(
-                    click.style(
-                        "No local models available. Please fetch a model by running 'ersilia fetch' command",
-                        fg="red",
+            if hub:
+                catalog_table = mc.hub()
+            else:
+                catalog_table = mc.local()
+                if not catalog_table.data:
+                    click.echo(
+                        click.style(
+                            "No local models available. Please fetch a model by running 'ersilia fetch' command",
+                            fg="red",
+                        )
                     )
+                    return
+            if file_name is None:
+                catalog = (
+                    catalog_table.as_json() if as_json else catalog_table.as_table()
                 )
-                return
-
-        if file_name is None:
-            catalog = catalog_table.as_table() if as_table else catalog_table.as_json()
-        else:
-            catalog_table.write(file_name)
-            catalog = None
+            else:
+                catalog_table.write(file_name)
+                catalog = None
 
-        click.echo(catalog)
+            click.echo(catalog)
 
     return catalog

ersilia/cli/commands/example.py

@@ -37,7 +37,6 @@ def example_cmd():
     def example():
         pass
 
-
     @example.command()
     @click.argument("model", required=False, default=None, type=click.STRING)
     @click.option("--n_samples", "-n", default=5, type=click.INT)
@@ -50,6 +49,15 @@ def inputs(model, n_samples, file_name, simple, predefined):
         else:
             session = Session(config_json=None)
             model_id = session.current_model_id()
+            if model_id is None:
+                click.echo(
+                    click.style(
+                        "Error: No model id given and no model found running in this shell.",
+                        fg="red",
+                    ),
+                    err=True,
+                )
+                return
         eg = ExampleGenerator(model_id=model_id)
         if file_name is None:
             echo(
@@ -61,7 +69,6 @@ def inputs(model, n_samples, file_name, simple, predefined):
         else:
             eg.example(n_samples, file_name, simple, try_predefined=predefined)
 
-
     @example.command()
     @click.option("--n_samples", "-n", default=5, type=click.INT)
     @click.option("--file_name", "-f", default=None, type=click.STRING)
@@ -71,4 +78,4 @@ def models(n_samples, file_name):
         if file_name is None:
             echo(json.dumps(sampler.sample(n_samples), indent=4))
         else:
-            sampler.sample(n_samples, file_name)
+            sampler.sample(n_samples, file_name)

ersilia/hub/content/catalog.py

@@ -180,10 +180,19 @@ def __repr__(self):
 
 
 class ModelCatalog(ErsiliaBase):
-    def __init__(self, config_json=None, only_identifier=True):
+    LESS_FIELDS = ["Identifier", "Slug"]
+    MORE_FIELDS = LESS_FIELDS + [
+        "Title",
+        "Task",
+        "Input Shape",
+        "Output",
+        "Output Shape",
+    ]
+
+    def __init__(self, config_json=None, less=True):
         ErsiliaBase.__init__(self, config_json=config_json)
         self.mi = ModelIdentifier()
-        self.only_identifier = only_identifier
+        self.less = less
 
     def _is_eos(self, s):
         if self.mi.is_valid(s):
@@ -218,7 +227,8 @@ def _get_input(self, card):
     def _get_output(self, card):
         return self._get_item(card, "output")[0]
 
-    def _get_model_source(self, model_id):
+    def _get_model_source(self, card):
+        model_id = self._get_item(card, "identifier")
         model_source_file = os.path.join(self._model_path(model_id), MODEL_SOURCE_FILE)
         if os.path.exists(model_source_file):
             with open(model_source_file) as f:
@@ -238,80 +248,52 @@ def airtable(self):
         if webbrowser:
             webbrowser.open("https://airtable.com/shrUcrUnd7jB9ChZV")  # TODO Hardcoded
 
+    def _get_catalog(self, columns: list, model_cards: list):
+        """Get the catalog of models"""
+        R = []
+        columns = ["Index"] + columns
+
+        idx = 0
+        for card in model_cards:
+            status = self._get_status(card)
+            if status == "In Progress":
+                continue
+            idx += 1
+            r = [idx]
+            for field in columns[1:]:
+                if field == "Model Source":
+                    r += [self._get_model_source(card)]
+                else:
+                    r += [self._get_item(card, field)]
+            # Sort R by Identifier
+
+            R += [r]
+        R = sorted(R, key=lambda x: x[0])
+        return CatalogTable(data=R, columns=columns)
+
     def hub(self):
         """List models available in Ersilia model hub from the S3 JSON"""
-        mc = ModelCard()
         ji = JsonModelsInterface()
         models = ji.items_all()
-        R = []
-        for model in models:
-            status = self._get_status(model)
-            if status == "In Progress":
-                continue
-
-            identifier = self._get_item(model, "identifier")
-            if self.only_identifier:
-                R += [[identifier]]
-            else:
-                slug = self._get_slug(model)
-                title = self._get_title(model)
-                R += [[identifier, slug, title]]
-
-        columns = (
-            ["Identifier"] if self.only_identifier else ["Identifier", "Slug", "Title"]
-        )
-        return CatalogTable(R, columns=columns)
+        columns = self.LESS_FIELDS if self.less else self.MORE_FIELDS
+        table = self._get_catalog(columns, models)
+        return table
 
     def local(self):
         """List models available locally"""
         mc = ModelCard()
-        R = []
-        logger.debug("Looking for models in {0}".format(self._bundles_dir))
-        if self.only_identifier:
-            R = []
-            for model_id in os.listdir(self._bundles_dir):
-                if not self._is_eos(model_id):
-                    continue
-                R += [[model_id]]
-            columns = ["Identifier"]
-        else:
-            for model_id in os.listdir(self._bundles_dir):
-                if not self._is_eos(model_id):
-                    continue
-                card = mc.get(model_id)
-                slug = self._get_slug(card)
-                title = self._get_title(card)
-                status = self._get_status(card)
-                inputs = self._get_input(card)
-                output = self._get_output(card)
-                model_source = self._get_model_source(model_id)
-                service_class = self._get_service_class(card)
-                R += [
-                    [
-                        model_id,
-                        slug,
-                        title,
-                        status,
-                        inputs,
-                        output,
-                        model_source,
-                        service_class,
-                    ]
-                ]
-            columns = [
-                "Identifier",
-                "Slug",
-                "Title",
-                "Status",
-                "Input",
-                "Output",
-                "Model Source",
-                "Service Class",
-            ]
-        logger.info("Found {0} models".format(len(R)))
-        if len(R) == 0:
-            return CatalogTable(data=[], columns=columns)
-        return CatalogTable(data=R, columns=columns)
+        columns = self.LESS_FIELDS if self.less else self.MORE_FIELDS+["Model Source"]
+        cards = []
+        for model_id in os.listdir(self._bundles_dir):
+            if not self._is_eos(model_id):
+                continue
+            card = mc.get(model_id)
+            if not card:
+                continue
+            cards += [card]
+        table = self._get_catalog(columns, cards)
+        return table
+
 
     def bentoml(self):
         """List models available as BentoServices"""

ersilia/io/input.py

@@ -271,7 +271,7 @@ def example(self, n_samples, file_name, simple, try_predefined):
         if try_predefined is True and file_name is not None:
             self.logger.debug("Trying with predefined input")
             predefined_available = self.predefined_example(file_name)
-        
+
         if predefined_available:
             with open(file_name, "r") as f:
                 return f.read()

ersilia/utils/exceptions_utils/exceptions.py

@@ -93,7 +93,9 @@ def run_from_terminal(self):
         output_file = os.path.join(framework_dir, "example_output.csv")
         tmp_folder = make_temp_dir(prefix="ersilia-")
         log_file = os.path.join(tmp_folder, "terminal.log")
-        run_command("ersilia example inputs {0} -n 3 -f {1}".format(self.model_id, input_file))
+        run_command(
+            "ersilia example inputs {0} -n 3 -f {1}".format(self.model_id, input_file)
+        )
         cmd = "bash {0} {1} {2} {3} 2>&1 | tee -a {4}".format(
             exec_file, framework_dir, input_file, output_file, log_file
         )