Update app.py

app.py

@@ -5,6 +5,7 @@ import markdown
 import matplotlib.pyplot as plt
 import io
 import base64
+from fpdf import FPDF
 
 # Preload models
 models = {
@@ -35,9 +36,53 @@ def generate_score_chart(score):
     buf.seek(0)
     return base64.b64encode(buf.getvalue()).decode()
 
-def generate_report(answer, score, score_explanation, score_chart):
-    report = f"### Answer:\n\n{answer}\n\n### Confidence Score: {score}\n\n### Score Explanation:\n\n{score_explanation}\n\n![Score Chart](data:image/png;base64,{score_chart})"
-    return report
+def highlight_relevant_text(context, start, end):
+    highlighted_text = (
+        context[:start] + 
+        '<mark style="background-color: yellow;">' + 
+        context[start:end] + 
+        '</mark>' + 
+        context[end:]
+    )
+    return highlighted_text
+
+def generate_pdf_report(question, answer, score, score_explanation, score_chart, highlighted_context):
+    pdf = FPDF()
+    pdf.add_page()
+
+    pdf.set_font("Arial", size=12)
+    pdf.multi_cell(0, 10, f"Question: {question}")
+    pdf.ln()
+
+    pdf.set_font("Arial", size=12)
+    pdf.multi_cell(0, 10, f"Answer: {answer}")
+    pdf.ln()
+
+    pdf.set_font("Arial", size=12)
+    pdf.multi_cell(0, 10, f"Confidence Score: {score}")
+    pdf.ln()
+
+    pdf.set_font("Arial", size=12)
+    pdf.multi_cell(0, 10, f"Score Explanation: {score_explanation}")
+    pdf.ln()
+
+    pdf.set_font("Arial", size=12)
+    pdf.multi_cell(0, 10, "Highlighted Context:")
+    pdf.ln()
+    pdf.set_font("Arial", size=10)
+    pdf.multi_cell(0, 10, highlighted_context)
+    pdf.ln()
+
+    # Add score chart image to PDF
+    score_chart_image = io.BytesIO(base64.b64decode(score_chart))
+    pdf.image(score_chart_image, x=10, y=pdf.get_y(), w=100)
+
+    # Save PDF to memory
+    pdf_output = io.BytesIO()
+    pdf.output(pdf_output)
+    pdf_output.seek(0)
+
+    return pdf_output
 
 def answer_question(model_name, file, question, status):
     status = "Loading model..."
@@ -61,6 +106,11 @@ def answer_question(model_name, file, question, status):
     result = model(question=question, context=context)
     answer = result['answer']
     score = result['score']
+    start = result['start']
+    end = result['end']
+    
+    # Highlight relevant text
+    highlighted_context = highlight_relevant_text(context, start, end)
     
     # Generate the score chart
     score_chart = generate_score_chart(score)
@@ -68,11 +118,11 @@ def answer_question(model_name, file, question, status):
     # Explain score
     score_explanation = f"The confidence score ranges from 0 to 1, where a higher score indicates higher confidence in the answer's correctness. In this case, the score is {score:.2f}. A score closer to 1 implies the model is very confident about the answer."
     
-    # Generate the report
-    report = generate_report(answer, f"{score:.2f}", score_explanation, score_chart)
+    # Generate the PDF report
+    pdf_report = generate_pdf_report(question, answer, f"{score:.2f}", score_explanation, score_chart, highlighted_context)
     
     status = "Model loaded"
-    return answer, f"{score:.2f}", score_explanation, score_chart, report, status
+    return highlighted_context, f"{score:.2f}", score_explanation, score_chart, pdf_report, status
 
 # Define the Gradio interface
 with gr.Blocks() as interface:
@@ -96,11 +146,11 @@ with gr.Blocks() as interface:
         question_input = gr.Textbox(lines=2, placeholder="Enter your question here...", label="Question")
     
     with gr.Row():
-        answer_output = gr.Textbox(label="Answer")
+        answer_output = gr.HTML(label="Highlighted Answer")
         score_output = gr.Textbox(label="Confidence Score")
         explanation_output = gr.Textbox(label="Score Explanation")
         chart_output = gr.Image(label="Score Chart")
-        report_output = gr.Markdown(label="Report")
+        pdf_output = gr.File(label="Download PDF Report")
     
     with gr.Row():
         submit_button = gr.Button("Submit")
@@ -113,7 +163,7 @@ with gr.Blocks() as interface:
     submit_button.click(
         on_submit,
         inputs=[model_dropdown, file_input, question_input],
-        outputs=[answer_output, score_output, explanation_output, chart_output, report_output, status_output]
+        outputs=[answer_output, score_output, explanation_output, chart_output, pdf_output, status_output]
     )
 
 if __name__ == "__main__":