ML-Gateway

This project implements a FastAPI-based local server designed to load one or more pre-trained NLP models during startup and expose them through a clean, RESTful API for inference.

For example, it leverages the Hugging Face transformers library to load the CIRCL/vulnerability-severity-classification-RoBERTa-base, which specializes in classifying vulnerability descriptions according to their severity level. The server initializes this model once at startup, ensuring minimal latency during inference requests.

Clients interact with the server via dedicated HTTP endpoints corresponding to each loaded model. Additionally, the server automatically generates comprehensive OpenAPI documentation that details the available endpoints, their expected input formats, and sample responses—making it easy to explore and integrate the services.

The ultimate goal is to enrich vulnerability data descriptions through the application of a suite of NLP models, providing direct benefits to Vulnerability-Lookup and supporting other related projects.

Installation

git clone https://github.com/vulnerability-lookup/ML-Gateway
cd ML-Gateway/
poetry install

Running the server

For production on a 16-core machine, we recommend gunicorn with uvicorn workers and gunicorn's --preload flag:

OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 poetry run gunicorn api.main:app \
  -k uvicorn.workers.UvicornWorker \
  -w 4 --preload \
  -b 0.0.0.0:8000 \
  --graceful-timeout 2 --timeout 300 \
  --reuse-port --proxy-protocol

Why these settings on 16 cores:

• --preload imports the app (and therefore loads every model) once in the master process before forking. Workers inherit the loaded weights via copy-on-write, so the models are held in memory once instead of once per worker.
• -w 4 with OMP_NUM_THREADS=4 / MKL_NUM_THREADS=4 gives each worker 4 PyTorch intra-op threads. 4 × 4 = 16 keeps every core busy during inference while avoiding the memory overhead of 8+ worker processes.
• --reuse-port lets the kernel spread incoming connections across workers; --proxy-protocol preserves client IPs when fronted by a PROXY-protocol aware load balancer.

For development, a single uvicorn process is sufficient:

poetry run uvicorn api.main:app --host 0.0.0.0 --port 8000

Deployment is straightforward—no configuration files or databases are required. Inference with the pre-trained models runs efficiently without the need for a GPU.

We recommend running ML-Gateway on a separate server than Vulnerability-Lookup.

API Endpoint

Examples

curl -X 'POST' \
  'http://127.0.0.1:8000/classify/severity' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "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."
}'
{"severity":"Critical","confidence":0.9862,"model":"CIRCL/vulnerability-severity-classification-RoBERTa-base","model_revision":"93b8afccb438608a51c33983a0fafaedd0730f11"}

For classifying severity the default model involved is CIRCL/vulnerability-severity-classification-RoBERTa-base. But you can specify another:

curl -X 'POST' \
  'http://127.0.0.1:8000/classify/severity' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "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.",
  "model": "CIRCL/vulnerability-severity-classification-distilbert-base-uncased"
}'
{"severity":"Critical","confidence":0.7022,"model":"CIRCL/vulnerability-severity-classification-distilbert-base-uncased","model_revision":"<commit-sha>"}

If you need to use the model pre-trained for Chinese language:

curl -X 'POST'   'http://127.0.0.1:8000/classify/severity'   -H 'accept: application/json'   -H 'Content-Type: application/json'   -d '{
  "description": "TOTOLINK A3600R是中国吉翁电子（TOTOLINK）公司的一款6天线1200M无线路由器。TOTOLINK A3600R存在缓冲区溢出漏洞，该漏洞源于/cgi-bin/cstecgi.cgi文件的UploadCustomModule函数中的File参数未能正确验证输入数据的长度大小，攻击者可利用该漏洞在系统上执行任意代码或者导致拒绝服务。",
  "model": "CIRCL/vulnerability-severity-classification-chinese-macbert-base"
}'
{"severity":"高","confidence":0.9802,"model":"CIRCL/vulnerability-severity-classification-chinese-macbert-base","model_revision":"<commit-sha>"}

ATT&CK technique classification

