Spaces:
Sleeping
Sleeping
| 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) | |