Update app.py

app.py

@@ -29,6 +29,7 @@ TASK:
 1. Transcribe EXACTLY all the questions FIRST (with their total marks).   
 2. After ALL questions, transcribe the Markscheme exactly, preserving M/A/R notation in brackets.   
 3. Always number the questions sequentially (Question 1, Question 2, Question 3, …) **in the order they appear in the PDF**, even if the PDF shows a different number or leaves it blank. Do NOT skip or leave Question: blank. Never start a question other than question 1 ( even if it is labelled in pdf as 8 name it 1) 
+4. After the markscheme, DETECT and FLAG all questions in the markscheme where a graph/diagram is expected. For each, output the question number and the page number in the format below. 
 FORMAT:   
 ==== PAPER TOTAL MARKS ====   
 <total marks>   
@@ -52,7 +53,7 @@ Answer 2 :
 <exact MS for Q2 with notations>  
 (repeat for all answers)   
 ==== MARKSCHEME END ====   
-""" 
+==== GRAPH EXPECTED QUESTIONS ====\nGraph expected in:\n- Question <number> → Page <number>\n(one per line)\n==== END GRAPH EXPECTED ====\n""" 
 } 
 , 
 
@@ -94,8 +95,9 @@ At the very end, provide a summary table:
 Then show total clearly as a final line: 
 `Total: <obtained_marks>/<max_marks>` 
 NOTES: 
-- The assistant will receive two transcripts: (1) QP+MS transcription (questions then markscheme) and (2) AS transcription (student answers). Use the QP+MS transcript as the authoritative source of question wording, total marks, and verbatim markscheme entries (M/A/R mark IDs). 
+- The assistant will receive two transcripts: (1) QP+MS transcript (questions then markscheme) and (2) AS transcript (student answers). Use the QP+MS transcript as the authoritative source of question wording, total marks, and verbatim markscheme entries (M/A/R mark IDs). 
 - Match student answers to question IDs and grade according to the provided verbatim markscheme. 
