Update Tool

.gitignore

@@ -0,0 +1 @@
+*.code-workspace

README.md

@@ -1,13 +1,13 @@
 ---
-title: Roots Search Tool
-emoji: 🌸 🔎
-colorFrom: green
-colorTo: indigo
-sdk: gradio
-sdk_version: 3.7
+title: GÆA / gaia / gæa
+emoji: 🌏🌖
+colorFrom: blue
+colorTo: red
+sdk: streamlit
+sdk_version: 1.18.1
 app_file: app.py
 pinned: false
-license: apache-2.0
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+

app.py

@@ -1,260 +1,337 @@
-import http.client as http_client
 import json
-import logging
 import os
-import re
-import string
-import traceback
+import pprint
 
-import gradio as gr
+import streamlit as st
+import streamlit.components.v1 as components
 import requests
-from huggingface_hub import HfApi
-
-hf_api = HfApi()
-roots_datasets = {dset.id.split("/")[-1]:dset for dset in hf_api.list_datasets(author="bigscience-data", use_auth_token=os.environ.get("bigscience_data_token"))}
-
-def get_docid_html(docid):
-    data_org, dataset, docid = docid.split("/")
-    metadata = roots_datasets[dataset]
-    if metadata.private:
-        docid_html = (
-            f"<a "
-            f'class="underline-on-hover"'
-            f'title="This dataset is private. See the introductory text for more information"'
-            f'style="color:#AA4A44;"'
-            f'href="https://huggingface.co/datasets/bigscience-data/{dataset}"'
-            f'target="_blank"><b>🔒{dataset}</b></a><span style="color: #7978FF;">/{docid}</span>'
-        )
-    else:
-        docid_html = (
-            f"<a "
-            f'class="underline-on-hover"'
-            f'title="This dataset is licensed {metadata.tags[0].split(":")[-1]}"'
-            f'style="color:#2D31FA;"'
-            f'href="https://huggingface.co/datasets/bigscience-data/{dataset}"'
-            f'target="_blank"><b>{dataset}</b></a><span style="color: #7978FF;">/{docid}</span>'
-        )        
-    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
 
-
-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 = (
-            """
-                <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 ""
-        )
-        docid_html = get_docid_html(result["docid"])
-        results_html += """{}
-            <p style='font-size:14px; font-family: Arial; color:#7978FF; text-align: left;'>Document ID: {}</p>
-            <p style='font-size:12px; font-family: Arial; color:MediumAquaMarine'>Language: {}</p>
-            <p style='font-family: Arial;'>{}</p>
-            <br>
-        """.format(
-            meta_html, docid_html, result["lang"], tokens_html
-        )
-    return results_html + "<hr>"
-
-
-def scisearch(query, language, num_results=10):
+from typing import Union
+
+pp = pprint.PrettyPrinter(indent=2)
+
+os.environ["address"] = "http://34.79.83.149:8080"
+
+st.set_page_config(page_title="Gaia Search 🌖🌏", layout="wide")
+
+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"')
+
+
+corpus_name_map = {
+    "LAION": "laion",
+    "ROOTS": "roots",
+    "The Pile": "pile",
+    "C4": "c4",
+}
+
+st.sidebar.markdown(
+    """
+    <style>
+    .aligncenter {
+        text-align: center;
+        font-weight: bold;
+        font-size: 36px;
+    }
+    </style>
+    <p class="aligncenter">Gaia Search 🌖🌏</p>
+    <p>A search engine for large scale texual
+    corpora. Most of the datasets included in the tool are based on Common
+    Crawl. By using the tool, you are also bound by the Common Crawl terms
+    of use in respect of the content contained in the datasets.
+    </p>
+    """,
+    unsafe_allow_html=True,
+)
+
+st.sidebar.markdown(
+    """
+    <style>
+    .aligncenter {
+        text-align: center;
+    }
+    </style>
+    <p style='text-align: center'>
+    <a href="" style="color:#7978FF;">GitHub</a> | <a href="" style="color:#7978FF;" >Project Report</a> | <a href="" style="color:#7978FF;" >Colab</a> 
+    </p>
+    """,
+    unsafe_allow_html=True,
+)
+
+# <p class="aligncenter">
+#     <a href="" target="_blank">
+#         <img src="https://colab.research.google.com/assets/colab-badge.svg"/>
+#     </a>
+# </p>
+
+
+query = st.sidebar.text_input(label="Query", placeholder="Type your query here")
+corpus = st.sidebar.selectbox(
+    "Corpus",
+    tuple(corpus_name_map.keys()),
+    index=2,
+)
+max_results = st.sidebar.slider(
+    "Max Results",
+    min_value=1,
+    max_value=100,
+    step=1,
+    value=10,
+    help="Max Number of Documents to return",
+)
+
+# dark_mode_toggle = """
+#     <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";
+#             }
+#         };
+#         function myFunction() {
+#         var element = document.body;
+#         element.classList.toggle("dark-mode");
+#         }
+#     </script>
+#     <button onclick="myFunction()">Toggle dark mode</button>
+# """
+# st.sidebar.markdown(dark_mode_toggle, unsafe_allow_html=True)
+
+
+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, corpus, num_results=10):
     try:
