Update README.md

README.md

@@ -33,32 +33,29 @@ It is a classification model and is aimed to assist in classifying vulnerabiliti
 ## How to get started with the model
 
 ```python
->>> from transformers import AutoModelForSequenceClassification, AutoTokenizer
-... import torch
-...
-... labels = ["low", "medium", "high", "critical"]
-...
-... model_name = "CIRCL/vulnerability-severity-classification-distilbert-base-uncased"
-... tokenizer = AutoTokenizer.from_pretrained(model_name)
-... model = AutoModelForSequenceClassification.from_pretrained(model_name)
-... model.eval()
-...
-... test_description = "SAP NetWeaver Visual Composer Metadata Uploader is not protected with a proper authorization, allowing unauthenticated agent to upload potentially malicious executable binaries \
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+
+labels = ["low", "medium", "high", "critical"]
+
+model_name = "CIRCL/vulnerability-severity-classification-distilbert-base-uncased"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
+model.eval()
+
+test_description = "SAP NetWeaver Visual Composer Metadata Uploader is not protected with a proper authorization, allowing unauthenticated agent to upload potentially malicious executable binaries \
 that could severely harm the host system. This could significantly affect the confidentiality, integrity, and availability of the targeted system."
-... inputs = tokenizer(test_description, return_tensors="pt", truncation=True, padding=True)
-...
-... # Run inference
-... with torch.no_grad():
-...     outputs = model(**inputs)
-...     predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-...
-... # Print results
-... print("Predictions:", predictions)
-... predicted_class = torch.argmax(predictions, dim=-1).item()
-... print("Predicted severity:", labels[predicted_class])
-...
-Predictions: tensor([[4.9335e-04, 3.4782e-02, 2.6257e-01, 7.0215e-01]])
-Predicted severity: critical
+inputs = tokenizer(test_description, return_tensors="pt", truncation=True, padding=True)
+
+# Run inference
+with torch.no_grad():
+    outputs = model(**inputs)
+    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+
+# Print results
+print("Predictions:", predictions)
+predicted_class = torch.argmax(predictions, dim=-1).item()
+print("Predicted severity:", labels[predicted_class])
 ```
 
 ### Training hyperparameters