Update app.py

app.py

@@ -5,33 +5,79 @@ if root_dir not in sys.path:
     sys.path.append(root_dir)
 
 import gradio as gr
-import os
-import requests
 
-from gradio.themes import *
+from utils import set_text_bg_color
+from loop_retrieve_cards import get_models, get_datasets, get_readme_dict
+
+
+def match_card(input: str, card_id: str, card_type: str) -> str:
+    """
+    Search the input in a card. If the input string is contained in the card_id or its README, display this card.
+
+    Args:
+        input:  Input string
+        card_id: HuggingFace card id
+        card_type: Type of card, either "model" or "dataset"
+    """
+    display_str = ""
+    readme_dict = get_readme_dict()
+
+    if input.lower() in card_id.lower() or input.lower() in readme_dict[card_id].lower():
+        # Add card id
+        if card_type == "model":
+            display_str += f"## [{set_text_bg_color(input, card_id)}](https://huggingface.co/{card_id})\n\n"
+        else:
+            display_str += f"## [{set_text_bg_color(input, card_id)}](https://huggingface.co/datasets/{card_id})\n\n"
+
+        # Highlight lines that contain the input string
+        show_lines = []
+        for line in readme_dict[card_id].split("\n"):
+            if input.lower() in line.lower() and "<!--" not in line:
+                show_lines.append(set_text_bg_color(input, line))
+
+        # Add README
+        display_str += "\n\n".join(show_lines)
+
+        # Add a separator
+        display_str = f"\n\n{display_str}\n\n---\n\n"
+
+        # In case that the keyword is only contained in comments
+        if input.lower() not in card_id.lower() and len(show_lines) == 0:
+            display_str = ""
+
+    return display_str
+
+
+def show_card_info(input: str):
+    retrieval_str = ""
+
+    if input != "":
+        # Search models
+        retrieval_str += "# Models\n\n"
+        for model in get_models():
+            retrieval_str += match_card(input, model, "model")
+
+        # Search datasets
+        retrieval_str += "# Datasets\n\n"
+        for dataset in get_datasets():
+            retrieval_str += match_card(input, dataset, "dataset")
+
+    return gr.Markdown(retrieval_str, visible=True)
 
 
 # Build demo
 with gr.Blocks(title="SaprotHub", fill_width=True) as demo:
-    # models = retrieve_models()
-    # datasets = retrieve_datasets()
-    # print(models)
-    # print(datasets)
-    gr.Label("Subsection of text", visible=True)
-    # subsection_type = gr.Dropdown(models, label="Search models", interactive=True, scale=0)
-    # url = "https://hf-mirror.com/SaProtHub/Model-Thermostability-650M"
-    url = "https://www.uniprot.org/uniprotkb/P05067"
-    url = "https://huggingface.co/SaProtHub/Model-Thermostability-650M"
-    # url = "https://www.baidu.com"
-    response = requests.get(url)
-    # print(response.text)
-    # gr.HTML(response.text)
-    # iframe = f'<iframe srcdoc="{response.text}" width="100%"></iframe>'
-    iframe = f'<iframe src="{url}" width="100%" height="800"></iframe>'
-    # HTML = gr.HTML(open("/root/PycharmProjects/SaprotHub/hub_demo/test.html").read())
-    HTML = gr.HTML(response.text)
+    gr.Label("SaprotHub search", visible=True, show_label=False)
+    search_box = gr.Textbox(label="Search box", placeholder="Input keywords to search", interactive=True, scale=0, container=True)
+
+    # Display search results
+    search_hint = gr.Markdown("# Search results:", visible=True)
+    items = gr.Markdown(visible=False)
+
+    # Set events
+    search_box.change(show_card_info, inputs=[search_box], outputs=[items])
 
 
 if __name__ == '__main__':
     # Run the demo
-    demo.launch(server_name="0.0.0.0")
+    demo.launch(server_name="0.0.0.0")