Spaces:
Sleeping
Sleeping
File size: 3,963 Bytes
1b31f0c 4b2cf1f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | import os
import gradio as gr
from ingestion.build_index import build
from pipeline.rag_pipeline import rag_answer
DEFAULT_PDF_PATH = 'data/documents/1.pdf'
def load_knowledge_base(pdf_path):
if not pdf_path or not os.path.exists(pdf_path):
raise FileNotFoundError(f'PDF not found at "{pdf_path}".')
index, chunks = build(pdf_path)
document_name = os.path.basename(pdf_path)
return index, chunks, document_name
def load_default_pdf():
try:
index, chunks, document_name = load_knowledge_base(DEFAULT_PDF_PATH)
status = f'Ready. Loaded default document: {document_name}'
return index, chunks, document_name, status
except ValueError as e:
return None, None, '', f'Error: {str(e)}'
except Exception as e:
return None, None, '', f'Error loading PDF: {str(e)}'
def load_uploaded_pdf(file_path):
if not file_path:
return None, None, '', 'Upload a PDF first.'
try:
index, chunks, document_name = load_knowledge_base(file_path)
status = f'Ready. Loaded uploaded document: {document_name}'
return index, chunks, document_name, status
except ValueError as e:
return None, None, '', f'Error: {str(e)}'
except Exception as e:
return None, None, '', f'Error loading PDF: {str(e)}'
def ask(question, index, chunks, document_name):
if not question or not question.strip():
return 'Please enter a question.'
if index is None or not chunks:
return 'Load a PDF first, then ask a question.'
answer = rag_answer(question, index, chunks)
if document_name:
return f'Document: {document_name}\n\n{answer}'
return answer
def ask_and_clear(question, index, chunks, document_name):
result = ask(question, index, chunks, document_name)
return result, ''
with gr.Blocks(title='RAG PDF Assistant') as demo: # type: ignore
gr.Markdown('# RAG PDF Assistant')
gr.Markdown('Upload a PDF or use the bundled sample document, then ask questions against your RAG pipeline.')
index_state = gr.State(value=None)
chunks_state = gr.State(value=None)
document_name_state = gr.State(value='')
with gr.Row():
with gr.Column(scale=1):
file_input = gr.File(label='Upload PDF', file_types=['.pdf'], type='filepath')
load_upload_btn = gr.Button('Load Uploaded PDF', variant='secondary')
load_default_btn = gr.Button('Use Default PDF', variant='primary')
status_box = gr.Textbox(label='Status', interactive=False, lines=3)
with gr.Column(scale=2):
question_box = gr.Textbox(
label='Your Question',
placeholder='Ask something about the loaded PDF...',
lines=4,
)
ask_btn = gr.Button('Ask', variant='primary')
answer_box = gr.Textbox(label='Answer', lines=14, interactive=False)
load_default_btn.click(
fn=load_default_pdf,
inputs=None,
outputs=[index_state, chunks_state, document_name_state, status_box],
)
load_upload_btn.click(
fn=load_uploaded_pdf,
inputs=file_input,
outputs=[index_state, chunks_state, document_name_state, status_box],
)
ask_btn.click(
fn=ask_and_clear,
inputs=[question_box, index_state, chunks_state, document_name_state],
outputs=[answer_box, question_box],
)
question_box.submit(
fn=ask_and_clear,
inputs=[question_box, index_state, chunks_state, document_name_state],
outputs=[answer_box, question_box],
)
demo.load(
fn=load_default_pdf,
inputs=None,
outputs=[index_state, chunks_state, document_name_state, status_box],
)
if __name__ == '__main__':
is_space = bool(os.getenv('SPACE_ID'))
server_name = '0.0.0.0' if is_space else '127.0.0.1'
demo.launch(server_name=server_name, share=False, debug=False)
|