Upload 3 files

hfconstants.py

@@ -5,3 +5,18 @@ DS_SIZE_CATEGORIES = ["n<1K", "1K<n<10K", "10K<n<100K", "100K<n<1M", "1M<n<10M",
 SPACE_HARDWARES = ["cpu-basic", "zero-a10g", "cpu-upgrade", "t4-small", "l4x1", "a10g-large", "l40sx1", "a10g-small", "t4-medium", "cpu-xl", "a100-large"]
 
 SPACE_STAGES = ["RUNNING", "SLEEPING", "RUNTIME_ERROR", "PAUSED", "BUILD_ERROR", "CONFIG_ERROR", "BUILDING", "APP_STARTING", "RUNNING_APP_STARTING"]
+
+PIPELINES = ["text-generation", "text-to-image", "image-text-to-text", "fill-mask", "text-classification", "sentence-similarity", "automatic-speech-recognition",
+    "feature-extraction", "text-to-speech", "text2text-generation", "image-to-image", "text-to-video", "zero-shot-image-classification",
+    "image-classification", "image-to-video", "image-to-text", "token-classification", "translation", "time-series-forecasting", "audio-classification",
+    "object-detection", "audio-text-to-text", "zero-shot-classification", "text-to-audio", "image-segmentation", "image-feature-extraction",
+    "video-text-to-text", "image-to-3d", "any-to-any", "question-answering", "text-to-3d", "zero-shot-object-detection", "summarization",
+    "document-question-answering", "visual-question-answering", "depth-estimation", "voice-activity-detection", "audio-to-audio", "video-classification"]
+
+EMOJIS = {
+    "Pipeline": {"text-to-image": "🎨", "image-to-image": "🖼️", "text-to-speech": "🔊", "automatic-speech-recognition": "🎤",
+                 "text-to-audio": "🎵", "audio-text-to-text": "🎙️",
+                 "image-to-video": "🎞️", "text-to-video": "🎬",
+                 "text-generation": "🤖", "text2text-generation": "🤖"},
+    "SDK": {"docker": "🐳"}
+}

hfsearch.py

@@ -8,7 +8,7 @@ import pandas as pd
 import datetime
 import json
 import re
-from hfconstants import DS_SIZE_CATEGORIES, SPACE_HARDWARES, SPACE_STAGES
+from hfconstants import DS_SIZE_CATEGORIES, SPACE_HARDWARES, SPACE_STAGES, EMOJIS
 
 @spaces.GPU
 def dummy_gpu():
@@ -29,8 +29,10 @@ RESULT_ITEMS = {
     "LastMod.": [17, "str", True],
     "Library": [20, "markdown", False],
     "Pipeline": [21, "markdown", True],
+    "SDK": [24, "str", False],
     "Hardware": [25, "str", False],
     "Stage": [26, "str", False],
+    "Emoji": [35, "str", False],
     "NFAA": [40, "str", False],
 }
 
@@ -265,14 +267,12 @@ class Labels():
         return labels, label_types
 
     def get_widths(self):
-        labels = list(self.types.keys())
-        label_widths = [self.widths[s] for s in labels]
-        return label_widths
+        return self.widths.copy()
 
     def get_null_value(self, type: str):
         if type == "bool": return False
-        elif type == "number" or type == "date": return 0
-        else: return "None"
+        elif type == "number" or type == "date": return 0 #
+        else: return ""
 
 # https://huggingface.co/docs/huggingface_hub/package_reference/hf_api
 # https://huggingface.co/docs/huggingface_hub/package_reference/hf_api#huggingface_hub.ModelInfo
@@ -313,6 +313,7 @@ class HFSearchResult():
         elif isinstance(i, Collection): type = "collection"
         else: return
         self._set(type, "Type")
+        self._set("", "Emoji")
         if type in ["space", "model", "dataset"]:
             self._set(i.id, "ID")
             self._set(i.id.split("/")[0], "User")
@@ -333,9 +334,13 @@ class HFSearchResult():
                 if i.library_name is not None: self._set(i.library_name, "Library")
                 if i.pipeline_tag is not None: self._set(i.pipeline_tag, "Pipeline")
             if type == "space":
+                if i.sdk is not None: self._set(i.sdk, "SDK")
                 if i.runtime is not None:
                     self._set(i.runtime.hardware, "Hardware")
                     self._set(i.runtime.stage, "Stage")
+                if i.card_data is not None:
+                    card = i.card_data
+                    if card.title is not None: self._set(card.title, "Name")
         elif type == "paper": # https://github.com/huggingface/huggingface_hub/blob/v0.27.0/src/huggingface_hub/hf_api.py#L1428
             self._set(i.id, "ID")
             self._set(f"https://hf.co/papers/{i.id}", "URL")
@@ -479,8 +484,8 @@ class HFSearchResult():
                 sdf.loc[df["Gated"] == "auto", ["Gated"]] = 'color: dodgerblue'
                 sdf.loc[df["Gated"] == "manual", ["Gated"]] = 'color: crimson'
             if "Stage" in columns and "Hardware" in columns:
-                sdf.loc[(df["Stage"] == "RUNNING") & (df["Hardware"] != "zero-a10g") & (df["Hardware"] != "cpu-basic") & (df["Hardware"] != "None") & (df["Hardware"]), ["Hardware", "Type"]] = 'color: lime'
-                sdf.loc[(df["Stage"] == "RUNNING") & (df["Hardware"] == "zero-a10g"), ["Hardware", "Type"]] = 'color: green'
+                sdf.loc[(df["Stage"] == "RUNNING") & (df["Hardware"] != "zero-a10g") & (df["Hardware"] != "cpu-basic") & (df["Hardware"]), ["Hardware", "Type"]] = 'color: lime'
+                sdf.loc[(df["Stage"] == "RUNNING") & (df["Hardware"] == "zero-a10g"), ["Hardware", "Type"]] = 'color: limegreen'
                 sdf.loc[(df["Type"] == "space") & (df["Stage"] != "RUNNING")] = 'opacity: 0.5'
                 sdf.loc[(df["Type"] == "space") & (df["Stage"] != "RUNNING"), ["Type"]] = 'color: crimson'
                 sdf.loc[df["Stage"] == "RUNNING", ["Stage"]] = 'color: lime'
@@ -492,16 +497,28 @@ class HFSearchResult():
             return sdf
         
         def id_to_md(df: pd.DataFrame):
-            if df["Type"] == "collection": return f'[{df["User"]}: {df["Name"]}]({df["URL"]})'
-            elif df["Type"] == "paper": return f'[{df["Name"]} (arxiv:{df["ID"]})]({df["URL"]})'
-            else: return f'[{df["ID"]}]({df["URL"]})'
+            if df["Type"] == "collection": return f'[{df["User"]}/{df["Name"]}]({df["URL"]}){df["Emoji"]}'
+            elif df["Type"] == "space": return f'[{df["Name"]} ({df["ID"]})]({df["URL"]}){df["Emoji"]}'
+            elif df["Type"] == "paper": return f'[{df["Name"]} (arxiv:{df["ID"]})]({df["URL"]}){df["Emoji"]}'
+            else: return f'[{df["ID"]}]({df["URL"]}){df["Emoji"]}'
+        
+        def to_emoji(df: pd.DataFrame, label: str, key: str, emoji: str):
+            if df[label] == key: return f'{df["Emoji"]}{emoji}' if df["Emoji"] else f' {emoji}'
+            else: return df["Emoji"]
+        
+        def apply_emoji_df(df: pd.DataFrame):
+            for label, v in EMOJIS.items():
+                if label not in df.columns: continue
+                for key, emoji in v.items():
+                    df["Emoji"] = df.apply(to_emoji, axis=1, label=label, key=key, emoji=emoji)
+            return df
 
         def format_md_df(df: pd.DataFrame):
             df["ID"] = df.apply(id_to_md, axis=1)
             return df
         
         hide_labels = [l for l in labels if l in self.hide_labels or l not in self.show_labels]
-        df = format_md_df(pd.DataFrame(dflist, columns=labels))
+        df = format_md_df(apply_emoji_df(pd.DataFrame(dflist, columns=labels)))
         ref_df = df.copy()
         df = df.drop(hide_labels, axis=1).style.apply(highlight_df, axis=None, df=ref_df)
         return df
@@ -554,7 +571,8 @@ class HFSearchResult():
     def get_gr_df(self):
         df, labels, label_types = self.get()
         widths = self.labels.get_widths()
-        return gr.update(type="pandas", value=df, headers=labels, datatype=label_types, column_widths=widths, wrap=True)
+        column_widths = [widths[l] for l in labels]
+        return gr.update(type="pandas", value=df, headers=labels, datatype=label_types, column_widths=column_widths, wrap=True)
 
     def get_gr_hide_labels(self):
         return gr.update(choices=self.labels.get()[0], value=[], visible=True)