Update app.py

app.py

@@ -5,11 +5,16 @@ MODEL_NAME = "CIRCL/cwe-vulnerability-classification-codebert-base"
 
 classifier = pipeline("text-classification", model=MODEL_NAME, return_all_scores=True)
 
-def classify_cwe(text):
+def classify_cwe(text, min_confidence=0.05):  # 5% threshold
     results = classifier(text)[0]
     # Sort by confidence score descending
     sorted_results = sorted(results, key=lambda x: x["score"], reverse=True)
-    return {res["label"]: round(res["score"], 4) for res in sorted_results[:5]}
+    filtered = [res for res in sorted_results if res["score"] >= min_confidence]
+
+    if not filtered:
+        return {"⚠️ No confident prediction": "Try a more specific description."}
+
+    return {res["label"]: round(res["score"], 4) for res in filtered[:5]}
 
 interface = gr.Interface(
     fn=classify_cwe,