Update app.py

app.py

@@ -3,41 +3,65 @@ import gradio as gr
 import google.generativeai as genai
 from markdown_pdf import MarkdownPdf, Section
 import subprocess
-import json
-import traceback
-import re
-import concurrent.futures
-from pdf2image import convert_from_path
-from PIL import Image, ImageDraw, ImageFont
-import cv2
-import numpy as np
-import img2pdf
 
 # ---------- PROMPTS ----------
 PROMPTS = {
     "ALIGNMENT_PROMPT": {
         "role": "system",
-        "content": """Developer: Extract and align QP, MS, and AS into JSON. 
-## Instructions
-Output must be a valid JSON array, with one object per (sub-)question, in order. 
-Each object must have exactly these keys:
-- "question_number": string (e.g., "1", "1(a)", "2(b)(ii)")  
-- "qp": string (exact question text or "[Not found]")  
-- "ms": string (relevant markscheme text or "[Not found]")  
-- "as": string (final cleaned student answer; "[No response]" or "[illegible]" if needed)  
-### Numbering Rules
-- Always use **logical order of questions** (1, 2, 3, …) regardless of how they are labeled in the PDF.  
-- If the QP shows a mismatch (e.g., under "Question 1" the serial number says "12"), **still treat it as Q1**.  
-- Subparts must be written in standard form (e.g., "1(a)", "1(b)(ii)").  
-### Formatting Rules
-- Preserve math inside fenced code ```...```.  
-- If diagram/graph missing, write "[Graph omitted]".  
-- Do not add extra commentary outside JSON.  
-"""
+        "content": """Developer: Align QP, MS, and AS into structured JSON format.
+
+## Instructions:
+- Each question must include:
+  - `id` (question/sub-question number, e.g., "1", "2.a")
+  - `qp` (exact question wording)
+  - `total_marks` (integer)
+  - `ms` (markscheme with mark IDs and descriptions)
+  - `as` (student’s steps, numerical values, and notes)
+- Include `total_verification` in MS showing explicit mark breakdown.
+- The structure must be **valid JSON only**.
+
+## Example JSON:
+{
+  "questions": [
+    {
+      "id": "1",
+      "qp": "Ramiro walks to work each morning. During the first minute he walks 80 metres. In each subsequent minute he walks 90% of the distance walked during the previous minute.\\nThe distance between his house and work is 660 metres. Ramiro leaves his house at 08:00 and has to be at work by 08:15.\\nExplain why he will not be at work on time.",
+      "total_marks": 7,
+      "ms": {
+        "marks": [
+          { "id": "M1_1", "desc": "Recognise that the distance each minute forms a geometric sequence; show r = 0.9 (method mark)." },
+          { "id": "M1_2", "desc": "Recognise that total distance is the sum of a geometric sequence and give the sum formula (method mark)." },
+          { "id": "M1_3", "desc": "List at least 5 correct terms of the GP (method mark)." },
+          { "id": "A1_list", "desc": "List all 15 correct terms (accuracy mark)." },
+          { "id": "M1_4", "desc": "Attempt to find S_15 (method mark)." },
+          { "id": "A1_sum", "desc": "Correct numerical value for S_15 ≈ 635.287 (accuracy mark)." },
+          { "id": "R1", "desc": "Conclude: since S < 660, he will not be there on time (requires preceding A mark)." }
+        ],
+        "total_verification": "M1 + M1 + M1 + A1 + M1 + A1 + R1 = 7"
+      },
+      "as": {
+        "steps": [
+          "90% of 80 = 72 (2nd minute).",
+          "90% of 72 = 64.8 (3rd minute).",
+          "Sequence shown: 80, 72, 64.8, 58.32.",
+          "r = 72/80 = 0.9 ; also 64.8/72 = 0.9.",
+          "u_n = u_1 * r^(n-1).",
+          "S_n = u_1 * (r^n - 1)/(r - 1).",
+          "S_15 = 80 * (0.9^15 - 1)/(0.9 - 1).",
+          "S_15 = 635.29 (approx)."
+        ],
+        "numeric_S15": 635.29,
+        "notes": "Student found r and used the sum formula correctly, listed only 4 terms, got S15 ≈ 635.29 but did not explicitly state the final conclusion."
+      }
+    }
+  ]
+}"""
     },
+
     "GRADING_PROMPT": {
         "role": "system",
         "content": """Developer: You are an official examiner. Apply the following grading rules precisely.
+
 ### Abbreviations:
 - **M**: Marks for Method
 - **A**: Marks for Accuracy/Answer
@@ -45,6 +69,7 @@ Each object must have exactly these keys:
 - **AG**: Answer given in question—no marks
 - **FT**: Follow Through marks (if error carried forward correctly)
 - **MR**: Deduct for misread (once only)
+
 ---
 ## Grading Instructions
 1. Award marks using official annotations (e.g., M1, A2).  
@@ -54,14 +79,36 @@ Each object must have exactly these keys:
 5. Apply FT where appropriate.  
 6. Use proper notation: M1A0, A1, etc.  
 7. Any lost mark: use red `<span style="color:red">M0</span>` and make Reason red.  
+
 ---
 ## Output Format
-1. Produce a GitHub-flavored Markdown table with columns:  
-   | Student wrote | Marks Awarded | Reason |  
-   - Each row = one markable step/point, in order.  
-   - For blanks: “(no answer)” with marks lost.  
-2. After the table, write ONLY one line for total marks in the form:  Final Marks: X / Y
-"""
+Produce two sections per question/sub-question:
+
+---
+## Question X (and sub-question if applicable)
+
+### Markscheme vs Student Answer
+| Mark ID | Markscheme Expectation | Student’s Response | Awarded |
+|---------|------------------------|--------------------|---------|
+| M1_1    | Recognise GP, r=0.9    | "r = 72/80 = 0.9" | M1 |
+| M1_2    | Sum formula for GP     | "S_n = u1(r^n-1)/(r-1)" | M1 |
+| A1_list | 15 terms listed        | Only 4 terms shown | <span style="color:red">A0</span> |
+| …       | …                      | …                  | … |
+
+➡️ **Total: 6/7**
+
+---
+
+### Examiner’s Report
+At the very end, provide a summary table:
+
+| Question Number | Marks | Remark |
+|-----------------|-------|--------|
+| 1               | 6/7   | C      |
+| 2.a             | 9/9   | A      |
+
+Then show total clearly:  
+`Total: 15/16`"""
     }
 }
 
@@ -70,7 +117,6 @@ genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
 # ---------- HELPER: Save to PDF ----------
 def save_as_pdf(text, filename="output.pdf"):
-    print(f"📄 Saving grading report to PDF → {filename}")
     pdf = MarkdownPdf()
     pdf.add_section(Section(text, toc=False))
     pdf.save(filename)
@@ -78,17 +124,14 @@ def save_as_pdf(text, filename="output.pdf"):
 
 # ---------- HELPER: Compress PDF ----------
 def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
-    print(f"🗜️ Checking if compression needed for {input_path}...")
     if output_path is None:
         base, ext = os.path.splitext(input_path)
         output_path = f"{base}_compressed{ext}"
 
     if os.path.getsize(input_path) <= max_size:
-        print("✅ No compression needed")
         return input_path
 
     try:
-        print(f"⚡ Compressing {input_path} → {output_path}")
         gs_cmd = [
             "gs", "-sDEVICE=pdfwrite",
             "-dCompatibilityLevel=1.4",
@@ -98,13 +141,13 @@ def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
         ]
         subprocess.run(gs_cmd, check=True)
         if os.path.getsize(output_path) <= max_size:
-            print("✅ Compression successful")
+            print(f"✅ Compressed {input_path} → {output_path}")
             return output_path
         else:
-            print("⚠️ Compression didn’t shrink enough, using original")
+            print(f"⚠️ Compression failed to reduce below {max_size/1024/1024} MB")
             return input_path
     except Exception as e:
-        print(f"❌ Compression failed: {e}")
+        print(f"⚠️ Compression error: {e}")
         return input_path
 
 # ---------- HELPER: Create Model with Fallback ----------
@@ -116,207 +159,69 @@ def create_model():
         print("⚡ Falling back to gemini-2.5-flash model")
         return genai.GenerativeModel("gemini-2.5-flash", generation_config={"temperature": 0})
 
-# ---------- HELPER: Clean JSON Output ----------
-def clean_json_output(raw_text: str) -> str:
-    if not raw_text:
-        return ""
-    cleaned = re.sub(r"```(?:json)?", "", raw_text, flags=re.IGNORECASE)
-    cleaned = re.sub(r"(Raw output|json|JSON)[:\s]*", "", cleaned, flags=re.IGNORECASE)
-    cleaned = cleaned.strip()
-    cleaned = re.sub(r"```$", "", cleaned)
-    return cleaned
-
 # ---------- PIPELINE: ALIGN + GRADE ----------
-def align_and_grade(qp_file, ms_file, ans_file, imprint=False):
+def align_and_grade(qp_file, ms_file, ans_file):
     try:
-        print("\n🚀 Starting alignment + grading pipeline")
-
-        # Step 0: Compress
-        print("🔍 Step 0: Compressing PDFs...")
         qp_file = compress_pdf(qp_file, "qp_compressed.pdf")
         ms_file = compress_pdf(ms_file, "ms_compressed.pdf")
         ans_file = compress_pdf(ans_file, "ans_compressed.pdf")
 
-        # Step 1: Uploads
-        print("📤 Step 1: Uploading PDFs to Gemini...")
         qp_uploaded = genai.upload_file(path=qp_file, display_name="Question Paper")
         ms_uploaded = genai.upload_file(path=ms_file, display_name="Markscheme")
         ans_uploaded = genai.upload_file(path=ans_file, display_name="Answer Sheet")
 
         model = create_model()
 
-        # Step 2: Alignment
-        print("🧩 Step 2: Aligning QP, MS, and AS...")
+        # ---------------- STEP 1: ALIGN (JSON only) ----------------
         resp = model.generate_content([
             PROMPTS["ALIGNMENT_PROMPT"]["content"],
             qp_uploaded,
             ms_uploaded,
             ans_uploaded
         ])
-        aligned_json = getattr(resp, "text", None)
-        if not aligned_json and resp.candidates:
-            aligned_json = resp.candidates[0].content.parts[0].text
-
-        aligned_json = clean_json_output(aligned_json)
-        questions = json.loads(aligned_json)
-        print(f"✅ Parsed JSON with {len(questions)} questions")
-
-        # Step 3: Grading
-        print("📝 Step 3: Grading each question...")
-
-        def grade_one(idx_q):
-            idx, q = idx_q
-            print(f"   ➡️ Grading Question {q['question_number']}")
-            q_json = json.dumps(q, indent=2)
-            response = model.generate_content([
-                PROMPTS["GRADING_PROMPT"]["content"],
-                q_json
-            ])
-            grading_piece = getattr(response, "text", None)
-            if not grading_piece and response.candidates:
-                grading_piece = response.candidates[0].content.parts[0].text
-            return idx, q["question_number"], grading_piece
-
-        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
-            results = list(executor.map(grade_one, enumerate(questions)))
-        results.sort(key=lambda x: x[0])
-
-        # Step 4: Build report
-        print("📊 Step 4: Building grading report...")
-        grading_sections = []
-        grading_json = {"grading": []}
-        for _, qnum, grading_piece in results:
-            section = f"## Question {qnum}\n\n{grading_piece}"
-            grading_sections.append(section)
-
-            marks_list = re.findall(r"(M[01]|A[0-9]|R[01])", grading_piece)
-            grading_json["grading"].append({"question": qnum, "marks_awarded": marks_list})
-
-        grading_report = "\n\n".join(grading_sections)
-        base_name = os.path.splitext(os.path.basename(ans_file))[0]
-        grading_pdf_path = save_as_pdf(grading_report, f"{base_name}_graded.pdf")
+        json_output = getattr(resp, "text", None)
+        if not json_output and resp.candidates:
+            json_output = resp.candidates[0].content.parts[0].text
+
+        # ---------------- STEP 2: GRADING (Markdown + PDF) ----------------
+        response = model.generate_content([
+            PROMPTS["GRADING_PROMPT"]["content"],
+            json_output
+        ])
+        grading = getattr(response, "text", None)
+        if not grading and response.candidates:
+            grading = response.candidates[0].content.parts[0].text
 
-        imprint_pdf_path = None
-        if imprint:
-            print("✍ Step 5: Imprinting marks onto answer sheet...")
-            imprint_pdf_path = imprint_marks(ans_file, grading_json, model)
+        base_name = os.path.splitext(os.path.basename(ans_file))[0]
+        grading_pdf_path = save_as_pdf(grading, f"{base_name}_graded.pdf")
 
-        print("✅ Pipeline completed successfully")
-        return json.dumps(questions, indent=2), grading_report, grading_pdf_path, imprint_pdf_path
+        # Return JSON (alignment), Markdown grading, and PDF
+        return json_output, grading, grading_pdf_path
 
     except Exception as e:
-        print("❌ Fatal error in pipeline")
-        traceback.print_exc()
-        return f"❌ Error: {e}", None, None, None
-
-# ---------- PIPELINE: IMPRINT MARKS ----------
-def imprint_marks(ans_pdf, grading_json, model, grid_rows=20, grid_cols=14):
-    print("📄 Converting answer sheet to images with grid...")
-    output_dir = "grid_pages"
-    os.makedirs(output_dir, exist_ok=True)
-    pages = convert_from_path(ans_pdf, dpi=200)
-    page_images = []
-
-    for i, page in enumerate(pages):
-        img_path = os.path.join(output_dir, f"page_{i+1}_grid.png")
-        img = page.convert("RGB")
-        draw = ImageDraw.Draw(img)
-        w, h = img.size
-        cell_w, cell_h = w / grid_cols, h / grid_rows
-
-        try:
-            num_font = ImageFont.truetype("arial.ttf", 20)
-        except IOError:
-            num_font = ImageFont.load_default()
-
-        cell_num = 1
-        for r in range(grid_rows):
-            for c in range(grid_cols):
-                x = int(c * cell_w + cell_w / 2)
-                y = int(r * cell_h + cell_h / 2)
-                text = str(cell_num)
-                bbox = draw.textbbox((0, 0), text, font=num_font)
-                tw = bbox[2] - bbox[0]
-                th = bbox[3] - bbox[1]
-                draw.text((x - tw/2, y - th/2), text, fill="black", font=num_font)
-                cell_num += 1
-        img.save(img_path, "PNG")
-        page_images.append(img_path)
-    print("✅ Grid images prepared")
-
-    annotated_pages = []
-    for idx, page in enumerate(pages):
-        print(f"🔎 Asking Gemini for mapping on page {idx+1}...")
-        prompt = f"""
-You are an exam marker. The page is divided into a {grid_rows} x {grid_cols} grid with numbered cells.
-Return JSON: [{{"question": "1(a)", "cell_number": 15}}, ...]
-Grading JSON:
-{json.dumps(grading_json, indent=2)}
-"""
-        response = model.generate_content([prompt, Image.open(page_images[idx])])
-        mapping_text = getattr(response, "text", "")
-        match = re.search(r'\[.*\]', mapping_text, re.DOTALL)
-        mapping = json.loads(match.group(0)) if match else []
-        print(f"   ↪ Gemini returned {len(mapping)} mappings")
-
-        # Annotate
-        img = np.array(page.convert("RGB"))
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-        h, w, _ = img.shape
-        cell_w, cell_h = w / grid_cols, h / grid_rows
-
-        for item in mapping:
-            q = item["question"]
-            cell_number = item["cell_number"]
-            row = (cell_number - 1) // grid_cols
-            col = (cell_number - 1) % grid_cols
-
-            marks_list = next((g["marks_awarded"] for g in grading_json["grading"] if g["question"] == q), [])
-            marks_text = ",".join(marks_list)
-
-            x_c = int((col+1) * cell_w - cell_w/4)
-            y_c = int((row+0.5) * cell_h)
-            cv2.putText(img, marks_text, (x_c, y_c), cv2.FONT_HERSHEY_SIMPLEX,
-                        1.5, (0, 0, 255), 3, cv2.LINE_AA)
-
-        annotated_path = os.path.join(output_dir, f"annotated_{idx+1}.png")
-        cv2.imwrite(annotated_path, img)
-        annotated_pages.append(annotated_path)
-        print(f"🖊 Marks imprinted for page {idx+1}")
-
-    output_pdf = "answer_sheet_with_marks.pdf"
-    with open(output_pdf, "wb") as f:
-        f.write(img2pdf.convert(annotated_pages))
-    print(f"✅ Final imprinted PDF saved: {output_pdf}")
-    return output_pdf
+        return f"❌ Error: {e}", None, None
 
 # ---------- GRADIO APP ----------
-with gr.Blocks(title="LeadIB AI Grading with Optional Imprinting") as demo:
-    gr.Markdown("## LeadIB AI Grading\nUpload QP, MS, and AS. Get aligned JSON, grading report, and optionally imprint marks on the answer sheet.")
+with gr.Blocks(title="LeadIB AI Grading (Alignment + Auto-Grading)") as demo:
+    gr.Markdown("## 📘 LeadIB AI Grading\nUpload **Question Paper**, **Markscheme**, and **Student Answer Sheet**.\nThe system will first align into JSON, then auto-grade with detailed feedback.")
 
     with gr.Row():
-        qp_file = gr.File(label="Upload Question Paper (PDF)", type="filepath")
-        ms_file = gr.File(label="Upload Markscheme (PDF)", type="filepath")
-        ans_file = gr.File(label="Upload Student Answer Sheet (PDF)", type="filepath")
+        qp_file = gr.File(label="📄 Upload Question Paper (PDF)")
+        ms_file = gr.File(label="📄 Upload Markscheme (PDF)")
+        ans_file = gr.File(label="📝 Upload Student Answer Sheet (PDF)")
 
-    imprint_opt = gr.Checkbox(label="Imprint Marks on Answer Sheet?", value=False)
-    run_btn = gr.Button("Start Alignment + Auto-Grading")
+    run_button = gr.Button("🚀 Run Alignment + Grading")
 
     with gr.Row():
-        aligned_out = gr.Textbox(label="📄 Aligned QP | MS | AS (JSON)", lines=20)
+        json_output = gr.Textbox(label="📑 Step 1: Alignment (JSON)", lines=20)
+        grading_output = gr.Textbox(label="📝 Step 2: Grading (Markdown)", lines=20)
 
-    with gr.Row():
-        grading_out = gr.Textbox(label="✅ Grading Report (Markdown)", lines=20)
-
-    with gr.Row():
-        grading_pdf = gr.File(label="⬇️ Download Grading Report (PDF)")
-        imprint_pdf = gr.File(label="⬇️ Download Answer Sheet with Imprinted Marks (PDF)")
+    grading_pdf = gr.File(label="📥 Download Grading PDF")
 
-    run_btn.click(
+    run_button.click(
         fn=align_and_grade,
-        inputs=[qp_file, ms_file, ans_file, imprint_opt],
-        outputs=[aligned_out, grading_out, grading_pdf, imprint_pdf],
-        show_progress=True
+        inputs=[qp_file, ms_file, ans_file],
+        outputs=[json_output, grading_output, grading_pdf]
     )
 
 if __name__ == "__main__":