Update app.py

app.py

@@ -1,16 +1,16 @@
 import os
 import gradio as gr
 import google.generativeai as genai
-from google.generativeai.types import HarmCategory, HarmBlockThreshold
 from markdown_pdf import MarkdownPdf, Section
 
 # -------------------- CONFIG --------------------
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
 
-# ---------- PROMPTS (JSON Style) ----------
+# ---------- PROMPTS (JSON style) ----------
+
 PROMPTS = {
     "TRANSCRIPTION_PROMPT": {
-        "role": "Student Transcription",
+        "role": "system",
         "content": """Your Role: You are an expert technical transcriber specializing in mathematical and scientific documents. 
 Your mission is to convert handwritten solutions from a provided image or PDF into a clean, accurate, and logically structured Markdown format.
 Instructions:
@@ -22,7 +22,7 @@ Instructions:
 """
     },
     "MARKSCHEME_TRANSCRIPTION_PROMPT": {
-        "role": "Markscheme Transcription",
+        "role": "system",
         "content": """Your Role: You are an expert transcriber. 
 Convert the official marking scheme from the provided PDF into clean, structured Markdown.
 Instructions:
@@ -35,7 +35,7 @@ Instructions:
 """
     },
     "GRADING_PROMPT": {
-        "role": "Examiner Grading",
+        "role": "system",
         "content": """You are an official examiner. Use the following grading rules strictly.
 Abbreviations:  
 - M: Marks awarded for attempting to use a correct Method.  
@@ -76,9 +76,9 @@ Produce a GitHub-flavored Markdown table with 3 columns:
 |---------------|---------------|--------|  
 Special Formatting Rule:  
 - Whenever a mark is lost (M0, A0, R0 etc.), wrap it in red using: `<span style=\"color:red\">M0</span>`.  
-- Also wrap the corresponding Reason in red using `<span style=\"color:red\">reason text</span>`.  
+- Also wrap the corresponding Reason in red color.
 - Keep awarded marks (M1, A1, etc.) in plain text.  
-- If mixed (e.g., M1A0A1), only highlight the lost marks (`A0`) and its reason.  
+- If mixed (e.g., M1A0A1), only highlight the lost marks (`A0`) and its reason.
 After the table, provide:  
 ### Summary & Final Mark  
 - Total marks obtained vs total available  
@@ -95,27 +95,17 @@ def save_as_pdf(text, filename="output.pdf"):
     pdf.save(filename)
     return filename
 
-# ---------- COMMON MODEL CALL ----------
-def call_model(inputs):
-    model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
-    response = model.generate_content(
-        inputs,
-        safety_settings={
-            HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
-            HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
-        }
-    )
-    if getattr(response, "text", None):
-        return response.text
-    elif response.candidates:
-        return response.candidates[0].content.parts[0].text
-    return None
-
 # ---------- STEP 1: TRANSCRIBE STUDENT ----------
 def transcribe_student(ans_file):
     try:
         ans_uploaded = genai.upload_file(path=ans_file, display_name="Answer Sheet")
-        transcription = call_model([PROMPTS["TRANSCRIPTION_PROMPT"]["content"], ans_uploaded])
+        model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
+
+        resp = model.generate_content([PROMPTS["TRANSCRIPTION_PROMPT"]["content"], ans_uploaded])
+        transcription = getattr(resp, "text", None)
+        if not transcription and resp.candidates:
+            transcription = resp.candidates[0].content.parts[0].text
+
         pdf_path = save_as_pdf(transcription, "student_transcription.pdf")
         return transcription, pdf_path
     except Exception as e:
@@ -125,7 +115,13 @@ def transcribe_student(ans_file):
 def transcribe_ms(ms_file):
     try:
         ms_uploaded = genai.upload_file(path=ms_file, display_name="Markscheme")
-        ms_transcription = call_model([PROMPTS["MARKSCHEME_TRANSCRIPTION_PROMPT"]["content"], ms_uploaded])
+        model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
+
+        resp = model.generate_content([PROMPTS["MARKSCHEME_TRANSCRIPTION_PROMPT"]["content"], ms_uploaded])
+        ms_transcription = getattr(resp, "text", None)
+        if not ms_transcription and resp.candidates:
+            ms_transcription = resp.candidates[0].content.parts[0].text
+
         pdf_path = save_as_pdf(ms_transcription, "ms_transcription.pdf")
         return ms_transcription, pdf_path
     except Exception as e:
@@ -135,12 +131,19 @@ def transcribe_ms(ms_file):
 def grade(qp_file, ms_transcription, student_transcription):
     try:
         qp_uploaded = genai.upload_file(path=qp_file, display_name="Question Paper")
-        grading = call_model([
+        model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
+
+        response = model.generate_content([
             PROMPTS["GRADING_PROMPT"]["content"],
             qp_uploaded,
             "### Markscheme Transcription:\n" + ms_transcription,
             "### Student Transcription:\n" + student_transcription
         ])
+
+        grading = getattr(response, "text", None)
+        if not grading and response.candidates:
+            grading = response.candidates[0].content.parts[0].text
+
         pdf_path = save_as_pdf(grading, "grading.pdf")
         return grading, pdf_path
     except Exception as e: