add roots

README.md

@@ -4,7 +4,7 @@ emoji: 🌏🌖
 colorFrom: blue
 colorTo: red
 sdk: streamlit
-sdk_version: 1.17.0
+sdk_version: 1.18.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -6,6 +6,7 @@ import streamlit as st
 import streamlit.components.v1 as components
 import requests
 
+from typing import Union
 
 pp = pprint.PrettyPrinter(indent=2)
 
@@ -20,6 +21,7 @@ with open(os.path.join(os.getcwd(), ".streamlit/config.toml"), "w") as file:
 
 corpus_name_map = {
     "LAION": "laion",
+    "ROOTS": "roots",
     "The Pile": "pile",
     "C4": "c4",
 }
@@ -64,11 +66,11 @@ st.sidebar.markdown(
 # </p>
 
 
-query = st.sidebar.text_input(label="Query", value="")
+query = st.sidebar.text_input(label="Query", placeholder="Type your query here")
 corpus = st.sidebar.selectbox(
     "Corpus",
     tuple(corpus_name_map.keys()),
-    index=0,
+    index=2,
 )
 max_results = st.sidebar.slider(
     "Max Results",
@@ -127,10 +129,15 @@ def scisearch(query, corpus, num_results=10):
         if query == "" or query is None:
             return
 
-        post_data = {"query": query, "corpus": corpus, "k": num_results}
+        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,
@@ -170,9 +177,54 @@ def highlight_string(paragraph: str, highlight_terms: list) -> str:
     return process_pii(tokens_html)
 
 
-def process_results(corpus: str, hits: list, highlight_terms: list) -> str:
+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>"
+
+
+def process_results(corpus: str, hits: Union[list, dict], highlight_terms: list) -> str:
     hit_list = []
-    for i, hit in enumerate(hits):
+
+    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>
             """
@@ -200,76 +252,86 @@ def process_results(corpus: str, hits: list, highlight_terms: list) -> str:
     return " ".join(hit_list)
 
 
-if st.sidebar.button("Search", type="primary"):
-    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/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,
-    # )
-    # .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,
-    )
+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/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,
+        # )
+        # .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,
+        )