Update app.py

i hope this works

app.py

@@ -12,41 +12,61 @@ from pdf2image import convert_from_path
 from PIL import Image, ImageDraw, ImageFont
 import cv2
 import numpy as np
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from PyPDF2 import PdfReader, PdfWriter
 
-# ---------- PROMPTS (preserved exactly) ----------
+# ---------- CONFIG ----------
+genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
+GRID_ROWS, GRID_COLS = 20, 14
+
+# ---------- PROMPTS (updated) ----------
 PROMPTS = {
-    "ALIGNMENT_PROMPT": {
+    "QP_MS_TRANSCRIBE": {
         "role": "system",
-        "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...",
-      "total_marks": 7,
-      "ms": {
-        "marks": [
-          { "id": "M1_1", "desc": "Recognise GP (r=0.9)" }
-        ],
-        "total_verification": "M1 + A1 = 2"
-      },
-      "as": {
-        "steps": ["..."],
-        "notes": "..."
-      }
-    }
-  ]
-}"""
+        "content": """You are a high-quality OCR/Transcription assistant.
+
+INPUT: This file is a PDF that **first contains the Question Paper** and immediately after it **the Markscheme**.
+TASK: Produce an exact transcription in plain text with clear separators. For every question in the Question Paper extract and output:
+- Question ID (exact as printed, e.g., "1", "2(a)", "3.b", "4(ii)")
+- Question text (exact wording; do not change punctuation)
+- Total marks for the question (if printed; otherwise try to infer/leave blank)
+
+FOR THE MARKSCHEME: Transcribe the markscheme **verbatim** exactly as it appears. Do NOT alter mark IDs, abbreviations, indentation, or descriptions. The markscheme transcription must be faithful — errors in transcription should be kept as-is rather than "corrected".
+
+OUTPUT FORMAT:
+- Plain text with clearly delimited blocks. Use a pattern like:
+----
+QUESTION BEGIN
+ID: <id>
+QTEXT:
+<question text (multiline)>
+TOTAL_MARKS: <integer or empty>
+MARKSCHEME:
+<verbatim markscheme lines for this question (multiline)>
+QUESTION END
+----
+Repeat for every question in order. If some part is not available, leave the field empty but keep the block structure.
+"""
+    },
+
+    "AS_TRANSCRIBE": {
+        "role": "system",
+        "content": """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>".
+
+OUTPUT FORMAT:
+Produce plain text with clearly delimited answer blocks using the pattern:
+----
+ANSWER BEGIN
+ID: <id or INFERRED:... or EMPTY>
+ANSWER:
+<transcribed student answer text (multiline)>
+ANSWER END
+----
+Repeat for each student answer block found.
+"""
     },
 
     "GRADING_PROMPT": {
@@ -85,18 +105,14 @@ At the very end, provide a summary table:
 |-----------------|-------|--------|
 | 1               | 6/7   | C      |
 Then show total clearly:  
-`Total: 6/7`"""
+`Total: 6/7`
+
+NOTES: The assistant will receive a structured alignment JSON (questions list with qp, total_marks, ms verbatim, and as transcribed). Grade each question independently, using the markscheme provided in the `ms` field (verbatim) and the student's `as`. Provide full markdown output as described above.
+"""
     }
 }
 
-# -------------------- CONFIG --------------------
-# The Gemini API key must be set in the environment
-genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
-
-# Grid config for imprinting
-GRID_ROWS, GRID_COLS = 20, 14
-
-# ---------- HELPERS ----------
+# -------------------- HELPERS (unchanged unless needed) --------------------
 def save_as_pdf(text, filename="output.pdf"):
     pdf = MarkdownPdf()
     pdf.add_section(Section(text, toc=False))
@@ -104,10 +120,6 @@ def save_as_pdf(text, filename="output.pdf"):
     return filename
 
 def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
-    """
-    Compress PDF only if its size is larger than max_size (default 20MB).
-    Returns path to (possibly compressed) file.
-    """
     if output_path is None:
         base, ext = os.path.splitext(input_path)
         output_path = f"{base}_compressed{ext}"
@@ -150,52 +162,183 @@ def create_model():
         print("⚡ Falling back to gemini-2.5-flash model")
         return genai.GenerativeModel("gemini-2.5-flash", generation_config={"temperature": 0})
 
-# ---------- Extract marks per question (parse grading Markdown) ----------
 def extract_marks_from_grading(grading_text):
-    """
-    Parse the grading markdown produced by the GRADING_PROMPT and extract marks per question.
-    Returns dict: {"grading": [{"question": "1.a", "marks_awarded": ["M1","A1"]}, ...]}
-    """
     grading_json = {"grading": []}
-
-    # Split by question sections using "## Question" header
-    # We allow various header spacing, e.g. "## Question 1(a)" or "## Question 1(a)\n..."
+    # Split by question header
     question_blocks = re.split(r"##\s*Question\s+", grading_text)
-    for block in question_blocks[1:]:  # skip anything before first question header
-        # The first token up to newline is the question id line, e.g. "1(a)\n### Markscheme..."
+    for block in question_blocks[1:]:
         first_line = block.strip().splitlines()[0].strip()
-        # Extract the question id - keep typical formats like 1, 1(a), 2.b, 3.d(ii)
         q_id_match = re.match(r"([0-9]+(?:[a-zA-Z]|\([^\)]+\)|(?:\.[a-zA-Z0-9]+))*)", first_line)
         if not q_id_match:
-            # fallback: try to extract tokens until first space
             q_id = first_line.split()[0]
         else:
             q_id = q_id_match.group(1).strip()
-
-        # Now find all awarded marks in that block. Search the "Awarded" column entries like M1, A1, A0, R1 etc.
-        # We use a word-boundary regex to capture tokens.
         awarded = re.findall(r"\b(M\d+|A\d+|R\d+|M0|A0|R0)\b", block)
-        # Deduplicate preserving order
         seen = set()
         awarded_unique = []
         for m in awarded:
             if m not in seen:
                 awarded_unique.append(m)
                 seen.add(m)
-
         grading_json["grading"].append({
             "question": q_id,
             "marks_awarded": awarded_unique
         })
-
     return grading_json
 
-# ---------- Ask Gemini for mapping for one page image ----------
-def ask_gemini_for_mapping_for_page(model, image_path, grading_json, rows=GRID_ROWS, cols=GRID_COLS):
+# ---------- PDF merging helper ----------
+def merge_pdfs(paths, output_path):
+    writer = PdfWriter()
+    for p in paths:
+        reader = PdfReader(p)
+        for page in reader.pages:
+            writer.add_page(page)
+    with open(output_path, "wb") as f:
+        writer.write(f)
+    return output_path
+
+# ---------- Transcript parsing helpers ----------
+def parse_qp_ms_transcript(text):
+    """
+    Parse QP+MS transcript produced according to the QP_MS_TRANSCRIBE prompt blocks.
+    Expected block markers: QUESTION BEGIN ... QUESTION END with fields ID, QTEXT, TOTAL_MARKS, MARKSCHEME.
+    Return list of questions: {id, qp, total_marks, ms}
+    """
+    questions = []
+    # Try to find blocks using the explicit markers we requested
+    blocks = re.findall(r"QUESTION BEGIN(.*?)QUESTION END", text, flags=re.DOTALL | re.IGNORECASE)
+    if blocks:
+        for block in blocks:
+            id_match = re.search(r"ID:\s*(.+)", block)
+            qtext_match = re.search(r"QTEXT:\s*(.*?)\n(?:TOTAL_MARKS:|MARKSCHEME:)", block, flags=re.DOTALL)
+            tm_match = re.search(r"TOTAL_MARKS:\s*(.*)", block)
+            ms_match = re.search(r"MARKSCHEME:\s*(.*)", block, flags=re.DOTALL)
+            qid = id_match.group(1).strip() if id_match else ""
+            qtext = qtext_match.group(1).strip() if qtext_match else ""
+            total_marks = tm_match.group(1).strip() if tm_match else ""
+            # try to normalize total_marks to int if possible
+            try:
+                total_marks = int(re.search(r"\d+", total_marks).group(0)) if total_marks else None
+            except Exception:
+                total_marks = None
+            ms = ms_match.group(1).strip() if ms_match else ""
+            questions.append({
+                "id": qid,
+                "qp": qtext,
+                "total_marks": total_marks,
+                "ms": ms
+            })
+        return questions
+
+    # Fallback: If model didn't follow markers, try splitting by lines that look like question headers
+    # This is conservative: find headings like "1", "1.", "1(a)" at line starts
+    parts = re.split(r"(?m)^\s*(\d+(?:\([a-zA-Z0-9]+\)|[a-zA-Z]|\.[a-zA-Z0-9]+)?)\s*[\.\):\-]\s*", text)
+    # parts list pattern: [pretext, id1, body1, id2, body2, ...]
+    if len(parts) >= 3:
+        it = iter(parts)
+        pre = next(it)
+        while True:
+            try:
+                qid = next(it).strip()
+                body = next(it)
+            except StopIteration:
+                break
+            # try to separate question text and markscheme inside body using "Markscheme" keyword
+            ms_split = re.split(r"(?i)\bmarkscheme\b|(?i)\bmark scheme\b", body, maxsplit=1)
+            if len(ms_split) == 2:
+                qtext = ms_split[0].strip(":-\n ")
+                ms = ms_split[1].strip()
+            else:
+                # try to look for "Marks" summary then rest
+                m_search = re.search(r"(?i)\bmarks[:\s]*\d+", body)
+                if m_search:
+                    # take text before marks as qtext
+                    qtext = body[:m_search.start()].strip()
+                    ms = body[m_search.start():].strip()
+                else:
+                    # fallback: put entire body into qp and ms empty
+                    qtext = body.strip()
+                    ms = ""
+            # try to find total marks integer
+            tm = None
+            tm_found = re.search(r"(?i)(?:total\s*marks|marks|[\/]\s*\d+|out of)\s*[:\s]*?(\d+)", body)
+            if tm_found:
+                try:
+                    tm = int(tm_found.group(1))
+                except:
+                    tm = None
+            questions.append({
+                "id": qid,
+                "qp": qtext,
+                "total_marks": tm,
+                "ms": ms
+            })
+        return questions
+
+    # If nothing found, return one block with raw text as fallback
+    return [{"id": "1", "qp": text.strip(), "total_marks": None, "ms": ""}]
+
+def parse_as_transcript(text):
+    """
+    Parse AS transcript into answer blocks. Expected markers ANSWER BEGIN ... ANSWER END.
+    Return list: {id, ans}
+    """
+    answers = []
+    blocks = re.findall(r"ANSWER BEGIN(.*?)ANSWER END", text, flags=re.DOTALL | re.IGNORECASE)
+    if blocks:
+        for block in blocks:
+            id_match = re.search(r"ID:\s*(.+)", block)
+            ans_match = re.search(r"ANSWER:\s*(.*)", block, flags=re.DOTALL)
+            qid = id_match.group(1).strip() if id_match else ""
+            ans = ans_match.group(1).strip() if ans_match else block.strip()
+            answers.append({
+                "id": qid,
+                "as": ans
+            })
+        return answers
+
+    # Fallback: split by likely question labels in the student's transcription, e.g., "1.", "1)", "1a."
+    parts = re.split(r"(?m)^\s*(\d+(?:[a-zA-Z]|\([^\)]+\))?)\s*[\.\):\-]\s*", text)
+    if len(parts) >= 3:
+        it = iter(parts)
+        pre = next(it)
+        while True:
+            try:
+                qid = next(it).strip()
+                body = next(it)
+            except StopIteration:
+                break
+            answers.append({"id": qid, "as": body.strip()})
+        return answers
+
+    # If no structure at all, try to chunk by double newlines
+    chunks = [c.strip() for c in text.split("\n\n") if c.strip()]
+    for i, c in enumerate(chunks, start=1):
+        answers.append({"id": f"INFERRED:{i}", "as": c})
+    return answers
+
+# ---------- Gemini call wrapper ----------
+def gemini_generate_content(model, prompt_text, file_upload_obj=None):
     """
