Update README.md

README.md

@@ -12,4 +12,264 @@ metrics:
 - accuracy
 pipeline_tag: text-classification
 library_name: transformers
----
+---
+
+# UnixCoder-Primevul-BigVul Model Card
+
+## Model Overview
+
+`UnixCoder-Primevul-BigVul` is a multi-task model based on Microsoft's `unixcoder-base`, fine-tuned to detect vulnerabilities (`vul`) and classify Common Weakness Enumeration (CWE) types in code snippets. It was developed by [mahdin70](https://huggingface.co/mahdin70) and trained on a balanced dataset combining BigVul and PrimeVul datasets. The model performs binary classification for vulnerability detection and multi-class classification for CWE identification.
+
+- **Model Repository**: [mahdin70/UnixCoder-Primevul-BigVul](https://huggingface.co/mahdin70/UnixCoder-Primevul-BigVul)
+- **Base Model**: [microsoft/unixcoder-base](https://huggingface.co/microsoft/unixcoder-base)
+- **Tasks**: Vulnerability Detection (Binary), CWE Classification (Multi-class)
+- **License**: MIT (assumed; adjust if different)
+- **Date**: Trained and uploaded as of March 11, 2025
+
+## Model Architecture
+
+The model extends `unixcoder-base` with two task-specific heads:
+- **Vulnerability Head**: A linear layer mapping 768-dimensional hidden states to 2 classes (vulnerable or not).
+- **CWE Head**: A linear layer mapping 768-dimensional hidden states to 135 classes (134 CWE types + 1 for "no CWE").
+
+The architecture is implemented as a custom `MultiTaskUnixCoder` class in PyTorch, with the loss computed as the sum of cross-entropy losses for both tasks.
+
+## Training Dataset
+
+The model was trained on the `mahdin70/balanced_merged_bigvul_primevul` dataset (configuration: `10_per_commit`), which combines:
+- **BigVul**: A dataset of real-world vulnerabilities from open-source projects.
+- **PrimeVul**: A dataset focused on prime vulnerabilities in code.
+
+### Dataset Details
+- **Splits**:
+  - Train: 124,780 samples
+  - Validation: 26,740 samples
+  - Test: 26,738 samples
+- **Features**:
+  - `func`: Code snippet (text)
+  - `vul`: Binary label (0 = non-vulnerable, 1 = vulnerable)
+  - `CWE ID`: CWE identifier (e.g., CWE-89) or None for non-vulnerable samples
+- **Preprocessing**:
+  - CWE labels were encoded using a `LabelEncoder` with 134 unique CWE classes identified across the dataset.
+  - Non-vulnerable samples assigned a CWE label of -1 (mapped to 0 in the model).
+
+The dataset is balanced to ensure a fair representation of vulnerable and non-vulnerable samples, with a maximum of 10 samples per commit where applicable.
+
+## Training Details
+
+### Training Arguments
+The model was trained using the Hugging Face `Trainer` API with the following arguments:
+- **Output Directory**: `./unixcoder_multitask`
+- **Evaluation Strategy**: Per epoch
+- **Save Strategy**: Per epoch
+- **Learning Rate**: 2e-5
+- **Batch Size**: 8 (per device, train and eval)
+- **Epochs**: 3
+- **Weight Decay**: 0.01
+- **Logging**: Every 10 steps, logged to `./logs`
+- **WANDB**: Disabled
+
+### Training Environment
+- **Hardware**: NVIDIA Tesla T4 GPU
+- **Framework**: PyTorch 2.5.1+cu121, Transformers 4.47.0
+- **Duration**: ~6 hours, 34 minutes, 53 seconds (23,397 steps)
+
+### Training Metrics
+Validation metrics across epochs:
+
+| Epoch | Training Loss | Validation Loss | Vul Accuracy | Vul Precision | Vul Recall | Vul F1   | CWE Accuracy |
+|-------|---------------|-----------------|--------------|---------------|------------|----------|--------------|
+| 1     | 0.3038        | 0.4997          | 0.9570       | 0.8082        | 0.5379     | 0.6459   | 0.1887       |
+| 2     | 0.6092        | 0.4859          | 0.9587       | 0.8118        | 0.5641     | 0.6657   | 0.2964       |
+| 3     | 0.4261        | 0.5090          | 0.9585       | 0.8114        | 0.5605     | 0.6630   | 0.3323       |
+
+- **Final Training Loss**: 0.4430 (average over all steps)
+
+## Evaluation
+
+The model was evaluated on the test split (26,738 samples) with the following metrics:
+- **Vulnerability Detection**:
+  - Accuracy: 0.9571
+  - Precision: 0.7947
+  - Recall: 0.5437
+  - F1 Score: 0.6457
+- **CWE Classification** (on vulnerable samples):
+  - Accuracy: 0.3288
+
+The model excels at identifying non-vulnerable code (high accuracy) but has moderate recall for vulnerabilities and lower CWE classification accuracy, indicating room for improvement in CWE prediction.
+
+## Usage
+
+### Installation
+Install the required libraries:
+```bash
+pip install transformers torch datasets huggingface_hub
+
+```
+Apologies for the oversight! Below is the corrected README.md with the entire content, including the "Sample Code Snippet" section through to the end, formatted properly in Markdown.
+
+markdown
+
+Collapse
+
+Wrap
+
+Copy
+# UnixCoder-Primevul-BigVul Model Card
+
+## Model Overview
+
+`UnixCoder-Primevul-BigVul` is a multi-task model based on Microsoft's `unixcoder-base`, fine-tuned to detect vulnerabilities (`vul`) and classify Common Weakness Enumeration (CWE) types in code snippets. It was developed by [mahdin70](https://huggingface.co/mahdin70) and trained on a balanced dataset combining BigVul and PrimeVul datasets. The model performs binary classification for vulnerability detection and multi-class classification for CWE identification.
+
+- **Model Repository**: [mahdin70/UnixCoder-Primevul-BigVul](https://huggingface.co/mahdin70/UnixCoder-Primevul-BigVul)
+- **Base Model**: [microsoft/unixcoder-base](https://huggingface.co/microsoft/unixcoder-base)
+- **Tasks**: Vulnerability Detection (Binary), CWE Classification (Multi-class)
+- **License**: MIT (assumed; adjust if different)
+- **Date**: Trained and uploaded as of March 11, 2025
+
+## Model Architecture
+
+The model extends `unixcoder-base` with two task-specific heads:
+- **Vulnerability Head**: A linear layer mapping 768-dimensional hidden states to 2 classes (vulnerable or not).
+- **CWE Head**: A linear layer mapping 768-dimensional hidden states to 135 classes (134 CWE types + 1 for "no CWE").
+
+The architecture is implemented as a custom `MultiTaskUnixCoder` class in PyTorch, with the loss computed as the sum of cross-entropy losses for both tasks.
+
+## Training Dataset
+
+The model was trained on the `mahdin70/balanced_merged_bigvul_primevul` dataset (configuration: `10_per_commit`), which combines:
+- **BigVul**: A dataset of real-world vulnerabilities from open-source projects.
+- **PrimeVul**: A dataset focused on prime vulnerabilities in code.
+
+### Dataset Details
+- **Splits**:
+  - Train: 124,780 samples
+  - Validation: 26,740 samples
+  - Test: 26,738 samples
+- **Features**:
+  - `func`: Code snippet (text)
+  - `vul`: Binary label (0 = non-vulnerable, 1 = vulnerable)
+  - `CWE ID`: CWE identifier (e.g., CWE-89) or None for non-vulnerable samples
+- **Preprocessing**:
+  - CWE labels were encoded using a `LabelEncoder` with 134 unique CWE classes identified across the dataset.
+  - Non-vulnerable samples assigned a CWE label of -1 (mapped to 0 in the model).
+
+The dataset is balanced to ensure a fair representation of vulnerable and non-vulnerable samples, with a maximum of 10 samples per commit where applicable.
+
+## Training Details
+
+### Training Arguments
+The model was trained using the Hugging Face `Trainer` API with the following arguments:
+- **Output Directory**: `./unixcoder_multitask`
+- **Evaluation Strategy**: Per epoch
+- **Save Strategy**: Per epoch
+- **Learning Rate**: 2e-5
+- **Batch Size**: 8 (per device, train and eval)
+- **Epochs**: 3
+- **Weight Decay**: 0.01
+- **Logging**: Every 10 steps, logged to `./logs`
+- **WANDB**: Disabled
+
+### Training Environment
+- **Hardware**: NVIDIA Tesla T4 GPU
+- **Framework**: PyTorch 2.5.1+cu121, Transformers 4.47.0
+- **Duration**: ~6 hours, 34 minutes, 53 seconds (23,397 steps)
+
+### Training Metrics
+Validation metrics across epochs:
+
+| Epoch | Training Loss | Validation Loss | Vul Accuracy | Vul Precision | Vul Recall | Vul F1   | CWE Accuracy |
+|-------|---------------|-----------------|--------------|---------------|------------|----------|--------------|
+| 1     | 0.3038        | 0.4997          | 0.9570       | 0.8082        | 0.5379     | 0.6459   | 0.1887       |
+| 2     | 0.6092        | 0.4859          | 0.9587       | 0.8118        | 0.5641     | 0.6657   | 0.2964       |
+| 3     | 0.4261        | 0.5090          | 0.9585       | 0.8114        | 0.5605     | 0.6630   | 0.3323       |
+
+- **Final Training Loss**: 0.4430 (average over all steps)
+
+## Evaluation
+
+The model was evaluated on the test split (26,738 samples) with the following metrics:
+- **Vulnerability Detection**:
+  - Accuracy: 0.9571
+  - Precision: 0.7947
+  - Recall: 0.5437
+  - F1 Score: 0.6457
+- **CWE Classification** (on vulnerable samples):
+  - Accuracy: 0.3288
+
+The model excels at identifying non-vulnerable code (high accuracy) but has moderate recall for vulnerabilities and lower CWE classification accuracy, indicating room for improvement in CWE prediction.
+
+## Usage
+
+### Installation
+Install the required libraries:
+```bash
+pip install transformers torch datasets huggingface_hub
+```
+
+### Sample Code Snippet
+Below is an example of how to use the model for inference on a code snippet:
+
+```python
+from transformers import AutoTokenizer, AutoModel
+import torch
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("microsoft/unixcoder-base")
+model = AutoModel.from_pretrained("mahdin70/UnixCoder-Primevul-BigVul", trust_remote_code=True)
+model.eval()
+
+# Example code snippet
+code = """
+void example(char *input) {
+    char buffer[10];
+    strcpy(buffer, input);
+}
+"""
+
+# Tokenize input
+inputs = tokenizer(code, return_tensors="pt", padding="max_length", truncation=True, max_length=512)
+
+# Move to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+inputs = {k: v.to(device) for k, v in inputs.items()}
+
+# Get predictions
+with torch.no_grad():
+    outputs = model(**inputs)
+    vul_logits = outputs["vul_logits"]
+    cwe_logits = outputs["cwe_logits"]
+
+    # Vulnerability prediction
+    vul_pred = torch.argmax(vul_logits, dim=1).item()
+    print(f"Vulnerability: {'Vulnerable' if vul_pred == 1 else 'Not Vulnerable'}")
+
+    # CWE prediction (if vulnerable)
+    if vul_pred == 1:
+        cwe_pred = torch.argmax(cwe_logits, dim=1).item() - 1  # Subtract 1 as -1 is "no CWE"
+        print(f"Predicted CWE: {cwe_pred if cwe_pred >= 0 else 'None'}")
+
+```
+
+### Output Example:
+
+```bash
+Vulnerability: Vulnerable
+Predicted CWE: 120  # Maps to CWE-120 (Buffer Overflow), depending on encoder
+```
+
+## Notes:
+
+The CWE prediction is an integer index (0 to 133). To map it to a specific CWE ID (e.g., CWE-120), you need the LabelEncoder used during training, available in the dataset preprocessing step.
+Ensure trust_remote_code=True as the model uses custom code from the repository.
+
+## Limitations
+- CWE Accuracy: The model struggles with precise CWE classification (32.88% accuracy), likely due to class imbalance or complexity in distinguishing similar CWE types.
+- Recall: Moderate recall (54.37%) for vulnerability detection suggests some vulnerable samples may be missed.
+- Generalization: Trained on BigVul and PrimeVul, performance may vary on out-of-domain codebases.
+
+## Future Improvements
+- Increase training epochs or dataset size for better CWE accuracy.
+- Experiment with class weighting to address CWE imbalance.
+- Fine-tune on additional datasets for broader generalization.