app.py

import os
import streamlit as st
import streamlit.components.v1 as components
from datasets import load_dataset
from pyserini.search.lucene import LuceneSearcher

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"')


st.sidebar.markdown(
    """
<style>
.aligncenter {
    text-align: center;
    font-weight: bold;
    font-size: 50px;
}
</style>
<p class="aligncenter">Streamlit Search</p>
""",
    unsafe_allow_html=True,
)

st.sidebar.markdown(
    """
<style>
.aligncenter {
    text-align: center;
}
</style>
""",
    unsafe_allow_html=True,
)

query = st.sidebar.text_input(label="Search query", value="")

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)


searcher = LuceneSearcher("index")
ds = load_dataset("imdb", split="train")


def search(query):
    hits = searcher.search(query, k=10)
    results = ds.select([int(hit.docid) for hit in hits])
    return results["text"]


if st.sidebar.button("Search"):
    results = search(query)
    results_html = ""
    for result in results:
        results_html += result + "<br><hr><br>"
    rendered_results = f"""
            <div id="searchresultsarea">
                <br>
                {results_html}
            </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,
    )
    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;
        }
        </style>
        """
        + rendered_results,
        height=800,
        scrolling=True,
    )