-    Sends a page image and grading JSON to Gemini asking for cell numbers for questions.
-    Returns a list like: [{"question": "1.a", "cell_number": 23}, ...]
+    Helper: send prompt_text and optionally a single uploaded file to model.generate_content.
+    Returns the textual response (str).
     """
+    inputs = [prompt_text]
+    if file_upload_obj:
+        inputs.append(file_upload_obj)
+    response = model.generate_content(inputs)
+    # Response handling as in original script
+    raw_text = getattr(response, "text", None)
+    if not raw_text and getattr(response, "candidates", None):
+        # new-style candidate chain
+        raw_text = response.candidates[0].content.parts[0].text
+    if not raw_text:
+        # attempt to stringify response
+        raw_text = str(response)
+    return raw_text
+
+# ---------- Imprinting and mapping helpers remain unchanged ----------
+def ask_gemini_for_mapping_for_page(model, image_path, grading_json, rows=GRID_ROWS, cols=GRID_COLS):
     prompt = f"""
 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}).
@@ -212,53 +355,35 @@ Return JSON only, like:
 Grading JSON:
 {json.dumps(grading_json, indent=2)}
 """
-    # Load image file
     img = Image.open(image_path)
-    # Send both prompt and image to Gemini
     response = model.generate_content([prompt, img])
     raw_text = getattr(response, "text", None)
     if not raw_text and getattr(response, "candidates", None):
         raw_text = response.candidates[0].content.parts[0].text
 
     print("🔎 Gemini mapping raw output (page):")
-    print(raw_text)
-
-    # Try to extract JSON from response
-    # Commonly model will return JSON; attempt to parse the first JSON array/list block
-    json_part = None
+    print(raw_text[:1000] + ("..." if len(raw_text) > 1000 else ""))
     try:
-        # naive: find first '[' and last ']' and json.loads
         start = raw_text.index('[')
         end = raw_text.rindex(']') + 1
         json_part = raw_text[start:end]
         mapping = json.loads(json_part)
         return mapping
     except Exception as e:
-        print("⚠️ Failed to parse mapping JSON directly:", e)
-        # try to find 'json\n{...}\n' patterns
         match = re.search(r'(\[.*\])', raw_text, re.DOTALL)
         if match:
             try:
                 mapping = json.loads(match.group(1))
                 return mapping
-            except Exception as e2:
-                print("⚠️ Second parse attempt failed:", e2)
-        # fallback empty list
+            except Exception:
+                pass
         return []
 
-# ---------- Imprinting Logic (uses mapping) ----------
 def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, rows=GRID_ROWS, cols=GRID_COLS):
-    """
-    Convert PDF to images, create grid-numbered images, ask Gemini for mapping per page,
-    and then annotate marks beside the mapped cells.
-    Returns path to final imprinted (and possibly compressed) PDF.
-    Prints imprint steps in console for each page/question.
-    """
     pages = convert_from_path(pdf_path, dpi=200)
     annotated_page_paths = []
     print(f"📄 Converted answer PDF to {len(pages)} page image(s) for imprinting.")
 
-    # create grid-numbered temporary images for sending to Gemini
     temp_grid_images = []
     for p_index, page in enumerate(pages):
         img = page.convert("RGB")
@@ -272,7 +397,6 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, rows=
             num_font = ImageFont.load_default()
 
         cell_num = 1
-        # We only need numbers for clarity when sending to model (but we won't draw gridlines)
         for r in range(rows):
             for c in range(cols):
                 x = int(c * cell_w + cell_w / 2)
@@ -288,52 +412,40 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, rows=
         img.save(temp_path, "PNG")
         temp_grid_images.append(temp_path)
 
-    # Now for each page, ask Gemini for mapping
     for p_index, grid_img_path in enumerate(temp_grid_images):
         print(f"\n🛰 Sending page {p_index+1} to Gemini for mapping...")
         mapping = ask_gemini_for_mapping_for_page(model, grid_img_path, grading_json, rows, cols)
         print(f"🔁 Parsed mapping for page {p_index+1}: {mapping}")
 
-        # Prepare a clean copy of the original page to annotate (no grid numbers)
         page_img = pages[p_index].convert("RGB")
         img_cv = np.array(page_img)
         img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
         h, w, _ = img_cv.shape
         cell_w_px, cell_h_px = w / cols, h / rows
 
-        # We will maintain a set of occupied cells to prefer right/left placement heuristics
         occupied = set()
 
-        # For each mapping entry, place the corresponding marks
         for item in mapping:
             qid = item.get("question")
             cell_number = item.get("cell_number")
             if qid is None or cell_number is None:
                 continue
 
-            # Find marks for this question from grading_json
             marks_list = next((g["marks_awarded"] for g in grading_json.get("grading", []) if g["question"] == qid), [])
             if not marks_list:
-                # possible the grading JSON uses slightly different formatting of q ids; try case-insensitive match
                 marks_list = next((g["marks_awarded"] for g in grading_json.get("grading", [])
                                    if g["question"].lower() == qid.lower()), [])
 
             marks_text = ",".join(marks_list) if marks_list else "?"
 
-            # Compute candidate cell coordinates
-            # Convert cell_number -> (row, col)
             row = (cell_number - 1) // cols
             col = (cell_number - 1) % cols
 
-            # Preference: place in cell to the right (col + 1), if within grid and not occupied.
             placed = False
             candidates = []
-            # Right cell
             if col + 1 < cols:
                 candidates.append((row, col + 1))
-            # same cell (fallback)
             candidates.append((row, col))
-            # left cell
             if col - 1 >= 0:
                 candidates.append((row, col - 1))
 
@@ -346,93 +458,177 @@ def imprint_marks_using_mapping(pdf_path, grading_json, output_pdf, model, rows=
                     break
 
             if chosen is None:
-                # all occupied? just pick original cell
                 chosen = (row, col)
 
-            # Convert chosen cell to pixel coordinates (approx center-right)
             r_c, c_c = chosen
-            x_c = int((c_c + 1) * cell_w_px - cell_w_px * 0.1)  # near right edge of cell
+            x_c = int((c_c + 1) * cell_w_px - cell_w_px * 0.1)
             y_c = int((r_c + 0.5) * cell_h_px)
 
-            # Print the imprint step to console
             print(f"Page {p_index+1} | Question {qid} -> mapped cell {cell_number} -> chosen cell ({r_c},{c_c})"
                   f" -> pixel coords ({x_c},{y_c}) | marks: {marks_text}")
 
-            # Draw the text on the image (scale font according to cell size)
             font_scale = max(0.6, min(1.6, cell_h_px / 60.0))
             thickness = max(1, int(font_scale * 2))
-            # Use cv2.putText (BGR)
             cv2.putText(img_cv, marks_text, (x_c, y_c), cv2.FONT_HERSHEY_SIMPLEX,
                         font_scale, (0, 0, 255), thickness, cv2.LINE_AA)
 
-        # Save annotated page
         annotated_path = f"annotated_page_{p_index+1}.png"
         cv2.imwrite(annotated_path, img_cv)
         annotated_page_paths.append(annotated_path)
         print(f"🖊 Annotated page saved: {annotated_path}")
 
-    # Merge annotated pages into a PDF
     with open(output_pdf, "wb") as f:
         f.write(img2pdf.convert(annotated_page_paths))
 
     print(f"📑 Imprinted PDF saved to: {output_pdf}")
-    # Compress output PDF only if > 20MB
     compressed = compress_pdf(output_pdf)
     if compressed != output_pdf:
         print(f"📦 Imprinted PDF compressed: {compressed}")
     return compressed
 
-# ---------- Main pipeline ----------
+# ---------- Main pipeline (rewritten) ----------
 def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
     """
