Update app.py

app.py

@@ -5,16 +5,11 @@ MODEL_NAME = "CIRCL/cwe-vulnerability-classification-codebert-base"
 
 classifier = pipeline("text-classification", model=MODEL_NAME, return_all_scores=True)
 
-def classify_cwe(text, min_confidence=0.05):  # 5% threshold
+def classify_cwe(text):
     results = classifier(text)[0]
     # Sort by confidence score descending
     sorted_results = sorted(results, key=lambda x: x["score"], reverse=True)
-    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]}
+    return {res["label"]: round(res["score"], 4) for res in sorted_results[:5]}
 
 interface = gr.Interface(
     fn=classify_cwe,
@@ -26,3 +21,4 @@ interface = gr.Interface(
 
 # Launch the Gradio app
 interface.launch()
+