import gradio as gr
from typing import TypedDict, List
from rank_bm25 import BM25Okapi


class Hit(TypedDict):
    cid: str
    score: float
    text: str

docs = [
    "Bacteria can be used to make cheese from milk. The bacteria turn the milk sugars into lactic acid. The acid is what causes the milk to curdle to form cheese. Bacteria are also involved in producing other foods. Yogurt is made by using bacteria to ferment milk. Fermenting cabbage with bacteria produces sauerkraut.",
    "Mesophiles grow best in moderate temperature, typically between 25°C and 40°C (77°F and 104°F). Mesophiles are often found living in or on the bodies of humans or other animals. The optimal growth temperature of many pathogenic mesophiles is 37°C (98°F), the normal human body temperature. Mesophilic organisms have important uses in food preparation, including cheese, yogurt, beer and wine.",
    "A wide range of friendly bacteria live in the gut. Bacteria begin to populate the human digestive system right after birth. Gut bacteria include Lactobacillus, the bacteria commonly used in probiotic foods such as yogurt, and E. coli bacteria. About a third of all bacteria in the gut are members of the Bacteroides species. Bacteroides are key in helping us digest plant food."
]


tokenized_docs = [doc.split() for doc in docs]
bm25 = BM25Okapi(tokenized_docs)


def search(query: str) -> List[Hit]:
    tokenized_query = query.split()
    scores = bm25.get_scores(tokenized_query)
    results = [
        {"cid": f"train-{i}", "score": score, "text": docs[i]} for i, score in enumerate(scores)
    ]
    return results


demo = gr.Interface(
    fn=search,  
    inputs=gr.Textbox(label="Enter your query"),  
    outputs=gr.JSON(label="Search Results"),  
    title="BM25 Search Engine Demo",
    description="Enter a query to search the SciQ dataset using a BM25-based search engine.",
)

if __name__ == "__main__":
    demo.launch()