app.py

import gradio as gr
from data_loader import load_documents
from chunker import build_chunk_dataset
from rag_pipeline import RAGSystem


def init_rag():
    docs = load_documents("ag_news")
    chunks = build_chunk_dataset(docs[:1000])  # обмеження для швидкості
    rag = RAGSystem(chunks)
    return rag


def ask_rag(query, api_key, use_bm25, use_dense, rag):
    if not api_key:
        return "Please provide Groq API Key", ""

    if not query:
        return "Please enter a question", ""

    answer, refs = rag.run(
        query=query,
        api_key=api_key,
        use_bm25=use_bm25,
        use_dense=use_dense
    )
    return answer, refs


with gr.Blocks(title="RAG News System (Groq)") as demo:
    gr.Markdown("""
    **RAG News Question Answering**
    **Retrieval-Augmented Generation system over AG News**

    - BM25 + Dense Retrieval
    - Cross-Encoder Reranking
    - Llama-3.1 via Groq
    """)

    # State
    rag_state = gr.State()

    with gr.Row():
        with gr.Column(scale=2):
            query = gr.Textbox(
                label="Question (English)",
                placeholder="Example: What are the latest technology trends?",
                lines=2
            )

            answer = gr.Markdown(label="Answer")

        with gr.Column(scale=1):
            api_key = gr.Textbox(
                label="Groq API Key",
                placeholder="gsk_...",
                type="password"
            )

            use_bm25 = gr.Checkbox(value=True, label="BM25 (Keyword Search)")
            use_dense = gr.Checkbox(value=True, label="Dense (Semantic Search)")

            ask_btn = gr.Button("Ask Groq")

            refs = gr.Textbox(
                label="Sources",
                lines=8
            )

    demo.load(
        fn=init_rag,
        outputs=rag_state
    )

    ask_btn.click(
        fn=ask_rag,
        inputs=[query, api_key, use_bm25, use_dense, rag_state],
        outputs=[answer, refs]
    )

if __name__ == "__main__":
    demo.launch()