elselse commited on
Commit
ddeca7e
·
unverified ·
1 Parent(s): b783ff6

Add application file

Browse files
Files changed (1) hide show
  1. app.py +26 -0
app.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline
3
+
4
+ # Load your model from Hugging Face
5
+ model_name = "CIRCL/vulnerability-severity-classification-chinese-macbert-base"
6
+ classifier = pipeline("text-classification", model=model_name, return_all_scores=True)
7
+
8
+ # Define severity labels
9
+ severity_labels = ["Low 低", "Medium 中", "High 高"]
10
+
11
+ def classify_text(text):
12
+ results = classifier(text)[0] # Extract the first (and only) result
13
+ scores = {severity_labels[i]: round(r["score"], 4) for i, r in enumerate(results)}
14
+ return scores
15
+
16
+ # Create the Gradio interface
17
+ interface = gr.Interface(
18
+ fn=classify_text,
19
+ inputs=gr.Textbox(lines=5, placeholder="Enter vulnerability description..."),
20
+ outputs=gr.Label(num_top_classes=4), # Display all scores
21
+ title="Vulnerability Severity Classification",
22
+ description="Enter a vulnerability description, and the model will classify its severity level."
23
+ )
24
+
25
+ # Launch the app
26
+ interface.launch()