db-d2/primevul-codebert-embeddings

---
license: mit
task_categories:
- text-classification
tags:
- code
- vulnerability-detection
- embeddings
- codebert
- positive-unlabeled-learning
language:
- code
size_categories:
- 100K<n<1M
---

# PrimeVul Embeddings for PU Learning

Pre-extracted [CLS] token embeddings from two code models for all functions in the PrimeVul v0.1 vulnerability detection dataset, plus the raw PrimeVul v0.1 JSONL source files.

## CodeBERT Embeddings (root .npz files)

Each .npz file contains frozen CodeBERT embeddings (768-dimensional vectors) for C/C++ functions, along with their labels and CWE type annotations. These were extracted once using a frozen CodeBERT model and are used for downstream PU (positive-unlabeled) learning experiments without requiring GPU access.

| File | Functions | Vulnerable | Shape |
|------|-----------|-----------|-------|
| train.npz | 175,797 | 4,862 (2.77%) | (175797, 768) |
| valid.npz | 23,948 | 593 (2.48%) | (23948, 768) |
| test.npz | 24,788 | 549 (2.21%) | (24788, 768) |
| test_paired.npz | 870 | 435 (50%) | (870, 768) |

Arrays in each .npz:

- embeddings: (N, 768) float32 -- CodeBERT [CLS] token vectors
- labels: (N,) int32 -- 0 = benign, 1 = vulnerable
- cwe_types: (N,) U20 string -- CWE category (e.g., "CWE-119") or "unknown"
- idxs: (N,) int64 -- original PrimeVul record index for traceability

### How to load

```python
import numpy as np

data = np.load("train.npz")
X = data["embeddings"] # (175797, 768)
y = data["labels"] # (175797,)
cwes = data["cwe_types"] # (175797,)
```

No special flags needed. All arrays use standard numpy dtypes (float32, int32, U20, int64).

## VulBERTa Embeddings (vulberta/ folder)

Same format as CodeBERT but extracted from claudios/VulBERTa-mlm, a RoBERTa model pretrained on C/C++ vulnerability code. Same functions, same labels, same idxs -- only the embedding vectors differ.

| File | Functions | Shape |
|------|-----------|-------|
| vulberta/train.npz | 175,797 | (175797, 768) |
| vulberta/valid.npz | 23,948 | (23948, 768) |
| vulberta/test.npz | 24,788 | (24788, 768) |
| vulberta/test_paired.npz | 870 | (870, 768) |

VulBERTa embeddings have higher L2 magnitude (~27 vs ~21 for CodeBERT) but the same 768 dimensions. Load the same way: np.load("vulberta/train.npz").

## Raw PrimeVul v0.1 data (raw/ folder)

The raw/ folder contains the original PrimeVul v0.1 JSONL files from the PrimeVul project. Each line is a JSON object with fields including func (source code), target (0/1 label), cwe (list of CWE strings), cve (CVE identifier), and project metadata.

| File | Records |
|------|---------|
| raw/primevul_train.jsonl | 175,797 |
| raw/primevul_valid.jsonl | 23,948 |
| raw/primevul_test.jsonl | 24,788 |
| raw/primevul_train_paired.jsonl | 9,724 |
| raw/primevul_valid_paired.jsonl | 870 |
| raw/primevul_test_paired.jsonl | 870 |

## Extraction details

### CodeBERT

- Model: microsoft/codebert-base (RoBERTa architecture, 125M parameters)
- Extraction: frozen model, [CLS] token from final layer
- Tokenization: max_length=512, truncation=True, padding=max_length
- Source data: PrimeVul v0.1 (chronological train/valid/test splits)
- Extracted on: Google Colab, A100 GPU, ~23 minutes for all splits

### VulBERTa

- Model: claudios/VulBERTa-mlm (RoBERTa architecture, 125M parameters, pretrained on C/C++ vulnerability code)
- Extraction: frozen model, [CLS] token from final layer
- Tokenization: max_length=512, truncation=True, padding=max_length
- Source data: PrimeVul v0.1 (same functions as CodeBERT)
- Extracted on: Google Colab, A100 GPU, ~23 minutes for all splits

## Citation

If you use this data, please cite the PrimeVul dataset:

```bibtex
@article{ding2024primevul,
title={Vulnerability Detection with Code Language Models: How Far Are We?},
author={Ding, Yangruibo and Fu, Yanjun and Ibrahim, Omniyyah and Sitawarin, Chawin and Chen, Xinyun and Alomair, Basel and Wagner, David and Ray, Baishakhi and Chen, Yizheng},
journal={arXiv preprint arXiv:2403.18624},
year={2024}
}
```

## License

MIT (same as PrimeVul)

---
许可证：MIT协议
任务类别：
- 文本分类
标签：
- 代码
- 漏洞检测
- 嵌入（embeddings）
- CodeBERT
- 正样本未标记学习（positive-unlabeled learning，PU学习）
语言：
- 代码
规模类别：
- 100K<n<1M
---

# 面向正样本未标记学习（positive-unlabeled learning，PU学习）的PrimeVul嵌入

