|
import os |
|
import gradio as gr |
|
import warnings |
|
from pathlib import Path |
|
import shutil |
|
|
|
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning) |
|
|
|
from rag_tool import RAGTool |
|
|
|
|
|
rag_tool = RAGTool() |
|
|
|
|
|
def upload_file(file): |
|
try: |
|
|
|
os.makedirs("./documents", exist_ok=True) |
|
|
|
|
|
file_path = Path(file.name) |
|
destination = Path("./documents") / file_path.name |
|
|
|
|
|
shutil.copy(file_path, destination) |
|
|
|
|
|
rag_tool.configure( |
|
documents_path=str(destination), |
|
embedding_model="sentence-transformers/all-MiniLM-L6-v2", |
|
persist_directory="./vector_store" |
|
) |
|
|
|
return f"File uploaded and processed: {file_path.name}" |
|
except Exception as e: |
|
return f"Error processing file: {str(e)}" |
|
|
|
|
|
def query_document(question): |
|
try: |
|
if not hasattr(rag_tool, 'vector_store') or rag_tool.vector_store is None: |
|
return "Please upload a document first." |
|
|
|
response = rag_tool(question) |
|
return response |
|
except Exception as e: |
|
return f"Error querying document: {str(e)}" |
|
|
|
|
|
with gr.Blocks(title="RAG Tool") as demo: |
|
gr.Markdown("# Document Question Answering System") |
|
gr.Markdown("Upload a document (PDF, TXT) and ask questions about it") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
file_input = gr.File(label="Upload Document") |
|
upload_button = gr.Button("Process Document") |
|
upload_result = gr.Textbox(label="Upload Status") |
|
|
|
with gr.Column(): |
|
query_input = gr.Textbox(label="Ask a Question", placeholder="What would you like to know?") |
|
query_button = gr.Button("Get Answer") |
|
response_output = gr.Textbox(label="Answer") |
|
|
|
|
|
upload_button.click( |
|
upload_file, |
|
inputs=file_input, |
|
outputs=upload_result |
|
) |
|
|
|
query_button.click( |
|
query_document, |
|
inputs=query_input, |
|
outputs=response_output |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
demo.launch() |