File size: 4,268 Bytes
4deeced db32a43 a46fa4b 4deeced 729180c 7130d3a 4deeced 7130d3a ee9269e a4047d8 4deeced ee9269e 4deeced e05a4e2 4deeced 203feb7 4deeced b8a9004 4deeced db32a43 4deeced a46fa4b 4deeced 7130d3a ee9269e 7130d3a 4deeced a4047d8 4deeced e05a4e2 4deeced 4f2dec9 7130d3a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import os
import torch
from unsloth import FastLanguageModel, is_bfloat16_supported
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset
import gradio as gr
import json
from huggingface_hub import HfApi
max_seq_length = 4096
dtype = None
load_in_4bit = True
hf_token = os.getenv("HF_TOKEN")
current_num = os.getenv("NUM")
print(f"stage ${current_num}")
api = HfApi(token=hf_token)
# models = f"dad1909/CyberSentinel-{current_num}"
model_base = "dad1909/cybersentinal-2.0"
print("Starting model and tokenizer loading...")
# Load the model and tokenizer
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_base,
max_seq_length=max_seq_length,
dtype=dtype,
load_in_4bit=load_in_4bit,
token=hf_token
)
print("Model and tokenizer loaded successfully.")
print("Configuring PEFT model...")
model = FastLanguageModel.get_peft_model(
model,
r=16,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
lora_alpha=16,
lora_dropout=0,
bias="none",
use_gradient_checkpointing="unsloth",
random_state=3407,
use_rslora=False,
loftq_config=None,
)
print("PEFT model configured.")
# Updated alpaca_prompt for different types
alpaca_prompt = {
"learning_from": """Below is a CVE definition.
### CVE definition:
{}
### detail CVE:
{}""",
"definition": """Below is a definition about software vulnerability. Explain it.
### Definition:
{}
### Explanation:
{}""",
"code_vulnerability": """Below is a code snippet. Identify the line of code that is vulnerable and describe the type of software vulnerability.
### Code Snippet:
{}
### Vulnerability solution:
{}"""
}
EOS_TOKEN = tokenizer.eos_token
def detect_prompt_type(instruction):
if instruction.startswith("what is code vulnerable of this code:"):
return "code_vulnerability"
elif instruction.startswith("Learning from"):
return "learning_from"
elif instruction.startswith("what is"):
return "definition"
else:
return "unknown"
def formatting_prompts_func(examples):
instructions = examples["instruction"]
outputs = examples["output"]
texts = []
for instruction, output in zip(instructions, outputs):
prompt_type = detect_prompt_type(instruction)
if prompt_type in alpaca_prompt:
prompt = alpaca_prompt[prompt_type].format(instruction, output)
else:
prompt = instruction + "\n\n" + output
text = prompt + EOS_TOKEN
texts.append(text)
return {"text": texts}
print("Loading dataset...")
dataset = load_dataset("dad1909/DCSV", split="train")
print("Dataset loaded successfully.")
print("Applying formatting function to the dataset...")
dataset = dataset.map(formatting_prompts_func, batched=True)
print("Formatting function applied.")
print("Initializing trainer...")
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
dataset_text_field="text",
max_seq_length=max_seq_length,
dataset_num_proc=2,
packing=False,
args=TrainingArguments(
per_device_train_batch_size=5,
gradient_accumulation_steps=5,
learning_rate=2e-4,
fp16=not is_bfloat16_supported(),
bf16=is_bfloat16_supported(),
warmup_steps=5,
logging_steps=10,
max_steps=200,
optim="adamw_8bit",
weight_decay=0.01,
lr_scheduler_type="linear",
seed=3407,
output_dir="outputs"
),
)
print("Trainer initialized.")
print("Starting training...")
trainer_stats = trainer.train()
print("Training completed.")
num = int(current_num)
num += 1
uploads_models = f"cybersentinal-2.0-{str(num)}"
print("Saving the trained model...")
model.save_pretrained_merged("model", tokenizer, save_method="merged_16bit")
print("Model saved successfully.")
print("Pushing the model to the hub...")
model.push_to_hub_merged(
uploads_models,
tokenizer,
save_method="merged_16bit",
token=hf_token
)
print("Model pushed to hub successfully.")
api.delete_space_variable(repo_id="dad1909/CyberCode", key="NUM")
api.add_space_variable(repo_id="dad1909/CyberCode", key="NUM", value=str(num)) |