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

def find(query):
    def get_detailed_instruct(task_description: str, query: str) -> str:
        return f'Instruct: {task_description}\nQuery: {query}'
    
    # Each query must come with a one-sentence instruction that describes the task
    task = 'Given a web search query, retrieve relevant passages that answer the query'
    queries = [
        get_detailed_instruct(task, query)
    ]
    
    quran = pd.read_csv('quran-simple-clean.txt', delimiter="|")
    documents = quran['text'].tolist()
    input_texts = queries + documents
    
    model = SentenceTransformer('intfloat/multilingual-e5-large-instruct')
    
    embeddings = model.encode(input_texts, convert_to_tensor=True, normalize_embeddings=True)
    scores = (embeddings[:1] @ embeddings[1:].T) * 100

    # insert the similarity value to dataframe & sort it
    quran['similarity'] = scores.tolist()[0]
    sorted_quran = quran.sort_values(by='similarity', ascending=False)
    
    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()