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_fine-tuning-use-cmlm-multilingual-quran-splitted.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
    file = open('quran-splitted.sav','rb')
    quran_splitted = pickle.load(file)
    quran_splitted['similarity'] = text_similarity
    sorted_quran = quran_splitted.sort_values(by='similarity', ascending=False)
    
    # insert the similarity value to dataframe & sort it
    quran = pd.read_csv('quran-simple-clean.txt', delimiter="|")
    
    results = ""
    i = 0
    while i<6:
        result = sorted_quran.iloc[i]
        result_quran = quran.loc[(quran['sura']==result['sura']) & (quran['aya']==result['aya'])]
        results = results + result_quran['text'].item()+" (Q.S "+str(result['sura']).rstrip('.0')+":"+str(result['aya']).rstrip('.0')+")\n"
        i=i+1
    
    return results

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