Update app.py

app.py

@@ -28,7 +28,8 @@ INPUT: This file is a PDF that first contains the Question Paper and immediately
 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.   
+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,73 @@ 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, qpms_text=None):
+    """
+    Construct the AS transcription prompt injecting the expected IDs block and graph detection instructions.
+    If qpms_text is provided, instruct the LLM to refer to it for ambiguous handwriting.
+    """
+    if not expected_ids:
+        ids_block = "{NA}"
+    else:
+        ids_block = "{\n" + "\n".join(expected_ids) + "\n}"
+    refer_text = ""
+    if qpms_text:
+        refer_text = (
+            "\nYou are also provided with the full transcript of the Question Paper and Markscheme (QP+MS). "
+            "If you encounter ambiguous handwriting (for example, if a number could be '-1.6' or '1.6'), refer to the QP+MS transcript to infer the student's intended answer. "
+            "However, if you are confident in your transcription, you may use your own judgment. "
+            "Always prioritize accuracy and context from the QP+MS transcript when in doubt.\n"
+        )
+    prompt = f"""You are a high-quality handwritten transcription assistant.
+INPUT: This PDF contains a student's handwritten answer sheet.{refer_text}
+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 ===="""
+    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): 
     """ 
@@ -283,17 +326,19 @@ def ask_gemini_for_mapping_for_page(model, image_path, grading_json, expected_id
     if expected_ids: 
         ids_block = "{\n" + "\n".join(expected_ids) + "\n}" 
     prompt = f"""
-You are an exam marker. Your role is to identify where each question begins on the page.
+You are an exam marker. Your task is to locate a blank cell adjacent to the answer step and place the marks there:
+Primary preference: Use the blank cell immediately to the right of the answer step.
+Fallback: If no blank cell is available on the right, use the blank cell immediately to the left..
 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):
+IMPORTANT: For your help i have provided u questions that u can expect in the image  :
 {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).
 - Do not place marks inside another question's answer area.
-- Prefer placing the marks in a BLANK cell immediately to the RIGHT of the answer step. If no blank cell is available to the right, then place in a blank cell to the LEFT.
+- Each question should have unique cell number 
+- If a question serial number is visible in the answer image, you must mandatorily identify the corresponding question using the grading JSON.
 - Never place marks above or below the answer.
-- If a question starts on a previous page, you may omit it for this page.
+- Only if there is no serial number u may omit to select cell number for mark placement 
 Return JSON only, like:
 [{{"question": "1.a", "cell_number": 15}}, ...]
 Grading JSON:
@@ -301,7 +346,8 @@ Grading JSON:
 """ 
     print(f"📡 Sending mapping request for image {image_path} to Gemini...") 
     img = Image.open(image_path) 
-    response = model.generate_content([prompt, img]) 
+    response = model.generate_content([prompt, img])
+    print("💬 Gemini response:", response)
     raw_text = getattr(response, "text", None) 
     if not raw_text and getattr(response, "candidates", None): 
         raw_text = response.candidates[0].content.parts[0].text 
@@ -446,93 +492,141 @@ 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.
+
+# ---------------- GRAPH PAGE EXTRACTION HELPER ----------------
+def extract_pdf_pages_as_images(pdf_path, page_numbers, prefix):
+    """
+    Extracts unique pages (1-based) from a PDF as images, saves as PNG, returns list of file paths.
+    Prints to console when extracting each page.
+    """
+    unique_pages = sorted(set(page_numbers))
+    images = convert_from_path(pdf_path, dpi=200, first_page=min(unique_pages), last_page=max(unique_pages))
+    out_paths = []
+    for idx, page_num in enumerate(unique_pages):
+        # pdf2image returns images in order, but if not contiguous, we need to map
+        # So, get the image for this page (1-based)
+        img_idx = page_num - min(unique_pages)
+        img = images[img_idx]
+        out_path = f"{prefix}_page_{page_num}.png"
+        img.save(out_path, "PNG")
+        print(f"📤 Extracted graph page {page_num} from {pdf_path} as {out_path}")
+        out_paths.append(out_path)
+    return out_paths
+
+# ---------------- 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)
+        ms_graph_pages = list(ms_graph_mapping.values())
+        ms_graph_images = []
+        if ms_graph_pages:
+            ms_graph_images = extract_pdf_pages_as_images(merged_qpms_path, ms_graph_pages, prefix="qpms_graph")
+
+        # 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, qpms_text) + "\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)
+        as_graph_pages = list(as_graph_mapping.values())
+        as_graph_images = []
+        if as_graph_pages:
+            as_graph_images = extract_pdf_pages_as_images(ans_path, as_graph_pages, prefix="as_graph")
+
+        # Step 3: (No graph bundle matching, just collect images)
+
+        # Step 4: Grading - send both transcripts to grading model, inject graph image 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 image note
+        if ms_graph_images or as_graph_images:
+            graph_note = "\n\n---\nSome questions require graphs. I’ve attached the relevant graph pages from QP+MS and from the Answer Sheet. Use them as visual context when grading.\n---\n"
+            grading_input += graph_note
+        grading_prompt_system = PROMPTS["GRADING_PROMPT"]["content"]
+        # Pass images as additional input to gemini_generate_content
+        grading_images = ms_graph_images + as_graph_images
+        grading_text = gemini_generate_content(model, grading_prompt_system + "\n\nPlease grade the following transcripts:\n" + grading_input, image_obj=grading_images if grading_images else None)
+        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: 
@@ -572,4 +666,4 @@ with gr.Blocks(title="LeadIB AI Grading (Final Flow — Verbose)") as demo:
     ) 
 
 if __name__ == "__main__": 
-    demo.launch() 
+    demo.launch()