tensorsoft/Mini-RAG-Chat-With-Your-Files-CPU-Only

1import gradio as gr2from retriever import TextRetriever3from generator import AnswerGenerator4 5# Initialize retriever and generator6retriever = TextRetriever()7generator = AnswerGenerator()8 9def process_document(document: str):10    """Process and store document embeddings."""11    retriever.add_document(document)12    return "Document processed successfully!"13 14def ask_question(question: str):15    """Retrieve relevant chunks and generate an answer."""16    if not retriever.chunks:17        return "Please process a document first.", ""18    19    # Retrieve top chunks20    relevant_chunks = retriever.retrieve(question, top_k=3)21    context = " ".join(relevant_chunks)22    23    # Generate answer24    answer = generator.generate_answer(context, question)25    26    # Format retrieved chunks for display27    chunks_display = "\n\n".join([f"Chunk {i+1}: {chunk}" for i, chunk in enumerate(relevant_chunks)])28    29    return answer, chunks_display30 31# Gradio interface32with gr.Blocks() as demo:33    gr.Markdown("# RAG Demo")34    35    with gr.Tab("Upload / Paste Document"):36        document_input = gr.Textbox(lines=10, placeholder="Paste your document here...")37        process_button = gr.Button("Process Document")38        process_output = gr.Textbox(label="Status")39        process_button.click(40            fn=process_document,41            inputs=document_input,42            outputs=process_output43        )44    45    with gr.Tab("Ask a Question"):46        question_input = gr.Textbox(lines=2, placeholder="Enter your question...")47        ask_button = gr.Button("Ask")48        answer_output = gr.Textbox(label="Answer")49        with gr.Accordion("Retrieved Chunks", open=False):50            chunks_output = gr.Textbox(label="Relevant Document Chunks")51        ask_button.click(52            fn=ask_question,53            inputs=question_input,54            outputs=[answer_output, chunks_output]55        )56 57# Launch the app58if __name__ == "__main__":59    demo.launch()