| import gradio as gr | |
| from transformers import pipeline | |
| # Load your model from Hugging Face | |
| model_name = "CIRCL/vulnerability-severity-classification-roberta-base" | |
| classifier = pipeline("text-classification", model=model_name, return_all_scores=True) | |
| # Define severity labels | |
| severity_labels = ["Low", "Medium", "High", "Critical"] | |
| def classify_text(text): | |
| results = classifier(text)[0] # Extract the first (and only) result | |
| scores = {severity_labels[i]: round(r["score"], 4) for i, r in enumerate(results)} | |
| return scores | |
| # Create the Gradio interface | |
| interface = gr.Interface( | |
| fn=classify_text, | |
| inputs=gr.Textbox(lines=5, placeholder="Enter vulnerability description..."), | |
| outputs=gr.Label(num_top_classes=4), # Display all scores | |
| title="Vulnerability Severity Classification", | |
| description="Enter a vulnerability description, and the model will classify its severity level." | |
| ) | |
| # Launch the app | |
| interface.launch() | |