from ragatouille import RAGPretrainedModel
import os 
import gradio as gr


path_to_index = 'colbert/indexes/akhooli/Arabic-ColBERT-100knew_index'
message = "waiting to load index ..."
if os.path.exists(path_to_index):
    RAG = RAGPretrainedModel.from_index(path_to_index)
    message = "index loaded!"
    print(message)
    
import gradio as gr

def process_results(results):
  for r in results:
    print(f"Sura: {r['document_id']} ({r['document_metadata']}) \n Text:{r['content']}")

k = 3 # How many documents you want to retrieve

def answer_fn(query):
    results = RAG.search(query= query)
    return process_results(results)

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