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="IMDB Search", layout="wide")

st.sidebar.markdown(
    """
    <style>
    .aligncenter {
        text-align: center;
        font-weight: bold;
        font-size: 28px;
    }
    </style>
    <p class="aligncenter">IMDB 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><br><hr><br><br>"
    rendered_results = f"""
        <div id="searchresultsarea">
            {results_html}
        </div>
    """

    components.html(
        rendered_results,
        height=800,
        scrolling=True,
    )