Update README.md

README.md

@@ -1,3 +1,95 @@
----
-license: mit
----
+---
+license: mit
+datasets:
+- mahdin70/cwe_enriched_balanced_bigvul_primevul
+metrics:
+- accuracy
+- precision
+- recall
+- f1
+base_model:
+- microsoft/codebert-base
+library_name: transformers
+---
+
+
+# CodeBERT-VulnCWE - Fine-Tuned CodeBERT for Vulnerability and CWE Classification
+
+## Model Overview
+This model is a fine-tuned version of **microsoft/codebert-base** on a curated and enriched dataset for vulnerability detection and CWE classification. It is capable of predicting whether a given code snippet is vulnerable and, if vulnerable, identifying the specific CWE ID associated with it.
+
+## Dataset
+The model was fine-tuned using the dataset [mahdin70/cwe_enriched_balanced_bigvul_primevul](https://huggingface.co/datasets/mahdin70/cwe_enriched_balanced_bigvul_primevul). The dataset contains both vulnerable and non-vulnerable code samples and is enriched with CWE metadata.
+
+### CWE IDs Covered:
+1. **CWE-119**: Improper Restriction of Operations within the Bounds of a Memory Buffer  
+2. **CWE-20**: Improper Input Validation  
+3. **CWE-125**: Out-of-bounds Read  
+4. **CWE-399**: Resource Management Errors  
+5. **CWE-200**: Information Exposure  
+6. **CWE-787**: Out-of-bounds Write  
+7. **CWE-264**: Permissions, Privileges, and Access Controls  
+8. **CWE-416**: Use After Free  
+9. **CWE-476**: NULL Pointer Dereference  
+10. **CWE-190**: Integer Overflow or Wraparound  
+11. **CWE-189**: Numeric Errors  
+12. **CWE-362**: Concurrent Execution using Shared Resource with Improper Synchronization  
+
+---
+
+## Model Training
+The model was trained for **3 epochs** with the following configuration:
+- **Learning Rate**: 2e-5  
+- **Weight Decay**: 0.01  
+- **Batch Size**: 8  
+- **Optimizer**: AdamW  
+- **Scheduler**: Linear  
+
+### Training Loss and Validation Metrics Per Epoch:
+| Epoch | Training Loss | Validation Loss | Vul Accuracy | Vul Precision | Vul Recall | Vul F1 | CWE Accuracy |
+|-------|---------------|-----------------|--------------|---------------|------------|--------|--------------|
+| 1     | 1.4663        | 1.4988          | 0.7887       | 0.8526        | 0.5498     | 0.6685 | 0.2932       |
+| 2     | 1.2107        | 1.3474          | 0.8038       | 0.8493        | 0.6002     | 0.7034 | 0.3688       |
+| 3     | 1.1885        | 1.3096          | 0.8034       | 0.8020        | 0.6541     | 0.7205 | 0.3963       |
+
+#### Training Summary:
+- **Total Training Steps**: 2958  
+- **Training Loss**: 1.3862  
+- **Training Time**: 3058.7 seconds (~51 minutes)  
+- **Training Speed**: 15.47 samples per second  
+- **Steps Per Second**: 0.967  
+
+
+## How to Use the Model
+```python
+from transformers import AutoModel, AutoTokenizer
+
+model = AutoModel.from_pretrained("mahdin70/CodeBERT-VulnCWE", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
+
+code_snippet = "int main() { int arr[10]; arr[11] = 5; return 0; }"
+inputs = tokenizer(code_snippet, return_tensors="pt")
+outputs = model(**inputs)
+
+vul_logits = outputs["vul_logits"]
+cwe_logits = outputs["cwe_logits"]
+
+vul_pred = vul_logits.argmax(dim=1).item()
+cwe_pred = kov_logits.argmax(dim=1).item()
+
+print(f"Vulnerability: {'Vulnerable' if vul_pred == 1 else 'Non-vulnerable'}")
+print(f"CWE ID: {cwe_pred if vul_pred == 1 else 'N/A'}")
+```
+
+## Limitations and Future Improvements
+- The model achieves a CWE classification accuracy of 39.63% on the validation set, indicating significant room for improvement. Advanced architectures, better data balancing, or additional pretraining could enhance performance.
+- The model's vulnerability detection F1-score (72.05% on validation) is moderate but could be improved with further tuning or a larger dataset.
+- The model may struggle with edge cases or CWEs not well-represented in the training data.
+- Test set evaluation metrics are pending. Running the model on the test set will provide a clearer picture of its generalization.
+
+
+## Notes
+- Ensure the `trust_remote_code=True` flag is used when loading the model, as it relies on custom code for the `MultiTaskCodeBERT` architecture.
+- The model expects input code snippets tokenized using the CodeBERT tokenizer (`microsoft/codebert-base`).
+- For best results, preprocess code snippets consistently with the training dataset (e.g., max length of 512 tokens).
+