app.py

import os
import pandas as pd
import gradio as gr
from persian_rag import PersianRAG

os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"


# Function to load CSV and initialize PersianRAG
def init_rag(knowledge_file, embedding_model, llm_model, device, retrieved_docs):
    knowledge = pd.read_csv(knowledge_file)
    rag_system = PersianRAG(knowledge, embedding_model=embedding_model, llm_model=llm_model, device=device,
                            retrieved_docs=retrieved_docs)
    return rag_system


# Function to handle querying
def query_rag(rag_system, query):
    return rag_system.rag(query)


# Gradio interface to upload CSV and configure RAG system
def rag_interface(knowledge_file, query, embedding_model, llm_model, device, retrieved_docs):
    rag_system = init_rag(knowledge_file, embedding_model, llm_model, device, retrieved_docs)
    return query_rag(rag_system, query)


# Create Gradio interface
interface = gr.Interface(
    fn=rag_interface,
    inputs=[
        gr.File(label="Upload Knowledge Base CSV"),
        gr.Textbox(label="Enter your query"),
        gr.Dropdown(choices=["LABSE", "paraphrase-multilingual-mpnet-base-v2"], value="LABSE", label="Embedding Model"),
        gr.Textbox(value="MehdiHosseiniMoghadam/AVA-Mistral-7B-V2", label="LLM Model Name"),
        gr.Dropdown(choices=["cuda", "cpu"], value="cuda", label="Device"),
        gr.Slider(minimum=1, maximum=5, step=1, value=3, label="Number of Retrieved Documents")],
    outputs="text",
    title="Persian RAG System",
    description="Upload a CSV file as the knowledge base, ask a question, and get an answer.")

# Launch the Gradio interface
interface.launch(share=False)