Update app.py

app.py

@@ -4,7 +4,6 @@ import google.generativeai as genai
 from reportlab.platypus import SimpleDocTemplate, Paragraph
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.lib.pagesizes import A4
-from google.generativeai.types import HarmCategory, HarmBlockThreshold, SafetySetting
 
 # -------------------- CONFIG --------------------
 genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
@@ -76,6 +75,14 @@ def safe_generate(model, inputs, fallback_prompt=None):
     except Exception as e:
         return None, f"❌ Exception: {e}"
 
+# ---------- COMMON SAFETY SETTINGS ----------
+safety_settings = [
+    {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
+    {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
+    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
+    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
+]
+
 # ---------- STEP 1: TRANSCRIPTION ----------
 def transcribe(ans_file):
     try:
@@ -83,12 +90,7 @@ def transcribe(ans_file):
         model = genai.GenerativeModel(
             "gemini-2.5-pro",
             generation_config={"temperature": 0},
-            safety_settings=[
-                SafetySetting(category=HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=HarmBlockThreshold.BLOCK_NONE),
-                SafetySetting(category=HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=HarmBlockThreshold.BLOCK_NONE),
-                SafetySetting(category=HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=HarmBlockThreshold.BLOCK_NONE),
-                SafetySetting(category=HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=HarmBlockThreshold.BLOCK_NONE),
-            ]
+            safety_settings=safety_settings
         )
 
         transcription, error = safe_generate(model, [TRANSCRIPTION_PROMPT, ans_uploaded], fallback_prompt="Convert the PDF into structured plain text with questions separated.")
@@ -108,12 +110,7 @@ def grade(qp_file, ms_file, transcription):
         model = genai.GenerativeModel(
             "gemini-2.5-pro",
             generation_config={"temperature": 0},
-            safety_settings=[
-                SafetySetting(category=HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT, threshold=HarmBlockThreshold.BLOCK_NONE),
-                SafetySetting(category=HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT, threshold=HarmBlockThreshold.BLOCK_NONE),
-                SafetySetting(category=HarmCategory.HARM_CATEGORY_HATE_SPEECH, threshold=HarmBlockThreshold.BLOCK_NONE),
-                SafetySetting(category=HarmCategory.HARM_CATEGORY_HARASSMENT, threshold=HarmBlockThreshold.BLOCK_NONE),
-            ]
+            safety_settings=safety_settings
         )
 
         grading, error = safe_generate(model, [GRADING_PROMPT, qp_uploaded, ms_uploaded, transcription], fallback_prompt="Grade the answers according to the marking scheme. Show marks step by step.")
@@ -161,4 +158,4 @@ with gr.Blocks(title="LeadIB AI Grading") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()