POST /classify/attack-techniques ranks MITRE ATT&CK (Enterprise) techniques for a vulnerability description with CIRCL/vulnerability-attack-technique-classification-roberta-base. Unlike severity classification this is a multi-label task: every technique in the model's vocabulary is scored independently (sigmoid), so the scores do not sum to 1 and several techniques can clear the 0.5 prediction threshold at once. The top_k field (default 10) controls how many ranked techniques are returned; the full ranking is computed and cached once per description, so varying top_k does not re-run inference.

curl -X 'POST' \
  'http://127.0.0.1:8000/classify/attack-techniques' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "description": "Zoho ManageEngine ServiceDesk Plus before 11306 is vulnerable to unauthenticated remote code execution.",
  "top_k": 3
}'
{"techniques":[{"technique":"T1190","name":"Exploit Public-Facing Application","score":0.7147,"predicted":true},{"technique":"T1059","name":"Command and Scripting Interpreter","score":0.7063,"predicted":true},{"technique":"T1505","name":"Server Software Component","score":0.6834,"predicted":true}],"model":"CIRCL/vulnerability-attack-technique-classification-roberta-base","model_revision":"3b98a1d962b5c35c5dc82ac188f8d2a9aae63eda"}

Response fields:

Field	Type	Description
techniques	array	Top-k techniques ranked by score, best first. Empty when classification failed.
techniques[].technique	string	MITRE ATT&CK technique ID (e.g. T1190).
techniques[].name	string | null	Official ATT&CK technique name, from the bundled ATT&CK name table.
techniques[].score	float	Sigmoid probability, rounded to four decimals.
techniques[].predicted	bool	true when the score is at least 0.5, the threshold used by the model's training metrics.
model / model_revision / error		Same provenance and error semantics as /classify/severity.

Integration with Vulnerability-Lookup

The HTML frontend templates of Vulnerability-Lookup use asynchronous JavaScript calls to interact with a dedicated API endpoint:

fetch("https://vulnerability.circl.lu/api/vlai/severity-classification", {
    method: "POST",
    headers: {
    "Content-Type": "application/json"
    },
    body: JSON.stringify({ description: "Description of the vulnerability…" })
})
.then(response => response.json())
.then(result => {
    console.log(result["severity"] + " (confidence: " + result["confidence"] + ")");
})
.catch((error) => {
    console.error("Error:", error);
});

When a request is received, the Vulnerability-Lookup backend forwards it to the ML-Gateway API. The ML-Gateway performs the inference, and the resulting severity classification, along with a confidence score, is returned and displayed to the user.

Response fields

Field	Type	Description
severity	string | null	Predicted severity label (e.g. Low, Medium, High, Critical). null when classification could not be performed.
confidence	float	Softmax probability of the predicted class, rounded to four decimals.
model	string	Hugging Face model identifier that produced the prediction.
model_revision	string | null	Commit SHA of the model snapshot loaded from the Hugging Face Hub. null when the snapshot does not carry revision metadata (e.g. a local path). Useful to pin and audit the exact weights that produced a result.
error	string	Present only when classification failed; carries a human-readable message.

Refreshing the models

poetry run ml-gw-cli refresh-all

Funding

AIPITCH (AI-Powered Innovative Toolkit for Cybersecurity Hubs) is a co-funded EU project supported by the European Cybersecurity Competence Centre (ECCC) under the DIGITAL-ECCC-2024-DEPLOY-CYBER-06-ENABLINGTECH program and CIRCL.

The project brings together an international consortium to develop AI-based tools that enhance the capabilities of operational cybersecurity teams. These tools are designed to support critical services, with a focus on national security teams, while also being applicable to internal security teams in companies and institutions.

License

ML-Gateway is licensed under GNU Affero General Public License version 3.

Copyright (c) 2025-2026 Computer Incident Response Center Luxembourg (CIRCL)
Copyright (C) 2025-2026 Cédric Bonhomme - https://github.com/cedricbonhomme