elselse commited on
Commit
bb7eccc
·
verified ·
1 Parent(s): 5aa904f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -7
app.py CHANGED
@@ -5,16 +5,11 @@ MODEL_NAME = "CIRCL/cwe-vulnerability-classification-codebert-base"
5
 
6
  classifier = pipeline("text-classification", model=MODEL_NAME, return_all_scores=True)
7
 
8
- def classify_cwe(text, min_confidence=0.05): # 5% threshold
9
  results = classifier(text)[0]
10
  # Sort by confidence score descending
11
  sorted_results = sorted(results, key=lambda x: x["score"], reverse=True)
12
- filtered = [res for res in sorted_results if res["score"] >= min_confidence]
13
-
14
- if not filtered:
15
- return {"⚠️ No confident prediction": "Try a more specific description."}
16
-
17
- return {res["label"]: round(res["score"], 4) for res in filtered[:5]}
18
 
19
  interface = gr.Interface(
20
  fn=classify_cwe,
@@ -26,3 +21,4 @@ interface = gr.Interface(
26
 
27
  # Launch the Gradio app
28
  interface.launch()
 
 
5
 
6
  classifier = pipeline("text-classification", model=MODEL_NAME, return_all_scores=True)
7
 
8
+ def classify_cwe(text):
9
  results = classifier(text)[0]
10
  # Sort by confidence score descending
11
  sorted_results = sorted(results, key=lambda x: x["score"], reverse=True)
12
+ return {res["label"]: round(res["score"], 4) for res in sorted_results[:5]}
 
 
 
 
 
13
 
14
  interface = gr.Interface(
15
  fn=classify_cwe,
 
21
 
22
  # Launch the Gradio app
23
  interface.launch()
24
+