+- For questions where a graph is expected and the student attempted a graph, you will be provided with the relevant markscheme and answer sheet graph images/pages. Use these for grading those questions with visual context. For all other questions, proceed as usual. 
 - Produce full markdown as above. Ensure mark IDs used in the grading are present and consistent with the markscheme. 
 """ 
     } 
@@ -221,32 +223,64 @@ def extract_question_ids_from_qpms(text):
         print("⚠️ No question IDs extracted; will send NA placeholder.") 
     return ids 
 
-def build_as_prompt_with_expected_ids(expected_ids): 
-    """ 
-    Construct the AS transcription prompt injecting the expected IDs block. 
-    """ 
-    if not expected_ids: 
-        ids_block = "{NA}" 
-    else: 
-        ids_block = "{\n" + "\n".join(expected_ids) + "\n}" 
-    prompt = f"""You are a high-quality handwritten transcription assistant. 
-INPUT: This PDF contains a student's handwritten answer sheet. 
-TASK: Transcribe the student's answers exactly (as text). Preserve step order and line breaks. Attempt to assign each answer to a question ID if the student has labelled it (e.g., "1", "1a", "2(b)", "3"). If the student hasn't labelled answers, segment contiguous answer blocks and attempt to infer question IDs from context — but mark inferred IDs clearly as "INFERRED: <id>" 
-Enclose all mathematical expressions in Markdown fenced code blocks (``` triple backticks). 
-If a diagram/graph is omitted, write [Graph omitted]. 
-Unreadable parts: [illegible]. 
-Unanswered: [No response]. 
-Do NOT recreate diagrams. 
-Ensure consistency and determinism in formatting so subsequent models can grade directly from this aligned format. 
-Expected questions (if missing, write NA): 
-{ids_block} 
------------------------ 
-OUTPUT FORMAT: 
-Question <id> 
-AS: 
-<transcribed answer or placeholder> 
-""" 
-    return prompt 
+# Update AS prompt builder to include graph detection
+
+def build_as_prompt_with_expected_ids(expected_ids):
+    """
+    Construct the AS transcription prompt injecting the expected IDs block and graph detection instructions.
+    """
+    if not expected_ids:
+        ids_block = "{NA}"
+    else:
+        ids_block = "{\n" + "\n".join(expected_ids) + "\n}"
+    prompt = f"""You are a high-quality handwritten transcription assistant.
+INPUT: This PDF contains a student's handwritten answer sheet.
+TASK: Transcribe the student's answers exactly (as text). Preserve step order and line breaks. Attempt to assign each answer to a question ID if the student has labelled it (e.g., "1", "1a", "2(b)", "3"). If the student hasn't labelled answers, segment contiguous answer blocks and attempt to infer question IDs from context — but mark inferred IDs clearly as "INFERRED: <id>"
+Enclose all mathematical expressions in Markdown fenced code blocks (``` triple backticks).
+If a diagram/graph is omitted, write [Graph omitted].
+Unreadable parts: [illegible].
+Unanswered: [No response].
+Do NOT recreate diagrams.
+Ensure consistency and determinism in formatting so subsequent models can grade directly from this aligned format.
+Expected questions (if missing, write NA):
+{ids_block}
+-----------------------
+OUTPUT FORMAT:
+Question <id>
+AS:
+<transcribed answer or placeholder>
+==== GRAPH FOUND ANSWERS ====\nGraph found in:\n- Answer <number> → Page <number>\n(one per line)\n==== END GRAPH FOUND ====\n"""
+    return prompt
+
+# Robust parsing functions for graph detection
+
+def extract_graph_questions_from_ms(ms_text):
+    """
+    Parse LLM output for Markscheme to extract questions/pages where a graph is expected.
+    Returns dict: {question_number: ms_page_number}
+    """
+    matches = re.findall(r"==== GRAPH EXPECTED QUESTIONS ====\\s*Graph expected in:(.*?)==== END GRAPH EXPECTED ====" , ms_text, re.DOTALL)
+    mapping = {}
+    if matches:
+        for line in matches[0].splitlines():
+            m = re.match(r"-\s*Question\s*(\d+)\s*[\u2192\-\:]\s*Page\s*(\d+)", line.strip())
+            if m:
+                mapping[int(m.group(1))] = int(m.group(2))
+    return mapping
+
+def extract_graph_answers_from_as(as_text):
+    """
+    Parse LLM output for Answer Sheet to extract answers/pages where a graph was found.
+    Returns dict: {answer_number: as_page_number}
+    """
+    matches = re.findall(r"==== GRAPH FOUND ANSWERS ====\\s*Graph found in:(.*?)==== END GRAPH FOUND ====" , as_text, re.DOTALL)
+    mapping = {}
+    if matches:
+        for line in matches[0].splitlines():
+            m = re.match(r"-\s*Answer\s*(\d+)\s*[\u2192\-\:]\s*Page\s*(\d+)", line.strip())
+            if m:
+                mapping[int(m.group(1))] = int(m.group(2))
+    return mapping
 
 def extract_marks_from_grading(grading_text): 
     """ 