-    Runs: compress (if needed) -> upload files -> alignment -> grading -> extract marks ->
-          optional imprint (per-page mapping + annotation).
-    Returns: (alignment_json_text, grading_markdown_text, grading_pdf_path, imprinted_pdf_path or None)
+    New flow:
+    1) compress as needed
+    2) merge QP + MS -> merged_qpms.pdf
+    3) upload merged_qpms and ans separately
+    4) send two parallel transcription requests:
+       - merged_qpms with QP_MS_TRANSCRIBE prompt
+       - ans with AS_TRANSCRIBE prompt
+    5) parse transcripts to get per-question qp, ms, and per-answer as
+    6) align locally by question ID
+    7) send aligned structure to grading prompt
+    8) extract marks and optionally imprint
     """
     try:
-        # Step 0: compress only if >20MB
         qp_path = compress_pdf(qp_path)
         ms_path = compress_pdf(ms_path)
         ans_path = compress_pdf(ans_path)
 
+        # Merge QP + MS into single PDF
+        merged_qpms_path = os.path.splitext(qp_path)[0] + "_merged_qp_ms.pdf"
+        merge_pdfs([qp_path, ms_path], merged_qpms_path)
+        print(f"📎 Merged QP + MS -> {merged_qpms_path}")
+
         # Upload files to Gemini
         print("🔼 Uploading files to Gemini...")
-        qp_uploaded = genai.upload_file(path=qp_path, display_name="Question Paper")
-        ms_uploaded = genai.upload_file(path=ms_path, display_name="Markscheme")
+        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")
 
         model = create_model()
 
-        # Step 1: ALIGN (JSON only)
-        print("1️⃣ Sending ALIGNMENT_PROMPT to Gemini (alignment step)...")
-        resp = model.generate_content([
-            PROMPTS["ALIGNMENT_PROMPT"]["content"],
-            qp_uploaded,
-            ms_uploaded,
-            ans_uploaded
-        ])
-
-        json_output = getattr(resp, "text", None)
-        if not json_output and getattr(resp, "candidates", None):
-            json_output = resp.candidates[0].content.parts[0].text
-
-        # Ensure we have text
-        if not json_output:
-            raise RuntimeError("No alignment JSON returned from Gemini.")
-
-        print("✅ Alignment JSON received (truncated preview):")
-        print((json_output[:1000] + '...') if len(json_output) > 1000 else json_output)
-
-        # Step 2: GRADING (Markdown)
-        print("2️⃣ Sending GRADING_PROMPT to Gemini (grading step)...")
-        response = model.generate_content([
-            PROMPTS["GRADING_PROMPT"]["content"],
-            json_output
-        ])
+        # Prepare prompts
+        qpms_prompt = PROMPTS["QP_MS_TRANSCRIBE"]["content"]
+        as_prompt = PROMPTS["AS_TRANSCRIBE"]["content"]
+
+        # Send both requests in parallel
+        print("📡 Sending transcription requests (QP+MS & AS) in parallel...")
+        transcripts = {}
+        with ThreadPoolExecutor(max_workers=2) as ex:
+            futures = {
+                ex.submit(gemini_generate_content, model, qpms_prompt, merged_uploaded): "qpms",
+                ex.submit(gemini_generate_content, model, as_prompt, ans_uploaded): "as"
+            }
+            for fut in as_completed(futures):
+                key = futures[fut]
+                try:
+                    res_text = fut.result()
+                except Exception as e:
+                    res_text = f"❌ Error during transcription: {e}"
+                transcripts[key] = res_text
+                print(f"✅ Transcription complete for: {key} (chars: {len(res_text)})")
+
+        qpms_text = transcripts.get("qpms", "")
+        as_text = transcripts.get("as", "")
+
+        # Debug: save transcripts for review
+        with open("debug_qpms_transcript.txt", "w", encoding="utf-8") as f:
+            f.write(qpms_text)
+        with open("debug_as_transcript.txt", "w", encoding="utf-8") as f:
+            f.write(as_text)
+
+        # Parse transcripts
+        print("🔧 Parsing QP+MS transcript...")
+        qpms_questions = parse_qp_ms_transcript(qpms_text)
+        print(f"Found {len(qpms_questions)} questions in QP+MS transcript.")
+
+        print("🔧 Parsing Answer Sheet transcript...")
+        as_answers = parse_as_transcript(as_text)
+        print(f"Found {len(as_answers)} answer blocks in AS transcript.")
+
+        # Build alignment: map by normalized IDs
+        def normalize_id(qid):
+            if not qid:
+                return ""
+            s = qid.strip().lower()
+            s = re.sub(r"[\.\)\(:\s]+", "", s)
+            return s
+
+        answers_map = {}
+        for a in as_answers:
+            nid = normalize_id(a.get("id", ""))
+            if nid == "":
+                # if empty id, try to infer using INFERRED: or use a running fallback index
+                nid = a.get("id", "")
+            # store first matching block (if multiple blocks for same id, append)
+            if nid in answers_map:
+                answers_map[nid] += "\n\n" + a.get("as", "")
+            else:
+                answers_map[nid] = a.get("as", "")
+
+        aligned_questions = []
+        for q in qpms_questions:
+            qid = q.get("id", "")
+            nid = normalize_id(qid)
+            # try direct id match
+            student_ans = answers_map.get(nid)
+            # try alternative matches (e.g., '1a' vs '1(a)')
+            if student_ans is None:
+                for k in answers_map:
+                    if k.startswith(nid) or nid.startswith(k) or (nid and nid.replace(" ", "") in k):
+                        student_ans = answers_map[k]
+                        break
+            # fallback: look for first answer that contains the question id as text (loose)
+            if student_ans is None:
+                for k, v in answers_map.items():
+                    if qid and qid.lower() in k:
+                        student_ans = v
+                        break
+
+            aligned_questions.append({
+                "id": qid,
+                "qp": q.get("qp", ""),
+                "total_marks": q.get("total_marks"),
+                "ms": q.get("ms", ""),           # verbatim markscheme block
+                "as": student_ans if student_ans is not None else ""
+            })
+
+        # If any answer blocks left unmatched, optionally append them as INFERRED entries
+        matched_ids = set([normalize_id(q["id"]) for q in aligned_questions])
+        for k, v in answers_map.items():
+            if k not in matched_ids:
+                aligned_questions.append({
+                    "id": k,
+                    "qp": "",
+                    "total_marks": None,
+                    "ms": "",
+                    "as": v
+                })
+
+        # Build alignment JSON text to send to grading model
+        alignment_payload = {"questions": aligned_questions}
+        alignment_json_text = json.dumps(alignment_payload, indent=2, ensure_ascii=False)
+        print("📦 Built alignment JSON (truncated):")
+        print(alignment_json_text[:1000] + ("..." if len(alignment_json_text) > 1000 else ""))
+
+        # Step: grading
+        print("2️⃣ Sending grading prompt to Gemini...")
+        # We send both the system grading prompt and the alignment JSON as content
+        response = model.generate_content([PROMPTS["GRADING_PROMPT"]["content"], alignment_json_text])
         grading_text = getattr(response, "text", None)
         if not grading_text and getattr(response, "candidates", None):
             grading_text = response.candidates[0].content.parts[0].text
-
         if not grading_text:
             raise RuntimeError("No grading output returned from Gemini.")
 
@@ -444,7 +640,7 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
         grading_pdf_path = save_as_pdf(grading_text, f"{base_name}_graded.pdf")
         print(f"📄 Grading PDF saved: {grading_pdf_path}")
 
-        # Step 2.5: Extract marks per question from grading text
+        # Extract marks for imprinting
         grading_json = extract_marks_from_grading(grading_text)
         print("🔧 Extracted grading JSON (per-question marks):")
         print(json.dumps(grading_json, indent=2))
@@ -452,20 +648,19 @@ def align_and_grade_pipeline(qp_path, ms_path, ans_path, imprint=False):
         imprinted_pdf_path = None
         if imprint:
             print("✍ Imprint option enabled. Starting imprinting process...")
-            # Convert answer PDF to grid pages, ask Gemini for mapping per page and annotate
             imprinted_pdf_path = f"{base_name}_imprinted.pdf"
             imprinted_pdf_path = imprint_marks_using_mapping(ans_path, grading_json, imprinted_pdf_path, model)
             print(f"✅ Imprinting finished. Imprinted PDF at: {imprinted_pdf_path}")
 
-        return json_output, grading_text, grading_pdf_path, imprinted_pdf_path
+        return alignment_json_text, grading_text, grading_pdf_path, imprinted_pdf_path
 
     except Exception as e:
         print("❌ Pipeline error:", e)
         return f"❌ Error: {e}", None, None, None
 
-# ---------- GRADIO APP ----------
-with gr.Blocks(title="LeadIB AI Grading (Alignment + Auto-Grading + Imprint)") as demo:
-    gr.Markdown("## 📘 LeadIB AI Grading\nUpload **Question Paper**, **Markscheme**, and **Student Answer Sheet**.\nSystem aligns → grades → optionally imprints marks (per-question, per-page mapping).")
+# ---------------- GRADIO UI (adapted) ----------------
+with gr.Blocks(title="LeadIB AI Grading (New Flow: Parallel Transcription + Align + Grade)") as demo:
+    gr.Markdown("## 📘 LeadIB AI Grading — Updated Flow\nUpload **Question Paper**, **Markscheme**, and **Student Answer Sheet**.\nSystem: merge QP+MS -> transcribe QP+MS and AS in parallel -> align locally -> grade -> (optional) imprint marks.")
 
     with gr.Row():
         qp_file = gr.File(label="📄 Upload Question Paper (PDF)")
@@ -476,14 +671,13 @@ with gr.Blocks(title="LeadIB AI Grading (Alignment + Auto-Grading + Imprint)") a
     run_button = gr.Button("🚀 Run Alignment + Grading")
 
     with gr.Row():
-        json_output_box = gr.Textbox(label="📑 Step 1: Alignment (JSON)", lines=20)
-        grading_output_box = gr.Textbox(label="📝 Step 2: Grading (Markdown)", lines=20)
+        json_output_box = gr.Textbox(label="📑 Step: Alignment (JSON)", lines=20)
+        grading_output_box = gr.Textbox(label="📝 Step: Grading (Markdown)", lines=20)
 
     grading_pdf_file = gr.File(label="📥 Download Grading PDF")
     imprint_pdf_file = gr.File(label="📥 Download Imprinted PDF (Optional)")
 
     def run_pipeline(qp_file_obj, ms_file_obj, ans_file_obj, imprint_flag):
-        # Gradio File objects have .name attribute when saved locally
         qp_path = qp_file_obj.name
         ms_path = ms_file_obj.name
         ans_path = ans_file_obj.name
@@ -492,7 +686,7 @@ with gr.Blocks(title="LeadIB AI Grading (Alignment + Auto-Grading + Imprint)") a
             qp_path, ms_path, ans_path, imprint=imprint_flag
         )
 
-        # For Gradio file outputs: return file paths (or None)
+        # For Gradio file outputs: return text/paths
         return alignment_text, grading_text, grading_pdf_path, imprinted_pdf_path
 
     run_button.click(