refactor

app.py

@@ -1,7 +1,6 @@
 import json
 import os
 import pprint
-import re
 
 import streamlit as st
 import streamlit.components.v1 as components
@@ -12,74 +11,111 @@ pp = pprint.PrettyPrinter(indent=2)
 
 os.environ["address"] = "http://34.79.83.149:8080"
 
-st.set_page_config(page_title="Gaia Search", layout="wide")
+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",
+    "The Pile": "pile",
+    "C4": "c4",
+}
+
 st.sidebar.markdown(
     """
-<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>
-""",
+    <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="" >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>
-""",
+    <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,
 )
 
-query = st.sidebar.text_input(label="Search query", value="")
+# <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", value="")
 corpus = st.sidebar.selectbox(
     "Corpus",
-    ("laion", "pile", "c4"),
+    tuple(corpus_name_map.keys()),
     index=0,
 )
 max_results = st.sidebar.slider(
-    "Maximum Number of Results",
+    "Max Results",
     min_value=1,
     max_value=100,
     step=1,
     value=10,
-    help="Maximum Number of Documents to return",
+    help="Max 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>
+
+# 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)
 
@@ -101,23 +137,12 @@ def scisearch(query, corpus, num_results=10):
         )
 
         payload = json.loads(output.text)
-        results = payload["results"]
-        highlight_terms = payload["highlight_terms"]
-        return 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)
 
 
-
-
 PII_TAGS = {"KEY", "EMAIL", "USER", "IP_ADDRESS", "ID", "IPv4", "IPv6"}
 PII_PREFIX = "PI:"
 
@@ -134,126 +159,115 @@ def process_pii(text):
 
 
 def highlight_string(paragraph: str, highlight_terms: list) -> str:
-    # TODO:
-    # 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
+    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 process_results(hits: list, highlight_terms: list) -> str:
+def process_results(corpus: str, 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>, Score: {round(hit['score'], 2)}</p>
-                    """
-        if "meta" in hit:
-            if  hit["meta"] is not None and "docs" in hit["meta"]:
+            <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"""
-                                <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>
-                    """
+                    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>"""
+            )
         hit_list.append(res_head)
     return " ".join(hit_list)
 
 
-if st.sidebar.button("Search"):
-    hits, highlight_terms = scisearch(query, corpus, max_results)
-    html_results = process_results(hits, highlight_terms)
+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,
-    )
+        <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: '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";
+            #searchresultsarea {
+                font-family: "Source Sans Pro", sans-serif;
             }
-        };
-        function myFunction() {
-           var element = document.body;
-           element.classList.toggle("dark-mode");
-        }
-        </script>
-        <button onclick="myFunction()">Toggle dark mode</button>
+            #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,