@@ -286,7 +320,6 @@ def ask_gemini_for_mapping_for_page(model, image_path, grading_json, expected_id
 You are an exam marker. Your role is to identify where each question begins on the page.
 The page is divided into a {rows} x {cols} grid. Each cell has a RUNNING NUMBER label (1..{rows*cols}).
 For each question in the grading JSON, return the cell NUMBER where the FIRST STEP of that question begins.
-
 IMPORTANT: Only spot and return cell numbers for the following question IDs (one per line):
 {ids_block}
 If you see a sub-question (e.g., ii) above a main question (e.g., Q4), infer it belongs to the previous question (e.g., Q3.ii).
@@ -450,93 +483,124 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, expec
     print("📑 Imprinted PDF saved to:", compressed) 
     return compressed 
 
-# ---------------- MAIN PIPELINE ---------------- 
-def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False): 
-    """ 
-    Final pipeline implementing requested flow and verbose console logging. 
-    """ 
-    try: 
-        print("🔁 Starting pipeline...") 
-        # Step 0: compress as needed 
-        qp_path = compress_pdf(qp_path) 
-        ms_path = compress_pdf(ms_path) 
-        ans_path = compress_pdf(ans_path) 
-
-        # Merge QP + MS 
-        merged_qpms_path = os.path.splitext(qp_path)[0] + "_merged_qp_ms.pdf" 
-        merge_pdfs([qp_path, ms_path], merged_qpms_path) 
-        print("📎 Merged QP + MS ->", merged_qpms_path) 
-
-        # Upload files to Gemini 
-        print("🔼 Uploading files to Gemini...") 
-        merged_uploaded = genai.upload_file(path=merged_qpms_path, display_name="QP+MS (merged)") 
-        ans_uploaded = genai.upload_file(path=ans_path, display_name="Answer Sheet") 
-        print("✅ Upload complete.") 
-
-        # Create model and print which selected 
-        model = create_model() 
-
-        # Step 1.i: QP+MS transcription (first) 
-        print("1.i) Transcribing QP+MS (questions first, then full markscheme)...") 
-        qpms_prompt = PROMPTS["QP_MS_TRANSCRIPTION"]["content"] 
-        qpms_text = gemini_generate_content(model, qpms_prompt, file_upload_obj=merged_uploaded) 
-        print("📄 QP+MS transcription received. Saving debug file: debug_qpms_transcript.txt") 
-        with open("debug_qpms_transcript.txt", "w", encoding="utf-8") as f: 
-            f.write(qpms_text) 
-
-        # Step 2: extract serial numbers (question IDs) using regex from qpms_text 
-        extracted_ids = extract_question_ids_from_qpms(qpms_text) 
-        if not extracted_ids: 
-            extracted_ids = ["NA"] 
-
-        # Step 1.ii: Build AS prompt injecting extracted IDs and transcribe AS 
-        print("1.ii) Building AS transcription prompt with expected question IDs and sending to Gemini...") 
-        as_prompt = build_as_prompt_with_expected_ids(extracted_ids) 
-        as_text = gemini_generate_content(model, as_prompt, file_upload_obj=ans_uploaded) 
-        print("📝 AS transcription received. Saving debug file: debug_as_transcript.txt") 
-        with open("debug_as_transcript.txt", "w", encoding="utf-8") as f: 
-            f.write(as_text) 
-
-        # Step 3: Grading - send both transcripts to grading model 
-        print("2) Preparing grading input and sending to Gemini for grading...") 
-        grading_input = ( 
-            "=== QP+MS TRANSCRIPT BEGIN ===\n" 
-            + qpms_text 
-            + "\n=== QP+MS TRANSCRIPT END ===\n\n" 
-            + "=== ANSWER SHEET TRANSCRIPT BEGIN ===\n" 
-            + as_text 
-            + "\n=== ANSWER SHEET TRANSCRIPT END ===\n" 
-        ) 
-        grading_prompt_system = PROMPTS["GRADING_PROMPT"]["content"] 
-        grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input) 
-        print("🧾 Grading output received. Saving debug file: debug_grading.md") 
-        with open("debug_grading.md", "w", encoding="utf-8") as f: 
-            f.write(grading_text) 
-
-        # Save grading PDF 
-        base_name = os.path.splitext(os.path.basename(ans_path))[0] 
-        grading_pdf_path = save_as_pdf(grading_text, f"{base_name}_graded.pdf") 
-        print("📄 Grading PDF saved:", grading_pdf_path) 
-
-        # Step 4: Extract marks for imprinting 
-        grading_json = extract_marks_from_grading(grading_text) 
-        with open("debug_grading_json.json", "w", encoding="utf-8") as f: 
-            json.dump(grading_json, f, indent=2, ensure_ascii=False) 
-        print("🔧 Grading marks extraction complete.") 
-
-        imprinted_pdf_path = None 
-        if imprint: 
-            print("✍ Imprint option enabled. Starting imprinting process (parallel mapping requests)...") 
-            imprinted_pdf_path = f"{base_name}_imprinted.pdf" 
-            imprinted_pdf_path = imprint_marks_using_mapping(ans_path, grading_json, imprinted_pdf_path, model, extracted_ids) 
-            print("✅ Imprinting finished. Imprinted PDF at:", imprinted_pdf_path) 
-
-        print("🏁 Pipeline finished successfully.") 
-        return qpms_text, as_text, grading_text, grading_pdf_path, imprinted_pdf_path 
-
-    except Exception as e: 
-        print("❌ Pipeline error:", e) 
-        return f"❌ Error: {e}", None, None, None, None 
+# ---------------- GRAPH DETECTION HELPERS ----------------
+# These functions are now robustly handled by the new_code, so they are no longer needed here.
+
+# ---------------- PIPELINE UPDATE FOR GRAPH-AWARE GRADING ----------------
+def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
+    """
+    Final pipeline implementing requested flow and verbose console logging.
+    Now includes Graph-Aware Grading logic.
+    """
+    try:
+        print("🔁 Starting pipeline...")
+        # Step 0: compress as needed
+        qp_path = compress_pdf(qp_path)
+        ms_path = compress_pdf(ms_path)
+        ans_path = compress_pdf(ans_path)
+
+        # Merge QP + MS
+        merged_qpms_path = os.path.splitext(qp_path)[0] + "_merged_qp_ms.pdf"
+        merge_pdfs([qp_path, ms_path], merged_qpms_path)
+        print("📎 Merged QP + MS ->", merged_qpms_path)
+
+        # Upload files to Gemini
+        print("🔼 Uploading files to Gemini...")
+        merged_uploaded = genai.upload_file(path=merged_qpms_path, display_name="QP+MS (merged)")
+        ans_uploaded = genai.upload_file(path=ans_path, display_name="Answer Sheet")
+        print("✅ Upload complete.")
+
+        # Create model and print which selected
+        model = create_model()
+
+        # Step 1.i: QP+MS transcription (first)
+        print("1.i) Transcribing QP+MS (questions first, then full markscheme, with graph detection)...")
+        qpms_prompt = PROMPTS["QP_MS_TRANSCRIPTION"]["content"] + "\nAt the end, also list all questions in the markscheme where a graph is expected, in the format:\nGraph expected in:\n- Question <number> → Page <number>\n(One per line, after ==== MARKSCHEME END ====)"
+        qpms_text = gemini_generate_content(model, qpms_prompt, file_upload_obj=merged_uploaded)
+        print("📄 QP+MS transcription received. Saving debug file: debug_qpms_transcript.txt")
+        with open("debug_qpms_transcript.txt", "w", encoding="utf-8") as f:
+            f.write(qpms_text)
+
+        # Step 1.i.a: Extract graph-expected questions from MS
+        ms_graph_mapping = extract_graph_questions_from_ms(qpms_text)
+        print("🖼️ Graph-expected questions in MS:", ms_graph_mapping)
+
+        # Step 2: extract serial numbers (question IDs) using regex from qpms_text
+        extracted_ids = extract_question_ids_from_qpms(qpms_text)
+        if not extracted_ids:
+            extracted_ids = ["NA"]
+
+        # Step 1.ii: Build AS prompt injecting extracted IDs and transcribe AS
+        print("1.ii) Building AS transcription prompt with expected question IDs and graph detection, sending to Gemini...")
+        as_prompt = build_as_prompt_with_expected_ids(extracted_ids) + "\nAt the end, also list all answers where a graph is found, in the format:\nGraph found in:\n- Answer <number> → Page <number>\n(One per line, after all answers)"
+        as_text = gemini_generate_content(model, as_prompt, file_upload_obj=ans_uploaded)
+        print("📝 AS transcription received. Saving debug file: debug_as_transcript.txt")
+        with open("debug_as_transcript.txt", "w", encoding="utf-8") as f:
+            f.write(as_text)
+
+        # Step 2.a: Extract graph-attempted answers from AS
+        as_graph_mapping = extract_graph_answers_from_as(as_text)
+        print("🖼️ Graph-attempted answers in AS:", as_graph_mapping)
+
+        # Step 3: Graph Matching
+        graph_bundles = []
+        for ans_num, as_page in as_graph_mapping.items():
+            if ans_num in ms_graph_mapping:
+                graph_bundles.append({
+                    "question": ans_num,
+                    "ms_page": ms_graph_mapping[ans_num],
+                    "as_page": as_page
+                })
+        print("🔗 Graph bundles for grading:", graph_bundles)
+
+        # Step 4: Grading - send both transcripts to grading model, inject graph bundle info
+        print("2) Preparing grading input and sending to Gemini for grading...")
+        grading_input = (
+            "=== QP+MS TRANSCRIPT BEGIN ===\n"
+            + qpms_text
+            + "\n=== QP+MS TRANSCRIPT END ===\n\n"
+            + "=== ANSWER SHEET TRANSCRIPT BEGIN ===\n"
+            + as_text
+            + "\n=== ANSWER SHEET TRANSCRIPT END ===\n"
+        )
+        # Inject graph bundle note
+        if graph_bundles:
+            graph_note = "\n\n---\nFor the following questions, a graph was expected and the student attempted it. Please use the provided images for grading these questions:\n"
+            for bundle in graph_bundles:
+                graph_note += f"- Question {bundle['question']}:\n  - Markscheme graph (Page {bundle['ms_page']})\n  - Student’s graph (Page {bundle['as_page']})\n"
+            graph_note += "\nGrade these with visual context. For all other questions, proceed as usual.\n---\n"
+            grading_input += graph_note
+
+        grading_prompt_system = PROMPTS["GRADING_PROMPT"]["content"]
+        grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input)
+        print("🧾 Grading output received. Saving debug file: debug_grading.md")
+        with open("debug_grading.md", "w", encoding="utf-8") as f:
+            f.write(grading_text)
+
+        # Save grading PDF
+        base_name = os.path.splitext(os.path.basename(ans_path))[0]
+        grading_pdf_path = save_as_pdf(grading_text, f"{base_name}_graded.pdf")
+        print("📄 Grading PDF saved:", grading_pdf_path)
+
+        # Step 4: Extract marks for imprinting
+        grading_json = extract_marks_from_grading(grading_text)
+        with open("debug_grading_json.json", "w", encoding="utf-8") as f:
+            json.dump(grading_json, f, indent=2, ensure_ascii=False)
+        print("🔧 Grading marks extraction complete.")
+
+        imprinted_pdf_path = None
+        if imprint:
+            print("✍ Imprint option enabled. Starting imprinting process (parallel mapping requests)...")
+            imprinted_pdf_path = f"{base_name}_imprinted.pdf"
+            imprinted_pdf_path = imprint_marks_using_mapping(ans_path, grading_json, imprinted_pdf_path, model, extracted_ids)
+            print("✅ Imprinting finished. Imprinted PDF at:", imprinted_pdf_path)
+
+        print("🏁 Pipeline finished successfully.")
+        return qpms_text, as_text, grading_text, grading_pdf_path, imprinted_pdf_path
+
+    except Exception as e:
+        print("❌ Pipeline error:", e)
+        return f"❌ Error: {e}", None, None, None, None
 
 # ---------------- GRADIO UI ---------------- 
 with gr.Blocks(title="LeadIB AI Grading (Final Flow — Verbose)") as demo: