Update app.py

app.py

@@ -8,7 +8,6 @@ from pydantic import BaseModel
 from transformers import BertTokenizer, AutoModelForSequenceClassification
 from arabert.preprocess import ArabertPreprocessor
 
-
 MODEL_REPO = "kkAsmaa/ChildShield"
 MODEL_NAME = "aubmindlab/bert-base-arabertv02-twitter"
 SUB_FOLDER = "ChildShield"
@@ -20,7 +19,6 @@ model = AutoModelForSequenceClassification.from_pretrained(MODEL_REPO, token=HF_
 model.eval()
 arabic_prep = ArabertPreprocessor(model_name=MODEL_NAME)
 
-
 app = FastAPI(title="ChildShield Backend API")
 
 class InputData(BaseModel):
@@ -104,7 +102,6 @@ def predict_safety_api(text):
 
     final_prediction = "UNSAFE" if is_blocked else "SAFE"
 
-
     print("\n📊 ===== CHILDSHIELD REPORT =====")
     print(f"📥 Original Text:\n{text[:100]}")
     print(f"\n🧹 Cleaned Text:\n{cleaned_text[:100]}")
@@ -121,7 +118,23 @@ def predict_safety_api(text):
     except Exception as e:
         print(f"⚠️ [Logging Warning] Could not write to log file: {e}")
 
-    
+    unsafe_confidence_score = round(highest_unsafe_prob, 4)
+    safe_confidence_score = round(1.0 - highest_unsafe_prob, 4)
+    if is_blocked:
+        return {
+            "original_text": text,
+            "cleaned_text": cleaned_text,
+            "total_tokens": total_tokens_count,
+            "window_size": window_size,
+            "overlap": overlap,
+            "total_windows": total_windows_count,
+            "triggered_windows": triggered_windows,
+            "windows_analysis": windows_analysis,
+            "final_prediction": final_prediction,
+            "blocked": is_blocked,
+            "confidence": unsafe_confidence_score 
+        }
+
     return {
         "original_text": text,
         "cleaned_text": cleaned_text,
@@ -133,16 +146,14 @@ def predict_safety_api(text):
         "windows_analysis": windows_analysis,
         "final_prediction": final_prediction,
         "blocked": is_blocked,
-        "highest_unsafe_confidence": round(highest_unsafe_prob, 4)
+        "confidence": safe_confidence_score 
     }
 
-
 @app.post("/predict")
 def predict(data: InputData):
     result = predict_safety_api(data.text)
     return result
 
-
 gradio_interface = gr.Interface(    
     fn=predict_safety_api,
     inputs=gr.Textbox(lines=4, placeholder="Enter Arabic text to analyze..."),
@@ -150,9 +161,7 @@ gradio_interface = gr.Interface(
     title="ChildShield Production API Gate (Arabic Version)🛡️"
 )
 
-
 app = gr.mount_gradio_app(app, gradio_interface, path="/")
 
-
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)