Update README.md

README.md

@@ -1,51 +1,58 @@
 ---
 library_name: peft
 base_model: deepseek-ai/deepseek-coder-6.7b-instruct
+license: mit
+language:
+- en
 ---
 
 # Model Card for Model ID
-
-<!-- Provide a quick summary of what the model is/does. -->
-
+Fine-tuned version of `deepseek-coder-6.7b-instruct` aiming to improve vulnerability detection in solidity smart contracts and provide informative explanations on what the vulnerabilities are, and how to solve them.
 
 
 ## Model Details
 
 ### Model Description
 
-<!-- Provide a longer summary of what this model is. -->
+Given the following prompt below:
+```
+Below are one or more Solidity codeblocks. The codeblocks might contain vulnerable code.
+If there is a vulnerability please provide a description of the vulnearblity in terms of the code that is responsible for it.
+Describe how an attacker would be able to take advantage of the vulnerability so the explanation is even more clear.
 
+Output only the description of the vulnerability and the attacking vector. No additional information is needed.
 
+If there is no vulnerability output "There is no vulnearbility".
 
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
+Codeblocks:
+{}
+```
 
-### Model Sources [optional]
+When 1 or more codeblocks are provided to the model using this prompt, the model will output:
+1. Wether there is a vulnerability or not.
+2. What the vulnerability is.
+3. How an attacker would take advantage of the detected vulnerability.
 
-<!-- Provide the basic links for the model. -->
+Afterwards, the above output can be chained to produce a solution - the context has the code, the vulnerability and the attacking vector so deducing a solution becomes a more straight-forward task.
+Additionally, the same fine-tuned model can be used for the solution recommendation as the fine-tuning is low-rank (LoRA) and a lot of the model ability is preserved. 
 
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
-
-## Uses
 
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+- **Developed by:** [Kristian Apostolov]
+- **Shared by:** [Kristian Apostolov]
+- **Model type:** [Decoder]
+- **Language(s) (NLP):** [English]
+- **License:** [MIT]
+- **Finetuned from model:** [deepseek-ai/deepseek-coder-6.7b-instruct]
 
-### Direct Use
+### Model Sources [optional]
 
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+- **Repository:** [https://huggingface.co/msc-smart-contract-auditing/deepseek-coder-6.7b-vulnerability-detection]
 
-[More Information Needed]
+## Uses
 
-### Downstream Use [optional]
+Provide code from a smart contract for a preliminary audit.
 
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+### Direct Use
 
 [More Information Needed]
 
@@ -53,75 +60,112 @@ base_model: deepseek-ai/deepseek-coder-6.7b-instruct
 
 <!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
 
-[More Information Needed]
+Malicious entity could detect 0-day vulnerability and take advantage of it.
 
 ## Bias, Risks, and Limitations
 
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
-
-[More Information Needed]
+The training data could be improved. Audits sometimes describe vulnerabilities which are not necessarily contained in the code itself, but are a part of a larger context. 
 
 ### Recommendations
 
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
-
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. 
 
 ## How to Get Started with the Model
 
 Use the code below to get started with the model.
 
-[More Information Needed]
+```python
+model_name = 'msc-smart-contract-auditing/deepseek-coder-6.7b-vulnerability'
+tokenizer = AutoTokenizer.from_pretrained( # For some reason the tokenizer didn't safe properly
+    "deepseek-ai/deepseek-coder-6.7b-instruct", 
+    trust_remote_code=True, 
+    force_download=True,
+)
 
-## Training Details
+prompt = \
+"""
+Below are one or more Solidity codeblocks. The codeblocks might contain vulnerable code.
+If there is a vulnerability please provide a description of the vulnearblity in terms of the code that is responsible for it.
+Describe how an attacker would be able to take advantage of the vulnerability so the explanation is even more clear.
 
-### Training Data
+Output only the description of the vulnerability and the attacking vector. No additional information is needed.
 
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+If there is no vulnerability output "There is no vulnearbility".
 
-[More Information Needed]
+Codeblocks:
+{}
 
-### Training Procedure
+"""
 
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+codeblocks = "Your code here"
 
-#### Preprocessing [optional]
+messages = [
+    { 'role': 'user', 'content': prompt.format(codeblocks) }
+]
 
-[More Information Needed]
+inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
+outputs = model.generate(inputs, max_new_tokens=512, do_sample=True, top_k=25, top_p=0.95, num_return_sequences=1, eos_token_id=tokenizer.eos_token_id)
+description = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True) 
 
+print(description)
+```
+## Training Details
 
-#### Training Hyperparameters
+### Training Data
 
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+https://huggingface.co/datasets/msc-smart-contract-auditing/audits-with-reasons
 
-#### Speeds, Sizes, Times [optional]
+### Training Procedure
 
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+lora_config = LoraConfig(
+    r=16,  # rank
+    lora_alpha=32,  # scaling factor
+    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj",],
+    lora_dropout=0.05,  # dropout rate for LoRA layers
+)
+
+model = get_peft_model(model, lora_config)
+
+from transformers import Trainer, TrainingArguments
+
+trainer = Trainer(
+    model=model,
+    args=TrainingArguments(
+        per_device_train_batch_size = 2,
+        gradient_accumulation_steps = 4,
+        warmup_steps = 5,
+        num_train_epochs = 1,
+        learning_rate = 2e-4,
+        fp16 = True,
+        logging_steps = 1,
+        optim = "adamw_8bit",
+        weight_decay = 0.01,
+        lr_scheduler_type = "linear",
+        seed = 3407,
+        output_dir = "outputs",
+    ),
+    train_dataset=train_prompts,
+    eval_dataset=test_prompts,
+)
 
-[More Information Needed]
+#### Training Hyperparameters
 
-## Evaluation
+- **Training regime:** fp16 mixed precision
 
-<!-- This section describes the evaluation protocols and provides the results. -->
+## Evaluation
 
 ### Testing Data, Factors & Metrics
 
 #### Testing Data
 
-<!-- This should link to a Dataset Card if possible. -->
-
-[More Information Needed]
+https://huggingface.co/datasets/msc-smart-contract-auditing/audits-with-reasons
 
 #### Factors
 
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-
 [More Information Needed]
 
 #### Metrics
 
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-
 [More Information Needed]
 
 ### Results
@@ -132,12 +176,6 @@ Use the code below to get started with the model.
 
 
 
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-[More Information Needed]
-
 ## Environmental Impact
 
 <!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->