# Author: Bastien
# Date: 5/3/2024
# Project: RAG-RESEARCH-PROJECT | BSc Computer Science - Semester 6

import os

os.system("pip install --upgrade pip")
#os.system("pip install faiss-cpu")
#os.system("pip install transformers")
#os.system("pip install torch")


# Import of required libraries 
from transformers import pipeline
from rag_functions import construct_prompt, encode_query, generate_response, clean_output_text, build_index, retrieve_documents, prompt_model
import gradio as gr
import pandas as pd
import os
import faiss
import torch
import time

preprocessed_docs_path = './preprocessed_docs.csv'
embeddings_path = './embeddings.pt'
index_path = './faiss_index'

# Load the Pre-processed docs from CSV
preprocessed_docs = pd.read_csv(preprocessed_docs_path)

# Load embeddings
doc_embeddings = torch.load(embeddings_path)

# Load FAISS index
index = faiss.read_index(index_path)

# Define a Gradio interface
def chat_interface(question, history_df):
    response = prompt_model(question, index, preprocessed_docs)
    # Insert the new question and response at the beginning of the DataFrame
    history_df = pd.concat([pd.DataFrame({"Question": [question], "Answer": [response]}), history_df], ignore_index=True)
    return response, history_df

with gr.Blocks() as demo:
    with gr.Row():
        question = gr.Textbox(label="Your Question", placeholder="Type Here...")
        submit_btn = gr.Button("Ask")
    response = gr.Textbox(label="Answer", interactive=False)
    
    # Initialize an empty DataFrame to keep track of question-answer history
    history_display = gr.Dataframe(headers=["Question", "Answer"], value=[], interactive=False)

    submit_btn.click(fn=chat_interface, inputs=[question, history_display], outputs=[response, history_display])

if __name__ == "__main__":
    demo.launch()