Update app.py

app.py

@@ -3,8 +3,9 @@ import gradio as gr
 import pandas as pd
 from huggingface_hub import InferenceClient
 from threading import Timer
+from tqdm import tqdm
 
-HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
+HUGGINGFACE_TOKEN =os.environ.get("HUGGINGFACE_TOKEN")
 def get_available_free(use_cache = False):
     if use_cache:
         if os.path.exists(str(os.getcwd())+"/data.csv"):
@@ -25,15 +26,13 @@ def get_available_free(use_cache = False):
     }
 
     all_models = list(set(models + models_vision + models_others))
-    print(all_models)
-    for m in all_models:
+    for m in tqdm(all_models):
         text_available = False
         chat_available = False
         vision_available = False
         if m in models_vision:
             vision_available = True
         pro_sub = False
-        print(m)
         try:
             InferenceClient(m, timeout=10, token=HUGGINGFACE_TOKEN).text_generation("Hi.", max_new_tokens=1)
             text_available = True
@@ -84,14 +83,23 @@ def update_data(use_cache = False):
     
     return df
 
-def display_table(search_query="", use_cache = False):
+def display_table(search_query="", filters=[], use_cache=False):
     df = update_data(use_cache)
     search_query = str(search_query)
+    
     if search_query:
         filtered_df = df[df["Model"].str.contains(search_query, case=False)]
     else:
         filtered_df = df
-    
+
+    if filters:
+        if "Free" in filters:
+            filtered_df = filtered_df[filtered_df["API"] == "Free"]
+        if "Text Completion" in filters:
+            filtered_df = filtered_df[filtered_df["Text Completion"] == "✓"]
+        if "Chat Completion" in filters:
+            filtered_df = filtered_df[filtered_df["Chat Completion"] == "✓"]
+
     styled_df = filtered_df.style.apply(apply_row_styles, axis=1, subset=["Model", "API", "Text Completion", "Chat Completion", "Vision"])
     return styled_df
 
@@ -149,9 +157,14 @@ with gr.Blocks() as demo:
     gr.Markdown("## HF Serverless LLM Inference API Status")
     gr.Markdown(description)
     search_box = gr.Textbox(label="Search for a model", placeholder="Type model name here...")
+    filter_box = gr.CheckboxGroup(choices=["Free", "Text Completion", "Chat Completion"], label="Filters")
     table = gr.Dataframe(value=display_table(use_cache=True), headers="keys")
-    
-    search_box.change(fn=search_models, inputs=search_box, outputs=table)
+
+    def update_filters(query, filters):
+        return search_models(query, filters, use_cache=True)
+
+    search_box.change(fn=update_filters, inputs=[search_box, filter_box], outputs=table)
+    filter_box.change(fn=update_filters, inputs=[search_box, filter_box], outputs=table)
     
     def update_every_two_hours(first_run):
         search_models(search_box.value, first_run)