Sumeetgpt/code-risk-classifier

Code Risk Classifier

Fine-tuned SetFit model for code security risk classification. Part of the Spidercob DLP platform — reduces false positives in supply-chain and secret scanning.

Model Description

Classifies code snippets into one of four risk categories:

Label	Description
REAL_SECRET	Actual hardcoded credentials, API keys, or tokens
VULNERABLE_LOGIC	Code patterns associated with known CVEs or insecure practices
TEST_MOCK	Test fixtures, mock data, example credentials in test files
SAFE_CODE	Benign code with no security concerns

Base model: sentence-transformers/all-MiniLM-L6-v2
Framework: SetFit (few-shot fine-tuning)
Test accuracy: 98.09%

Usage

from setfit import SetFitModel

model = SetFitModel.from_pretrained("Sumeetgpt/code-risk-classifier")
predictions = model.predict([
    'API_KEY = "sk-prod-abc123real"',
    'API_KEY = "test_placeholder_key"',
])
print(predictions)  # ['REAL_SECRET', 'TEST_MOCK']

Intended Use

Used inside the Spidercob supply-chain scanner to filter CVE findings and secret detections. Prevents alert fatigue from test fixtures and example configs while surfacing genuine security risks.

Training Data

Trained on examples from:

• WebGoat (vulnerable Java web application)
• DVWA (Damn Vulnerable Web Application)
• truffleHog secret detection corpus
• factory_boy, faker, pytest fixtures
• Django, FastAPI secure coding examples

Limitations

• Optimized for short code snippets (single lines to a few lines of context)
• Works best with Python, JavaScript, Java — limited training on other languages
• Not a standalone SAST tool — designed as a downstream classifier for pattern-matched findings