import gradio as gr
import pickle
from sentence_transformers import SentenceTransformer
import pandas as pd

def find(query):
    # transform query from user
    model = SentenceTransformer('Bofandra/fine-tuning-use-cmlm-multilingual-quran-translation')
    encoded_query_text = model.encode(query)

    # get encoded quran text
    file = open('encoded_quran_cmlm_finetuned.sav','rb')
    encoded_quran_text = pickle.load(file)
    file.close()

    # compare query to each quran verse
    i = 0
    text_similarity = []
    for encoded_quran_ayat in encoded_quran_text:
      similarity = encoded_query_text @ encoded_quran_ayat.T
      text_similarity.append(similarity)
      i=i+1
      print(i)

    # insert the similarity value to dataframe & sort it
    quran = pd.read_csv('quran-simple-clean.txt', delimiter="|")
    quran['similarity'] = text_similarity
    sorted_quran = quran.sort_values(by='similarity', ascending=False)
    
    results = ""
    i = 0
    while i<6:
        results = results + sorted_quran.iloc[i]['text']+" (Q.S "+str(sorted_quran.iloc[i]['sura'])+":"+str(sorted_quran.iloc[i]['aya'])+")\n"
        i=i+1
    
    return results

demo = gr.Interface(fn=find, inputs="textbox", outputs="textbox")
    
if __name__ == "__main__":
    demo.launch()