Update app.py

app.py

@@ -1,96 +1,56 @@
 import os
 import gradio as gr
 import google.generativeai as genai
-from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
+from reportlab.platypus import SimpleDocTemplate, Paragraph
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.lib.pagesizes import A4
-from reportlab.lib.units import cm
 
 # Configure Gemini API
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
 # ---------- PROMPTS ----------
-TRANSCRIPTION_PROMPT = """
-Persona:
-You are an expert transcriptionist specializing in scientific and mathematical documents. 
-Your primary goal is to convert handwritten mathematical work into a perfectly formatted, 
-machine-readable Markdown document using LaTeX for all mathematical notation.
+TRANSCRIPTION_PROMPT = """ ... """  # unchanged
+GRADING_PROMPT = """ ... """        # unchanged
 
-Rules:
-- Transcribe exactly what is written, do not correct errors.
-- Use $...$ for inline math, $$...$$ for block math.
-- Ignore struck-through text.
-- Preserve structure: bold for Q numbers (**1.**), step-by-step math with \\begin{align*}.
-- If a symbol is ambiguous, mark as [x?].
-Output must be a clean Markdown string.
-"""
-
-GRADING_PROMPT = """
-You are an official examiner. Grade the student transcription using the question paper 
-and the official marking scheme.
-
-Rules:
-1. Apply marks exactly as per the markscheme (M1, A1, etc.).
-2. M marks must be earned before A marks.
-3. Justify each awarded or withheld mark with clear reasoning.
-4. Classify all errors as Conceptual Error, Silly Mistake, or None.
-5. Follow dependency between M and A strictly.
-6. Do not give marks outside the markscheme.
-Output must be a structured grading report with reasoning.
-"""
+# ---------- HELPER: Save to PDF ----------
+def save_as_pdf(text, filename="output.pdf"):
+    styles = getSampleStyleSheet()
+    doc = SimpleDocTemplate(filename, pagesize=A4)
+    story = [Paragraph(p, styles["Normal"]) for p in text.split("\n")]
+    doc.build(story)
+    return filename
 
 # ---------- STEP 1: TRANSCRIPTION ----------
 def transcribe(ans_file):
-    yield "⏳ Processing transcription..."
     try:
         ans_uploaded = genai.upload_file(path=ans_file.name, display_name="Answer Sheet")
         model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
+
         resp = model.generate_content([TRANSCRIPTION_PROMPT, ans_uploaded])
+
         transcription = getattr(resp, "text", None) or resp.candidates[0].content.parts[0].text
-        yield transcription
+        pdf_path = save_as_pdf(transcription, "transcription.pdf")
+        return transcription, pdf_path
     except Exception as e:
-        yield f"❌ Error during transcription: {e}"
+        return f"❌ Error during transcription: {e}", None
 
 # ---------- STEP 2: GRADING ----------
 def grade(qp_file, ms_file, transcription):
-    yield "⏳ Grading in progress..."
     try:
         qp_uploaded = genai.upload_file(path=qp_file.name, display_name="Question Paper")
         ms_uploaded = genai.upload_file(path=ms_file.name, display_name="Marking Scheme")
-
         model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
-        resp = model.generate_content([GRADING_PROMPT, qp_uploaded, ms_uploaded, transcription])
-        grading = getattr(resp, "text", None) or resp.candidates[0].content.parts[0].text
-
-        # Save grading report as PDF
-        pdf_path = "/tmp/grading_report.pdf"
-        doc = SimpleDocTemplate(pdf_path, pagesize=A4)
-        styles = getSampleStyleSheet()
-        elements = []
 
-        elements.append(Paragraph("📘 AI Teacher Assistant - Grading Report", styles["Title"]))
-        elements.append(Spacer(1, 0.5*cm))
-        for para in grading.split("\n"):
-            if para.strip():
-                elements.append(Paragraph(para, styles["Normal"]))
-                elements.append(Spacer(1, 0.2*cm))
-
-        doc.build(elements)
+        resp = model.generate_content([GRADING_PROMPT, qp_uploaded, ms_uploaded, transcription])
 
-        yield grading, pdf_path
+        grading = getattr(resp, "text", None) or resp.candidates[0].content.parts[0].text
+        pdf_path = save_as_pdf(grading, "grading.pdf")
+        return grading, pdf_path
     except Exception as e:
-        yield f"❌ Error during grading: {e}", None
+        return f"❌ Error during grading: {e}", None
 
 # ---------- GRADIO APP ----------
 with gr.Blocks(title="📘 AI Teacher Assistant") as demo:
-    # Inject MathJax for LaTeX rendering
-    gr.HTML("""
-    <script src="https://polyfill.io/v3/polyfill.min.js?features=es6"></script>
-    <script id="MathJax-script" async
-      src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js">
-    </script>
-    """)
-
     gr.Markdown("## 📘 AI Teacher Assistant\nUpload exam documents to transcribe and grade student answers step by step.")
 
     with gr.Row():
@@ -100,16 +60,29 @@ with gr.Blocks(title="📘 AI Teacher Assistant") as demo:
 
     # Step 1: Transcription
     transcribe_btn = gr.Button("Step 1: Transcribe Answer Sheet")
-    transcription_out = gr.Markdown(label="📄 Student Transcription")
+    with gr.Row():
+        transcription_out = gr.Markdown(label="📄 Student Transcription")
+        transcription_pdf = gr.File(label="⬇️ Download Transcription (PDF)")
 
     # Step 2: Grading
     grade_btn = gr.Button("Step 2: Grade the Student")
-    grading_out = gr.Textbox(label="✅ Grading Report (Step-by-Step)", lines=20)
-    download_pdf = gr.File(label="⬇️ Download Grading Report (PDF)")
-
-    # Button Logic
-    transcribe_btn.click(fn=transcribe, inputs=[ans_file], outputs=[transcription_out])
-    grade_btn.click(fn=grade, inputs=[qp_file, ms_file, transcription_out], outputs=[grading_out, download_pdf])
+    with gr.Row():
+        grading_out = gr.Textbox(label="✅ Grading Report (Step-by-Step)", lines=20)
+        grading_pdf = gr.File(label="⬇️ Download Grading (PDF)")
+
+    # Button Logic with loading indicators
+    transcribe_btn.click(
+        fn=transcribe,
+        inputs=[ans_file],
+        outputs=[transcription_out, transcription_pdf],
+        show_progress=True
+    )
+    grade_btn.click(
+        fn=grade,
+        inputs=[qp_file, ms_file, transcription_out],
+        outputs=[grading_out, grading_pdf],
+        show_progress=True
+    )
 
 if __name__ == "__main__":
     demo.launch()