Update app.py

app.py

@@ -36,7 +36,10 @@ def predict_cwe(commit_message: str):
     # Map predictions to parent CWE (if available)
     mapped_results = {}
     for item in sorted_results[:5]:
-        mapped_results[item["label"]] = round(float(item["score"]), 4)
+        child_cwe = item["label"].replace("CWE-", "")
+        parent_cwe = child_to_parent.get(child_cwe, child_cwe)  # default to child if no parent
+        mapped_results[f"CWE-{parent_cwe}"] = round(float(item["score"]), 4)
+
     return mapped_results
 
 # Gradio UI