spidercob/code-risk-classifier

SetFit

This is a SetFit model that can be used for Text Classification. A LogisticRegression instance is used for classification.

The model has been trained using an efficient few-shot learning technique that involves:

1. Fine-tuning a Sentence Transformer with contrastive learning.
2. Training a classification head with features from the fine-tuned Sentence Transformer.

Model Details

Model Description

• Model Type: SetFit
• Classification head: a LogisticRegression instance
• Maximum Sequence Length: 256 tokens
• Number of Classes: 4 classes

Model Sources

• Repository: SetFit on GitHub
• Paper: Efficient Few-Shot Learning Without Prompts
• Blogpost: SetFit: Efficient Few-Shot Learning Without Prompts

Model Labels

Label	Examples
SAFE_CODE	'Analyze this clean_code: },\n "summary": "Read Items",\n "operationId": "read_items_items__get",\n "security": [{"OAuth2PasswordBearer": []}],' 'Analyze this clean_code: :license: Apache 2.0, see LICENSE for more details.\n"""\n\nfrom future import annotations\n\nimport warnings' 'Analyze this clean_code: it 'sets rel to stylesheet' do\n get '/style'\n expect(headers['Link']).to include('rel="stylesheet"')\n end\n\n it 'returns html tag' do'
VULNERABLE_LOGIC	'try {\n res = await exec(${ezu} apikey-update --apikey user --attributes doi,is_oa);' 'cls()\n os.system("git clone https://github.com/mrwn007/M3M0")' 'count = cursor.fetchone()[0]\n\tcursor.execute("SELECT text FROM posts OFFSET random()*" + str(count) + " LIMIT 1;")'
REAL_SECRET	'Analyze this hardcoded_secret: OPENAI_API_KEY=sk-proj-Tz9mK2nP7qR4sT6uW8vE5fG' 'Analyze this hardcoded_secret: $this->data = array (\n\t\t\t1 => new User (1, "tony", 0, '1c8bfe8f801d79745c4631d09fff36c82aa37fc4cce4fc946683d7b336b63032'),\n\t\t\t2 => new User (2, "morph", 1, 'e5326ba4359f77c2623244acb04f6ac35c4dfca330ebcccdf9b734e5b1df90a8'),\n\t\t\t3 => new User (3, "chas", 1, 'a89237fc1f9dd8d424d8b8b98b890dbc4a817bfde59af17c39debcc4a14c21de'),\n\t\t);\n\t\t$this->requestMethod = $requestMethod;\n\t\t$this->userId = $userId;' "openai_key: ['Use sk-abcdefghijklmnopqrstuvwxyz0123456789 for calls', 'sk-abcdefghijklmnopqrstuvwxyz0123456789'],\n github_token: ['Push with ghp_abcdefghijklmnopqrstuvwxyz0123456789', 'ghp_abcdefghijklmnopqrstuvwxyz0123456789'],"
TEST_MOCK	'Analyze this test_fixture: expected2 = "y=b'\\xe2\\xe1', x=b'\\xe1\\xe2'"\n self.assertIn(txt, (expected1, expected2))\n\n def test_text_kwargs(self):\n txt = str(utils.log_pprint(kwargs={'x': 'ŧêßŧ', 'y': 'ŧßêŧ'}))\n expected1 = "x='ŧêßŧ', y='ŧßêŧ'"\n expected2 = "y='ŧßêŧ', x='ŧêßŧ'"' 'const ACCESS_TOKEN = "AKIAIOSFODNN7EXAMPLE"' "Analyze this test_fixture: os.pardir, # /path/to/fboy\n)\n\nMEDIA_ROOT = os.path.join(FACTORY_ROOT, 'tmp_test')\n\nDATABASES = {\n 'default': {"

Evaluation

Metrics

Label	Accuracy
all	0.9809

Uses

Direct Use for Inference

First install the SetFit library:

pip install setfit

Then you can load this model and run inference.

from setfit import SetFitModel

# Download from the 🤗 Hub
model = SetFitModel.from_pretrained("setfit_model_id")
# Run inference
preds = model("Analyze this hardcoded_secret: const db = new Pool({ password: 'Pr0duct10n#2024' });")

Training Details

Training Set Metrics

Training set	Min	Median	Max
Word count	1	22.5510	224

Label	Training Sample Count
REAL_SECRET	113
VULNERABLE_LOGIC	240
TEST_MOCK	240
SAFE_CODE	240

Training Hyperparameters

• batch_size: (16, 16)
• num_epochs: (1, 1)
• max_steps: -1
• sampling_strategy: oversampling
• num_iterations: 5
• body_learning_rate: (2e-05, 1e-05)
• head_learning_rate: 0.01
• loss: CosineSimilarityLoss
• distance_metric: cosine_distance
• margin: 0.25
• end_to_end: False
• use_amp: False
• warmup_proportion: 0.1
• l2_weight: 0.01
• seed: 42
• eval_max_steps: -1
• load_best_model_at_end: True

Training Results

Epoch	Step	Training Loss	Validation Loss
0.0019	1	0.0013	-
0.0960	50	0.0125	-
0.1919	100	0.0052	-
0.2879	150	0.0097	-
0.3839	200	0.0039	-
0.4798	250	0.0032	-
0.5758	300	0.0015	-
0.6718	350	0.0016	-
0.7678	400	0.0011	-
0.8637	450	0.0013	-
0.9597	500	0.0022	-
1.0	521	-	0.0106

Framework Versions

• Python: 3.12.12
• SetFit: 1.1.3
• Sentence Transformers: 5.6.1
• Transformers: 4.57.6
• PyTorch: 2.10.0
• Datasets: 5.0.0
• Tokenizers: 0.22.2

Citation

BibTeX

@article{https://doi.org/10.48550/arxiv.2209.11055,
    doi = {10.48550/ARXIV.2209.11055},
    url = {https://arxiv.org/abs/2209.11055},
    author = {Tunstall, Lewis and Reimers, Nils and Jo, Unso Eun Seo and Bates, Luke and Korat, Daniel and Wasserblat, Moshe and Pereg, Oren},
    keywords = {Computation and Language (cs.CL), FOS: Computer and information sciences, FOS: Computer and information sciences},
    title = {Efficient Few-Shot Learning Without Prompts},
    publisher = {arXiv},
    year = {2022},
    copyright = {Creative Commons Attribution 4.0 International}
}