Update app.py

app.py

@@ -25,19 +25,31 @@ with open("child_to_parent_mapping.json", "r") as f:
 
 def predict_cwe(commit_message: str):
     """
-    Predict CWE(s) from a commit message and map to parent CWEs.
+    Predict CWE(s) from a commit message and map to parent CWEs if applicable.
     """
     results = classifier(commit_message)[0]
     sorted_results = sorted(results, key=lambda x: x["score"], reverse=True)
-        
+
     threshold = 0.2
     filtered_results = [item for item in sorted_results if item["score"] >= threshold]
 
-    # Map predictions to parent CWE (if available)
-    mapped_results = {}
-    for item in sorted_results[:5]:
-        mapped_results[item["label"]] = round(float(item["score"]), 4)
-    return mapped_results
+    parent_scores = {}
+
+    for item in filtered_results:
+        label = item["label"].replace("CWEEEE_", "") 
+        score = float(item["score"])
+
+        parent_label = child_to_parent.get(label, label)
+
+        if parent_label in parent_scores:
+            parent_scores[parent_label] += score
+        else:
+            parent_scores[parent_label] = score
+
+    # Sort by score descending and round
+    sorted_parent_scores = sorted(parent_scores.items(), key=lambda x: x[1], reverse=True)
+    return {f"CWE-{k}": round(v, 4) for k, v in sorted_parent_scores[:5]}
+
 
 # Gradio UI
 demo = gr.Interface(