-        query = " ".join(query.split())
+        print(query, corpus, num_results)
+        query = query.strip()
         if query == "" or query is None:
-            return ""
+            return
 
-        post_data = {"query": query, "k": num_results}
-        if language != "detect_language":
-            post_data["lang"] = language
+        post_data = {"query": query, "corpus": corpus, "k": num_results, "lang": "all"}
+        address = (
+            os.environ.get("address")
+            if corpus != "roots"
+            else "http://34.116.206.238:8080"
+        )
 
         output = requests.post(
-            os.environ.get("address"),
+            address,
             headers={"Content-type": "application/json"},
             data=json.dumps(post_data),
             timeout=60,
         )
 
         payload = json.loads(output.text)
-
-        if "err" in payload:
-            if payload["err"]["type"] == "unsupported_lang":
-                detected_lang = payload["err"]["meta"]["detected_lang"]
-                return f"""
-                    <p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'>
-                    Detected language <b>{detected_lang}</b> is not supported.<br>
-                    Please choose a language from the dropdown or type another query.
-                    </p><br><hr><br>"""
-
-        results = payload["results"]
-        highlight_terms = payload["highlight_terms"]
-
-        if language == "detect_language":
-            results = list(results.values())[0]
-            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
-
-        results = list(results.values())[0]
-        return process_results(results, highlight_terms)
+        return payload["results"], payload["highlight_terms"]
 
     except Exception as e:
-        results_html = f"""
-                <p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'>
-                Raised {type(e).__name__}</p>
-                <p style='font-size:14px; font-family: Arial; '>
-                Check if a relevant discussion already exists in the Community tab. If not, please open a discussion.
-                </p>
-            """
         print(e)
-        print(traceback.format_exc())
 
-    return results_html
 
+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,
+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
+            ),
         )
-
-        results = json.loads(output.text)
-    except:
-        print("Error flagging")
-    return ""
+    return text
 
 
-description = """# <p style="text-align: center;"> 🌸 🔎 ROOTS search tool 🔍 🌸 </p>
-The ROOTS corpus was developed during the [BigScience workshop](https://bigscience.huggingface.co/) for the purpose
-of training the Multilingual Large Language Model [BLOOM](https://huggingface.co/bigscience/bloom). This tool allows
-you to search through the ROOTS corpus. We serve a BM25 index for each language or group of languages included in
-ROOTS. You can read more about the details of the tool design
-[here](https://huggingface.co/spaces/bigscience-data/scisearch/blob/main/roots_search_tool_specs.pdf). For more
-information and instructions on how to access the full corpus check [this form](https://forms.gle/qyYswbEL5kA23Wu99)."""
+def highlight_string(paragraph: str, highlight_terms: list) -> str:
+    tokens = paragraph.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)
+    return process_pii(tokens_html)
+
+
+def extract_lang_from_docid(docid):
+    return docid.split("_")[1]
+
+
+def format_result(result, highlight_terms):
+    text = result["text"]
+    docid = result["docid"]
+    tokens_html = highlight_string(text, highlight_terms)
+    language = extract_lang_from_docid(docid)
+    result_html = """
+        <span style='font-size:14px; font-family: Arial; color:MediumAquaMarine'>Language: {} | </span>
+        <span style='font-size:14px; font-family: Arial; color:#7978FF; text-align: left;'>Document ID: {} | </span><br>
+        <span style='font-family: Arial;'>{}</span><br>
+        <br>
+    """.format(
+        language, docid, tokens_html
+    )
+    return "<p>" + result_html + "</p>"
 
 
-if __name__ == "__main__":
-    demo = gr.Blocks(
-        css=".underline-on-hover:hover { text-decoration: underline; } .flagging { font-size:12px; color:Silver; }"
-    )
+def process_results(corpus: str, hits: Union[list, dict], highlight_terms: list) -> str:
+    hit_list = []
 
