cedricbonhomme commited on
Commit
cff57aa
·
unverified ·
1 Parent(s): f00dd33

Added minum files for the app.

Browse files
Files changed (2) hide show
  1. app.py +27 -0
  2. requirements.txt +2 -0
app.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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-roberta-base"
6
+ classifier = pipeline("text-classification", model=model_name, return_all_scores=True)
7
+
8
+ # Define severity labels
9
+ severity_labels = ["Low", "Medium", "High", "Critical"]
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()
27
+
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ transformers
2
+ torch