Migrate from Gradio to Streamlit

app.py

@@ -6,107 +6,100 @@ import pprint
 import re
 import string
 
-import gradio as gr
+import streamlit as st
+import streamlit.components.v1 as components
 import requests
 
-pp = pprint.PrettyPrinter(indent=2)
-
 
-def get_docid_html(docid):
-    data_org, dataset, docid = docid.split("/")
+pp = pprint.PrettyPrinter(indent=2)
+st.set_page_config(page_title="Gaia Search", layout="wide")
 
-    docid_html = """<a
-        class="underline-on-hover"
-        title="I am hovering over the text"
-        style="color:#2D31FA;"
-        href="https://huggingface.co/datasets/bigscience-data/{}"
-        target="_blank">{}</a><span style="color: #7978FF;">/{}</span>""".format(
-        dataset, data_org + "/" + dataset, docid
+os.makedirs(os.path.join(os.getcwd(),".streamlit"), exist_ok = True)
+with open(os.path.join(os.getcwd(),".streamlit/config.toml"), "w") as file:
+    file.write(
+        '[theme]\nbase="light"'
     )
-    return docid_html
-
-
-PII_TAGS = {"KEY", "EMAIL", "USER", "IP_ADDRESS", "ID", "IPv4", "IPv6"}
-PII_PREFIX = "PI:"
-
-
-def process_pii(text):
-    for tag in PII_TAGS:
-        text = text.replace(
-            PII_PREFIX + tag,
-            """<b><mark style="background: Fuchsia; color: Lime;">REDACTED {}</mark></b>""".format(tag),
-        )
-    return text
 
+LANG_MAPPING = {'Arabic':'ar',
+                'Catalan':'ca',
+                'Code':'code',
+                'English':'en',
+                'Spanish':'es',
+                'French':'fr',
+                'Indonesian':'id',
+                'Indic':'indic',
+                'Niger-Congo':'nigercongo',
+                'Portuguese': 'pt',
+                'Vietnamese': 'vi',
+                'Chinese': 'zh',
+                'Detect Language':'detect_language',
+                'All':'all'}
 
-def process_meta_roots(result):
-    meta_html = (
-        """
-            <p class='underline-on-hover' style='font-size:12px; font-family: Arial; color:#585858; text-align: left;'>
-            <a href='{}' target='_blank'>{}</a></p>""".format(
-            result["meta"]["url"], result["meta"]["url"]
-        )
-        if "meta" in result and result["meta"] is not None and "url" in result["meta"]
-        else ""
-    )
 
 
+st.sidebar.markdown(
 """
-'meta': { 'docs': [ { 'TEXT': 'Hello World Example Hello '
-                                               'World Page Hello World.',
-                                       'URL': 'http://images.slideplayer.com/8/2335183/slides/slide_6.jpg',
-                                       '_id': 592573973},
-                                     { 'TEXT': 'Hello World Example Hello '
-                                               'World Page Hello World.',
-                                       'URL': 'http://images.slideplayer.com/8/2335183/slides/slide_9.jpg',
-                                       '_id': 1807595732},
-                                     { 'TEXT': 'Hello World Example Hello '
-                                               'World Page Hello World.',
-                                       'URL': 'http://images.slideplayer.com/8/2335183/slides/slide_10.jpg',
-                                       '_id': 1864921031},
-                                     { 'TEXT': 'Hello World Example Hello '
-                                               'World Page Hello World!',
-                                       'URL': 'http://images.slideplayer.com/8/2335183/slides/slide_5.jpg',
-                                       '_id': 1964462104},
-                                     { 'TEXT': 'Hello World Example Hello '
-                                               'World Page Hello World.',
-                                       'URL': 'http://images.slideplayer.com/8/2335183/slides/slide_8.jpg',
-                                       '_id': 2167992166}]},
+<style>
+.aligncenter {
+    text-align: center;
+    font-weight: bold;
+    font-size: 50px;
+}
+</style>
+<p class="aligncenter">Gaia Search 🌖🌏</p>
+<p style="text-align: center;"> A search engine for the LAION large scale image caption corpora</p>
+""",
+unsafe_allow_html=True,
+)
+
+st.sidebar.markdown(
 """
-def process_meta_laion(result):
-    meta_html = """"""
-    if "meta" not in result:
-        return meta_html
-    for doc in result["meta"]["docs"]:
-        # doc = json.loads(doc)
-        print(type(doc), doc)
-        print(doc["URL"])
-        meta_html += """<p class='underline-on-hover' style='font-size:12px; color:#7978FF; text-align: left;'>
-        <a href='{}' target='_blank'>{}</a></p>""".format(doc["URL"], doc["URL"])
-
-    return meta_html
-
-def process_results(results, highlight_terms):
-    if len(results) == 0:
-        return """<br><p style='font-family: Arial; color:Silver; text-align: center;'>
-                No results retrieved.</p><br><hr>"""
-
-    results_html = ""
-    for result in results:
-        tokens = result["text"].split()
-        tokens_html = []
-        for token in tokens:
-            if token in highlight_terms:
-                tokens_html.append("<b>{}</b>".format(token))
-            else:
-                tokens_html.append(token)
-        tokens_html = " ".join(tokens_html)
-        tokens_html = process_pii(tokens_html)
-        meta_html = process_meta_laion(result)
-        docid_html = """<p style='font-size:14px; font-family: Arial; color:#7978FF; text-align: left;'>Document ID: {}</p>""".format(result["docid"]) # get_docid_html(result["docid"])
-        language_html = """<p style='font-size:12px; font-family: Arial; color:MediumAquaMarine'>Language: {}</p>""".format(result["lang"])
-        results_html += tokens_html + meta_html + "<br>"
-    return results_html + "<hr>"
+<style>
+.aligncenter {
+    text-align: center;
+}
+</style>
+<p style='text-align: center'>
+<a href="" >GitHub</a> | <a href="" >Project Report</a>
+</p>
+<p class="aligncenter">
+  <a href="" target="_blank"> 
+        <img src="https://colab.research.google.com/assets/colab-badge.svg"/>
+    </a>
+</p>
+""",
+unsafe_allow_html=True,
+)
+
+query = st.sidebar.text_input(label='Search query', value='')
+language = st.sidebar.selectbox(
+    'Language',
+    ('Arabic', 'Catalan', 'Code', 'English', 'Spanish', 'French', 'Indonesian', 'Indic', 'Niger-Congo', 'Portuguese', 'Vietnamese', 'Chinese', 'Detect Language', 'All'),
+    index=3)
+max_results = st.sidebar.slider(
+    "Maximum Number of Results",
+    min_value=1,
+    max_value=100,
+    step=1,
+    value=10,
+    help="Maximum Number of Documents to return",
+)
+footer="""<style>
+.footer {
+position: fixed;
+left: 0;
+bottom: 0;
+width: 100%;
+background-color: white;
+color: black;
+text-align: center;
+}
+</style>
+<div class="footer">
+<p>Powered by <a href="https://huggingface.co/" >HuggingFace 🤗</a> and <a href="https://github.com/castorini/pyserini" >Pyserini 🦆</a></p>
+</div>
+"""
+st.sidebar.markdown(footer,unsafe_allow_html=True)
 
 
 def scisearch(query, language, num_results=10):
@@ -128,8 +121,6 @@ def scisearch(query, language, num_results=10):
 
         payload = json.loads(output.text)
 
-        pp.pprint(payload)
-
         if "err" in payload:
             if payload["err"]["type"] == "unsupported_lang":
                 detected_lang = payload["err"]["meta"]["detected_lang"]
@@ -141,38 +132,6 @@ def scisearch(query, language, num_results=10):
 
         results = payload["results"]
         highlight_terms = payload["highlight_terms"]
-
-        if language == "detect_language":
-            return (
-                (
-                    f"""<p style='font-family: Arial; color:MediumAquaMarine; text-align: center; line-height: 3em'>
-                Detected language: <b>{results[0]["lang"]}</b></p><br><hr><br>"""
-                    if len(results) > 0 and language == "detect_language"
-                    else ""
-                )
-                + process_results(results, highlight_terms)
-            )
-
-        if language == "all":
-            results_html = ""
-            for lang, results_for_lang in results.items():
-                if len(results_for_lang) == 0:
-                    results_html += f"""<p style='font-family: Arial; color:Silver; text-align: left; line-height: 3em'>
-                            No results for language: <b>{lang}</b><hr></p>"""
-                    continue
-
-                collapsible_results = f"""
-                    <details>
-                        <summary style='font-family: Arial; color:MediumAquaMarine; text-align: left; line-height: 3em'>
-                            Results for language: <b>{lang}</b><hr>
-                        </summary>
-                        {process_results(results_for_lang, highlight_terms)}
-                    </details>"""
-                results_html += collapsible_results
-            return results_html
-
-        return process_results(results, highlight_terms)
-
     except Exception as e:
         results_html = f"""
                 <p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'>
@@ -183,92 +142,132 @@ def scisearch(query, language, num_results=10):
             """
         print(e)
 
-    return results_html
+    return results, highlight_terms
 
+PII_TAGS = {"KEY", "EMAIL", "USER", "IP_ADDRESS", "ID", "IPv4", "IPv6"}
+PII_PREFIX = "PI:"
 
-def flag(query, language, num_results, issue_description):
-    try:
-        post_data = {"query": query, "k": num_results, "flag": True, "description": issue_description}
-        if language != "detect_language":
-            post_data["lang"] = language
-
-        output = requests.post(
-            os.environ.get("address"),
-            headers={"Content-type": "application/json"},
-            data=json.dumps(post_data),
-            timeout=120,
-        )
-
-        results = json.loads(output.text)
-    except:
-        print("Error flagging")
-    return ""
-
-
-description = """# <p style="text-align: center;">GAIA 🌖🌏</p>
-A large scale text corpora search engine."""
-
-
-if __name__ == "__main__":
-    demo = gr.Blocks(
-        css=".underline-on-hover:hover { text-decoration: underline; } .flagging { font-size:12px; color:Silver; }"
-    )
-
-    with demo:
-        with gr.Row():
-            gr.Markdown(value=description)
-        with gr.Row():
-            query = gr.Textbox(lines=2, placeholder="Type your query here...", label="Query")
-        with gr.Row():
-            lang = gr.Dropdown(
-                choices=[
-                    "ar",
-                    "ca",
-                    "code",
-                    "en",
-                    "es",
-                    "eu",
-                    "fr",
-                    "id",
-                    "indic",
-                    "nigercongo",
-                    "pt",
-                    "vi",
-                    "zh",
-                    "detect_language",
-                    "all",
-                ],
-                value="en",
-                label="Language",
-            )
-        with gr.Row():
-            k = gr.Slider(1, 100, value=10, step=1, label="Max Results")
-        with gr.Row():
-            submit_btn = gr.Button("Submit")
-        with gr.Row():
-            results = gr.HTML(label="Results")
-        flag_description = """
-            <p class='flagging'>
-            If you choose to flag your search, we will save the query, language and the number of results you requested.
-            Please consider adding any additional context in the box on the right.</p>"""
-        with gr.Column(visible=False) as flagging_form:
-            flag_txt = gr.Textbox(
-                lines=1,
-                placeholder="Type here...",
-                label="""If you choose to flag your search, we will save the query, language and the number of results
-                    you requested. Please consider adding relevant additional context below:""",
+def process_pii(text):
+    for tag in PII_TAGS:
+        text = text.replace(
+            PII_PREFIX + tag,
+            """<b><mark style="background: Fuchsia; color: Lime;">REDACTED {}</mark></b>""".format(tag),
             )
-            flag_btn = gr.Button("Flag Results")
-            flag_btn.click(flag, inputs=[query, lang, k, flag_txt], outputs=[flag_txt])
+    return text
 
-        def submit(query, lang, k):
-            if query == "":
-                return ["", ""]
-            return {
-                results: scisearch(query, lang, k),
-                flagging_form: gr.update(visible=True),
+def highlight_string(paragraph: str, highlight_terms: list) -> str:
+    for term in highlight_terms:
+        paragraph = re.sub(f"\\b{term}\\b", f"<b>{term}</b>", paragraph, flags=re.I)
+    paragraph = process_pii(paragraph)
+    return paragraph
+
+def process_results(hits: list, highlight_terms: list) -> str:
+    hit_list = []
+    for i, hit in enumerate(hits):
+        res_head = f"""
+                    <div class="searchresult">
+                        <h2>{i+1}. Document ID: {hit['docid']}</h2>
+                        <p>Language: <string>{hit['lang']}</string>, Score: {round(hit['score'], 2)}</p>              
+                    """
+        for subhit in hit['meta']['docs']:
+            res_head += f"""
+                        <button onclick="load_image({subhit['_id']})">Load Image</button><br>
+                        <p><img id='{subhit['_id']}' src='{subhit['URL']}'  style="width:400px;height:auto;display:none;"></p>
+                        <a href='{subhit['URL']}'>{subhit['URL']}</a>
+                        <p>{highlight_string(subhit['TEXT'], highlight_terms)}</p>
+                        """
+        res_head += f"""
+                    <p>{highlight_string(hit['text'], highlight_terms)}</p>
+                    </div>
+                    <hr>
+                    """
+        hit_list.append(res_head)
+    return " ".join(hit_list)
+
+
+if st.sidebar.button("Search"):
+    hits, highlight_terms = scisearch(query, LANG_MAPPING[language], max_results)
+    html_results = process_results(hits, highlight_terms)
+    rendered_results = f"""
+            <div id="searchresultsarea">
+                <br>
+                <p id="searchresultsnumber">About {max_results} results</p>
+                {html_results}
+            </div>
+                        """
+    st.markdown("""
+                <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.0.2/dist/css/bootstrap.min.css" rel="stylesheet" 
+                    integrity="sha384-EVSTQN3/azprG1Anm3QDgpJLIm9Nao0Yz1ztcQTwFspd3yD65VohhpuuCOmLASjC" crossorigin="anonymous">
+                """,
+                unsafe_allow_html=True)
+    st.markdown(
+        """
+        <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.0.0/css/all.min.css">                                  
+        """,
+        unsafe_allow_html=True)
+    st.markdown(
+        f"""
+        <div class="row no-gutters mt-3 align-items-center">
+            Gaia Search 🌖🌏
+            <div class="col col-md-4">
+                <input class="form-control border-secondary rounded-pill pr-5" type="search" value="{query}" id="example-search-input2">
+            </div>
+            <div class="col-auto">
+                <button class="btn btn-outline-light text-dark border-0 rounded-pill ml-n5" type="button">
+                    <i class="fa fa-search"></i>
+                </button>
+            </div>
+        </div> 
+        """,
+        unsafe_allow_html=True)
+    components.html(
+        """
+        <style>
+        #searchresultsarea {
+            font-family: 'Arial';
+        }
+        
+        #searchresultsnumber {
+            font-size: 0.8rem;
+            color: gray;
+        }
+        
+        .searchresult h2 {
+            font-size: 19px;
+            line-height: 18px;
+            font-weight: normal;
+            color: rgb(7, 111, 222);
+            margin-bottom: 0px;
+            margin-top: 25px;
+        }
+        
+        .searchresult a {
+            font-size: 12px;
+            line-height: 12px;
+            color: green;
+            margin-bottom: 0px;
+        }
+    
+        .dark-mode {
+          color: white;
+        }
+        </style>
+        <script>
+        function load_image(id){
+            console.log(id)
+            var x = document.getElementById(id);
+            console.log(x)
+            if (x.style.display === "none") {
+                x.style.display = "block";
+            } else {
+                x.style.display = "none";
             }
-
-        submit_btn.click(submit, inputs=[query, lang, k], outputs=[results, flagging_form])
-
-    demo.launch(enable_queue=True, debug=True)
+        };
+        function myFunction() {
+           var element = document.body;
+           element.classList.toggle("dark-mode");
+        }
+        </script>
+        <button onclick="myFunction()">Toggle dark mode</button>
+        """ + rendered_results, height=800, scrolling=True
+        )