Spaces:
Running
Running
File size: 1,830 Bytes
6ff89e0 daee42b 6ff89e0 daee42b |
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 |
import gradio as gr
from database.vaccine_coverage_db import VaccineCoverageDB
from rag.rag_pipeline import RAGPipeline
from utils.helpers import process_response
from config import DB_PATH, METADATA_FILE, PDF_DIR
# Initialize database and RAG pipeline
db = VaccineCoverageDB(DB_PATH)
rag = RAGPipeline(METADATA_FILE, PDF_DIR, use_semantic_splitter=True)
def query_rag(question, prompt_type):
if prompt_type == "Highlight":
response = rag.query(question, prompt_type="highlight")
else:
response = rag.query(question, prompt_type="evidence_based")
processed = process_response(response)
return processed["markdown"]
def save_pdf(item_key):
attachments = db.get_attachments_for_item(item_key)
if attachments:
attachment_key = attachments[0]["key"]
output_path = f"{attachment_key}.pdf"
if db.save_pdf_to_file(attachment_key, output_path):
return f"PDF saved successfully to {output_path}"
return "Failed to save PDF or no attachments found"
# Gradio interface
with gr.Blocks() as demo:
gr.Markdown("# Vaccine Coverage Study RAG System")
with gr.Tab("Query"):
question_input = gr.Textbox(label="Enter your question")
prompt_type = gr.Radio(["Highlight", "Evidence-based"], label="Prompt Type")
query_button = gr.Button("Submit Query")
output = gr.Markdown(label="Response")
query_button.click(
query_rag, inputs=[question_input, prompt_type], outputs=output
)
with gr.Tab("Save PDF"):
item_key_input = gr.Textbox(label="Enter item key")
save_button = gr.Button("Save PDF")
save_output = gr.Textbox(label="Save Result")
save_button.click(save_pdf, inputs=item_key_input, outputs=save_output)
if __name__ == "__main__":
demo.launch()
|