import gradio as gr
import pandas as pd
import requests
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from rag_embeddings import RagRetriever

# Load the RAG model
rag_retriever = RagRetriever("rag_sequence")

# Load the Mistral model
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/mistral-small-12L-4H-768d-albert")
model = AutoModelForSeq2SeqLM.from_pretrained("EleutherAI/mistral-small-12L-4H-768d-albert")

# Load the tabular data
data = pd.read_csv("data.csv")
#ADD DATASETS HERE

def langchain(user_prompt, master_prompt):
    # Retrieve data chunks using RAG embeddings
    retrieved_data = rag_retriever.retrieve(user_prompt, data, num_results=5)

    # Connect retrieved data chunks to user prompt and master prompt
    input_text = user_prompt + " " + master_prompt + " " + " ".join(retrieved_data)

    # Generate response using Mistral model
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
    generated_ids = model.generate(input_ids)
    response = tokenizer.decode(generated_ids[0], skip_special_tokens=True)

    return response

iface = gr.Interface(
    fn=langchain,
    inputs=["text", "text"],
    outputs="text",
    title="LangChain App",
    description="A Gradio app that retrieves specific datachunks using RAG embeding, from tabular csv data, and then connects those into the user prompt and the master prompt and then feed them into a Mistral model called from Hugging Face ran locally, then returns the response to the user via the gradio app GUI.",
)

iface.launch()

iface.launch(share=True)