CoolFace
Modelpublic

FBogaerts/plbart-multi_task-python-Finetuned

sourceHugging Facemitupdated 1y agoView on Hugging Face
0likes29downloads
README.md104 linesDownload Raw Back to root
1---2license: mit3base_model: uclanlp/plbart-multi_task-python4language:5- en6library_name: transformers7tags:8- text-generation9- code-generation10- vulnerability-injection11- security12- vaitp13- finetuned14pretty_name: "FBogaerts/plbart-multi_task-python-Finetuned Finetuned for Vulnerability Injection"15---16 17# FBogaerts/plbart-multi_task-python-Finetuned Finetuned for Vulnerability Injection (VAITP)18 19This model is a fine-tuned version of **uclanlp/plbart-multi_task-python** specialized for the task of security vulnerability injection in Python code. It has been trained to follow a specific instruction format to precisely modify code snippets and introduce vulnerabilities.20 21This model was developed as part of the research for our paper: *(coming soon)*.22 23The VAITP CLI Framework and related resources can be found at our [GitHub repository](coming soon).24 25## Model Description26 27This model was fine-tuned to act as a "Coder" LLM. It takes a specific instruction set and a piece of original Python code, and its objective is to return the modified code with the requested vulnerability injected.28 29The model excels when prompted using the specific format it was trained on.30 31## Intended Uses & Limitations32 33**Intended Use**34 35This model is intended for research purposes in the field of automated security testing, SAST/DAST tool evaluation, and the generation of training data for security-aware models. It should be used within a sandboxed environment to inject vulnerabilities into non-production code for analysis.36 37**Out-of-Scope Uses**38 39This model should **NOT** be used for:40-   Generating malicious code for use in real-world attacks.41-   Directly modifying production codebases.42-   Any application outside of controlled, ethical security research.43 44The generated code should always be manually reviewed before use.45 46## How to Use47 48This model expects a very specific prompt format, which we call the `FINETUNED_STYLE` in our paper. The format is:49 50`{instruction} _BREAK_ {original_code}`51 52Here is an example using `transformers`:53 54```python55from transformers import AutoTokenizer, AutoModelForCausalLM56 57model_name = "FBogaerts/plbart-multi_task-python-Finetuned"58tokenizer = AutoTokenizer.from_pretrained(model_name)59model = AutoModelForCausalLM.from_pretrained(model_name)60 61instruction = "Modify the function to introduce a OS Command Injection vulnerability. The vulnerable code must contain the pattern: 'User-controlled input is used in a subprocess call with shell=True'."62original_code = "import subprocess\ndef execute(cmd):\n    subprocess.run(cmd, shell=False)"63 64prompt = f"{instruction} _BREAK_ {original_code}"65 66inputs = tokenizer(prompt, return_tensors="pt")67outputs = model.generate(**inputs, max_new_tokens=256)68 69vulnerable_code = tokenizer.decode(outputs[0], skip_special_tokens=True)70# The model will output the full modified code block.71# Further cleaning may be needed to extract only the code.72print(vulnerable_code)73```74Training Procedure75 76Training Data77 78The model was fine-tuned on a dataset of 1,406 examples derived from the DeVAITP Vulnerability Corpus. Each example consists of a triplet: (instruction, original_code, vulnerable_code). The instructions were generated using the meta-prompting technique described in our paper, with meta-llama/Meta-Llama-3.1-8B-Instruct serving as the Planner model.79 80Training Hyperparameters81 82The model was fine-tuned using the following key hyperparameters:83 84    Framework: Hugging Face TRL85 86    Learning Rate:  2e-587 88    Number of Epochs: 189 90    Batch Size: 191 92    Hardware: Google Colab (L4 GPU)93 94Evaluation95 96(coming soon)97 98Citation99 100If you use this model in your research, please cite our paper:101(BibTeX entry will be provided upon publication)102 103 104