Clean up code, fix catalog fetching

ersilia-os · Sep 22, 2024 · 0ee5b45 · 0ee5b45
1 parent 6b6c6d3
commit 0ee5b45
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 32 deletions.
diff --git a/ersilia/hub/content/card.py b/ersilia/hub/content/card.py
@@ -736,8 +736,6 @@ def _load_data(self, model_id):
         if os.path.exists(card_path):
             with open(card_path, "r") as f:
                 card = json.load(f)
-                if "card" in card:
-                    return card["card"]
             return card
         else:
             return None

diff --git a/ersilia/hub/content/catalog.py b/ersilia/hub/content/catalog.py
@@ -73,41 +73,33 @@ def _is_eos(self, s):
             if not self.mi.is_test(s):
                 return True
         return False
+
+    def _get_item(self, card, item):
+        if "card" in card:
+            card = card["card"]
+        if item.lower() in card:
+            return card[item.lower()]
+        elif item.capitalize() in card:
+            return card[item.capitalize()]
+        elif item.title() in card:
+            return card[item.title()]
+        else:
+            return None
 
     def _get_title(self, card):
-        if "title" in card:
-            return card["title"]
-        if "Title" in card:
-            return card["Title"]
-        return None
+        return self._get_item(card, "title")
 
     def _get_slug(self, card):
-        if "slug" in card:
-            return card["slug"]
-        if "Slug" in card:
-            return card["Slug"]
-        return None
+        return self._get_item(card, "slug")
 
     def _get_status(self, card):
-        if "status" in card:
-            return card["status"]
-        if "Status" in card:
-            return card["Status"]
-        return None
+        return self._get_item(card, "status")
 
     def _get_input(self, card):
-        if "input" in card:
-            return card["input"][0]
-        if "Input" in card:
-            return card["Input"][0]
-        return None
+        return self._get_item(card, "input")[0]
 
     def _get_output(self, card):
-        if "output" in card:
-            return card["output"][0]
-        if "Output" in card:
-            return card["Output"][0]
-        return None
+        return self._get_item(card, "output")[0]
 
     def _get_model_source(self, model_id):
         model_source_file = os.path.join(self._model_path(model_id), MODEL_SOURCE_FILE)
@@ -208,11 +200,11 @@ def local(self):
                 if not self._is_eos(model_id):
                     continue
                 card = mc.get(model_id)
-                slug = self._get_slug(card["card"])
-                title = self._get_title(card["card"])
-                status = self._get_status(card["card"])
-                inputs = self._get_input(card["card"])
-                output = self._get_output(card["card"])
+                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]]

diff --git a/ersilia/hub/content/slug.py b/ersilia/hub/content/slug.py
@@ -53,6 +53,8 @@ def _remote_decode(self, model_id):
         if res is None:
             return None
         else:
+            if "card" in res:
+                res = res["card"]
             return res["Slug"].strip()
 
     def decode(self, model_id):