app.py

import gradio as gr
import polars as pl
import bm25s
import Stemmer
import stopwordsiso

# favourite_langs = {"English": "en", "Romanian": "ro", "German": "de", "-----": "-----"}
favourite_langs = {"English": "en", "Romanian": "ro", "German": "de"}
options = list(favourite_langs.keys())
models = ['ENRO', 'DERO']

def type_search(input_text, sselected_language, tselected_language, model_name, hits=10, search_type="Similarity", toggle_case=True):
    if search_type == "Word search":
        return search_text(input_text, sselected_language, tselected_language, model_name, hits, search_type, toggle_case)
    else: # "Best match search"
        return similarity_search(input_text, sselected_language, tselected_language, model_name, hits, search_type, toggle_case)

# English, Romanian
def search_text(input_text, sselected_language, tselected_language, model_name, hits=10, search_type="Similarity", toggle_case=True):
    # df = pl.read_csv('hf://datasets/TiberiuCristianLeon/2RO/ENRO/ENRO.tsv', separator='\t')
    # df = pl.read_parquet('hf://datasets/TiberiuCristianLeon/RSSNEWS/data/train-00000-of-00001.parquet')
    # df = pl.read_parquet('https://huggingface.co/datasets/TiberiuCristianLeon/2RO/resolve/refs%2Fconvert%2Fparquet/default/train/0000.parquet')
    path_to_model = f"https://huggingface.co/api/datasets/TiberiuCristianLeon/2RO/parquet/{model_name.lower()}/train/0.parquet"
    df = pl.read_parquet(path_to_model)
    if toggle_case:
        filtered = df.filter(pl.col(sselected_language).str.contains(input_text).alias("literal")) # case sensitive
    else:
        filtered = df.filter(pl.col(sselected_language).str.contains(f"(?i){input_text}").alias("literal")) # (?i) case insensitive
    # filtered = df.filter(pl.col(sselected_language).str.contains_any([input_text], ascii_case_insensitive=True).alias("contains_any"))
    print(toggle_case, filtered.head(hits))
    # print(filtered)
    # Extract rows
    list_of_arrays = filtered.select([sselected_language, tselected_language]).head(hits)
    # for dataframe type="numpy"
    # list_of_arrays = filtered.select([sselected_language, tselected_language]).head(hits).to_numpy()
    message_text = f'Done! Found {len(list_of_arrays)} entries'
    return list_of_arrays, message_text

def similarity_search(input_text, sselected_language, tselected_language, model_name, hits=10, search_type="Similarity", toggle_case=True):
    path_to_model = f"https://huggingface.co/api/datasets/TiberiuCristianLeon/2RO/parquet/{model_name.lower()}/train/0.parquet"
    df = pl.read_parquet(path_to_model)
    df = df.drop_nulls(subset=[sselected_language, sselected_language])
    # Extract both source and target columns
    source_corpus = df.select(sselected_language).to_series().to_list()
    target_corpus = df.select(tselected_language).to_series().to_list()
    # Filter out empty entries and keep track of valid indices
    valid_entries = [(src, tgt) for src, tgt in zip(source_corpus, target_corpus)]
    # Unpack filtered source and target texts
    filtered_source = [entry[0] for entry in valid_entries]
    filtered_target = [entry[1] for entry in valid_entries]
    # Run BM25 search on filtered source corpus
    index_name = f"index{sselected_language}{tselected_language}"
    list_of_arrays = bmretriever(filtered_source, input_text, sselected_language, filtered_target, index_name, hits)
    message_text = f'Done! Found {len(list_of_arrays)} entries'
    return list_of_arrays, message_text

def bmretriever(corpus, query, sselected_language, translations, index_name, k=10):
    stemmer = Stemmer.Stemmer(sselected_language.lower())
    try:
        corpus_tokens = bm25s.tokenize(corpus, stopwords=favourite_langs[sselected_language], stemmer=stemmer)
    except ValueError:
        stopwords = stopwordsiso.stopwords[favourite_langs[sselected_language]]
        corpus_tokens = bm25s.tokenize(corpus, stopwords=stopwords, stemmer=stemmer)            
    try:
        print('Loading saved retriever index')
        retriever = bm25s.BM25.load(index_name, load_corpus=True)
    except Exception as loadingerror:
        print(loadingerror)
        retriever = bm25s.BM25()
        retriever.index(corpus_tokens)
        retriever.save(index_name, corpus=corpus) # Save the corpus along with the model
    
    query_tokens = bm25s.tokenize(query, stemmer=stemmer)
    results, scores = retriever.retrieve(query_tokens, k=k, corpus=corpus)

    final_results = []
    for i in range(results.shape[1]):
        doc = results[0, i]
        score = scores[0, i]
        translation = translations[corpus.index(doc)]  # Match translation by index
        final_results.append((str(doc), str(translation)))
            # "score": round(float(score), 2)})
        print(f"Rank {i+1} (score: {score:.2f}): {doc} → {translation}")

    return final_results

# Define a function to swap dropdown values
def swap_languages(src_lang, tgt_lang):
    return tgt_lang, src_lang

def create_interface():
    with gr.Blocks() as interface:
        gr.Markdown("## Search Text in Dataset")

        with gr.Row():
            input_text = gr.Textbox(label="Enter text to search:", placeholder="Type your text here...", info="Press Enter key to start search")
        
        with gr.Row():
            sselected_language = gr.Dropdown(choices=options, value = options[0], label="Source language", interactive=True)
            tselected_language = gr.Dropdown(choices=options, value = options[1], label="Target language", interactive=True)
            swap_button = gr.Button("Swap Languages")
            swap_button.click(fn=swap_languages, inputs=[sselected_language, tselected_language], outputs=[sselected_language, tselected_language])
            search_type = gr.Radio(["Best match search", "Word search"], value="Best match search", label="Search type", info="Query word(s) or best match search with BM25")
            toggle_case = gr.Checkbox(info="Toggle case sensitive search", label="Case sensitive search", value=True, interactive=True, visible=True)
        model_name = gr.Dropdown(choices=models, label="Select a dataset", value = models[0], interactive=True)
        search_button = gr.Button("Search")

        translated_text = gr.Dataframe(label="Returned entries:", interactive=False, headers=[options[0], options [1]], datatype=["str", "str"], col_count=(2, "fixed"),
                                       wrap=True, show_row_numbers=False, show_copy_button=True)
        message_text = gr.Textbox(label="Messages:", placeholder="Display field for status and error messages", interactive=False)

        hits = gr.Slider(
            minimum=1,
            maximum=100,
            value=10,
            step=5,
            label="Number of returned hits")
        
        search_button.click(
            type_search,
            inputs=[input_text, sselected_language, tselected_language, model_name, hits, search_type, toggle_case], 
            outputs=[translated_text, message_text]
        )

        # Submit the form when Enter is pressed in the input_text textbox
        input_text.submit(
            type_search,
            inputs=[input_text, sselected_language, tselected_language, model_name, hits, search_type, toggle_case], 
            outputs=[translated_text, message_text]
        )

    return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.launch()