本数据集包含从两款代码模型中预提取的PrimeVul v0.1漏洞检测数据集中所有函数的[CLS] Token 嵌入（embeddings），同时附带原始PrimeVul v0.1版本的JSONL源文件。

## CodeBERT 嵌入（根目录下的.npz文件）

每个.npz文件均存储了针对C/C++函数的冻结版CodeBERT嵌入（768维向量），同时附带对应标签与CWE类型注释。本数据集通过冻结的CodeBERT模型一次性提取得到上述嵌入，可直接用于下游正样本未标记学习实验，无需使用GPU加速。

| 文件名 | 函数总数 | 含漏洞函数数 | 维度形状 |
|------|-----------|-----------|-------|
| train.npz | 175,797 | 4,862（占比2.77%） | (175797, 768) |
| valid.npz | 23,948 | 593（占比2.48%） | (23948, 768) |
| test.npz | 24,788 | 549（占比2.21%） | (24788, 768) |
| test_paired.npz | 870 | 435（占比50%） | (870, 768) |

每个.npz文件中包含的数组如下：
- embeddings：(N, 768) float32 —— CodeBERT的[CLS] Token 向量
- labels：(N,) int32 —— 0表示良性代码，1表示存在漏洞的代码
- cwe_types：(N,) U20 字符串 —— CWE分类（例如"CWE-119"）或"unknown"（未知）
- idxs：(N,) int64 —— 用于溯源的原始PrimeVul数据记录索引

### 加载方法
python
import numpy as np

data = np.load("train.npz")
X = data["embeddings"] # (175797, 768)
y = data["labels"] # (175797,)
cwes = data["cwe_types"] # (175797,)

无需特殊参数即可加载，所有数组均采用标准NumPy数据类型（float32、int32、U20、int64）。

## VulBERTa 嵌入（vulberta/ 目录）

其格式与CodeBERT嵌入完全一致，嵌入向量从claudios/VulBERTa-mlm模型中提取——该模型为基于C/C++漏洞代码预训练的RoBERTa架构模型。所有函数、标签与索引均与CodeBERT嵌入一致，仅嵌入向量存在差异。

| 文件名 | 函数总数 | 维度形状 |
|------|-----------|-------|
| vulberta/train.npz | 175,797 | (175797, 768) |
| vulberta/valid.npz | 23,948 | (23948, 768) |
| vulberta/test.npz | 24,788 | (24788, 768) |
| vulberta/test_paired.npz | 870 | (870, 768) |

VulBERTa嵌入的L2范数更高（约27，而CodeBERT嵌入约为21），但同样为768维。加载方法与CodeBERT嵌入一致：使用np.load("vulberta/train.npz")即可。

## 原始PrimeVul v0.1 数据（raw/ 目录）

raw/ 目录中包含PrimeVul项目官方发布的原始PrimeVul v0.1版本JSONL文件。每行均为一个JSON对象，包含以下字段：func（源代码）、target（0/1分类标签）、cwe（CWE字符串列表）、cve（CVE标识符）以及项目元数据。

| 文件名 | 记录总数 |
|------|---------|
| raw/primevul_train.jsonl | 175,797 |
| raw/primevul_valid.jsonl | 23,948 |
| raw/primevul_test.jsonl | 24,788 |
| raw/primevul_train_paired.jsonl | 9,724 |
| raw/primevul_valid_paired.jsonl | 870 |
| raw/primevul_test_paired.jsonl | 870 |

## 提取细节

### CodeBERT
- 模型：microsoft/codebert-base（RoBERTa架构，参数量1.25亿）
- 提取方式：采用冻结模型，提取最后一层的[CLS] Token 输出
- 分词设置：最大序列长度512，启用截断与最大长度填充
- 源数据：PrimeVul v0.1数据集（按时间顺序划分的训练/验证/测试集）
- 提取环境：Google Colab平台，搭载A100 GPU，所有数据集拆分的提取耗时约23分钟

### VulBERTa
- 模型：claudios/VulBERTa-mlm（RoBERTa架构，参数量1.25亿，基于C/C++漏洞代码预训练）
- 提取方式：采用冻结模型，提取最后一层的[CLS] Token 输出
- 分词设置：最大序列长度512，启用截断与最大长度填充
- 源数据：PrimeVul v0.1数据集（与CodeBERT嵌入使用的函数集完全一致）
- 提取环境：Google Colab平台，搭载A100 GPU，所有数据集拆分的提取耗时约23分钟

## 引用声明
若您使用本数据集，请引用PrimeVul数据集的相关论文：
bibtex
@article{ding2024primevul,
title={Vulnerability Detection with Code Language Models: How Far Are We?},
author={Ding, Yangruibo and Fu, Yanjun and Ibrahim, Omniyyah and Sitawarin, Chawin and Chen, Xinyun and Alomair, Basel and Wagner, David and Ray, Baishakhi and Chen, Yizheng},
journal={arXiv preprint arXiv:2403.18624},
year={2024}
}

## 许可证
MIT协议（与PrimeVul数据集一致）