-    with demo:
-        with gr.Row():
-            gr.Markdown(value=description)
-        with gr.Row():
-            query = gr.Textbox(lines=1, max_lines=1, 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:""",
+    if corpus == "roots":
+        result_page_html = ""
+        for lang, results_for_lang in hits.items():
+            print("Processing language", lang)
+            if len(results_for_lang) == 0:
+                result_page_html += """<div style='font-family: Arial; color:Silver; text-align: left; line-height: 3em'>
+                    No results for language: <b>{}</b></div>""".format(
+                    lang
+                )
+                continue
+            results_for_lang_html = ""
+            for result in results_for_lang:
+                result_html = format_result(result, highlight_terms)
+                results_for_lang_html += result_html
+            results_for_lang_html = f"""
+                <details>
+                    <summary style='font-family: Arial; color:MediumAquaMarine; text-align: left; line-height: 3em'>
+                        Results for language: <b>{lang}</b>
+                    </summary>
+                    {results_for_lang_html}
+                </details>"""
+            result_page_html += results_for_lang_html
+        return result_page_html
+
+    for hit in hits:
+        res_head = f"""
+            <p class="searchresult" style="color: #7978FF;">Document ID: {hit['docid']} | Score: {round(hit['score'], 2)}</p>
+            """
+        if corpus == "laion":
+            res_head += f"""
+                <p style="color: #7978FF;">Caption:</p>
+                <p>{highlight_string(hit['text'], highlight_terms)}</p>
+            """
+            if (
+                "meta" in hit
+                and hit["meta"] is not None
+                and "docs" in hit["meta"]
+                and len(hit["meta"]["docs"]) > 0
+            ):
+                res_head += """<p style="color: #7978FF;"> Image links:</p><ul>"""
+                for subhit in hit["meta"]["docs"]:
+                    res_head += f"""<li><a href={subhit["URL"]} target="_blank" style="color:#ffcdf8; ">{subhit["URL"]}</a></li>"""
+                res_head += "</ul>"
+            res_head += "<hr>"
+        else:
+            res_head += (
+                f"""<p>{highlight_string(hit['text'], highlight_terms)}</p></div><hr>"""
             )
-            flag_btn = gr.Button("Flag Results")
-            flag_btn.click(flag, inputs=[query, lang, k, flag_txt], outputs=[flag_txt])
-
-        def submit(query, lang, k):
-            query = query.strip()
-            if query is None or query == "":
-                return "", ""
-            return {
-                results: scisearch(query, lang, k),
-                flagging_form: gr.update(visible=True),
-            }
-
-        query.submit(fn=submit, inputs=[query, lang, k], outputs=[results, flagging_form])
-        submit_btn.click(submit, inputs=[query, lang, k], outputs=[results, flagging_form])
-
-    demo.launch(enable_queue=True, debug=True)
+        hit_list.append(res_head)
+    return " ".join(hit_list)
+
+
+submit_button = st.sidebar.button("Search", type="primary")
+
+if submit_button or query:
+    query = query.strip()
+    if query is None or query == "":
+        components.html(
+            """<p style='font-size:18px; font-family: Arial; color:MediumVioletRed; text-align: center;'>
+            Please provide a non-empty query.
+            </p><br><hr><br>"""
+        )
+    else:
+        hits, highlight_terms = scisearch(query, corpus_name_map[corpus], max_results)
+        html_results = process_results(corpus_name_map[corpus], 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/[email protected]/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,
+        # )
+        # .bk-root{position:relative;width:auto;height:auto;box-sizing:border-box;font-family:Helvetica, Arial, sans-serif;font-size:13px;}.bk-root .bk,.bk-root .bk:before,.bk-root .bk:after{box-sizing:inherit;margin:0;border:0;padding:0;background-image:none;font-family:inherit;font-size:100%;line-height:1.42857143;}.bk-root pre.bk{font-family:Courier, monospace;}
+        components.html(
+            """
+            <head>
+            <link href='https://fonts.googleapis.com/css?family=Source+Sans+Pro' rel='stylesheet' type='text/css'>
+            </head>
+            <style>
+                #searchresultsarea {
+                    font-family: "Source Sans Pro", sans-serif;
+                }
+                #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;
+                    color: #7978FF;"
+                }
+                .searchresult a {
+                    font-size: 12px;
+                    line-height: 12px;
+                    color: green;
+                    margin-bottom: 0px;
+                }
+            </style>
+            """
+            + rendered_results,
+            height=800,
+            scrolling=True,
+        )