lilbool/vuln-code-analysis

1from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments
2from datasets import load_dataset
3
4# 1. Preparar o Dataset
5def load_data_from_csv(csv_file):
6    dataset = load_dataset("csv", data_files=csv_file)
7    return dataset['train']
8
9# 2. Configurar o Tokenizer e Modelo
10def get_model_and_tokenizer():
11    model_name = "microsoft/codebert-base"
12    tokenizer = RobertaTokenizer.from_pretrained(model_name)
13    model = RobertaForSequenceClassification.from_pretrained(model_name, num_labels=2)  # 2 classes: safe/unsafe
14    return model, tokenizer
15
16# 3. Tokenizar os Dados
17def tokenize_function(example, tokenizer):
18    return tokenizer(example['content'], truncation=True, padding="max_length", max_length=512)
19
20# 4. Treinar o Modelo
21def train_model(dataset, tokenizer, model):
22    tokenized_data = dataset.map(lambda x: tokenize_function(x, tokenizer), batched=True)
23    training_args = TrainingArguments(
24        output_dir="./results",
25        evaluation_strategy="epoch",
26        save_strategy="epoch",
27        learning_rate=2e-5,
28        num_train_epochs=3,
29        per_device_train_batch_size=16,
30        per_device_eval_batch_size=16,
31        warmup_steps=500,
32        weight_decay=0.01,
33        logging_dir="./logs",
34        logging_steps=10,
35    )
36    trainer = Trainer(
37        model=model,
38        args=training_args,
39        train_dataset=tokenized_data,
40        eval_dataset=tokenized_data,
41        tokenizer=tokenizer,
42    )
43    trainer.train()
44
45if __name__ == "__main__":
46    # Carregar Dados e Modelo
47    dataset = load_data_from_csv("code_analysis_dataset.csv")
48    model, tokenizer = get_model_and_tokenizer()
49
50    # Treinar Modelo
51    train_model(dataset, tokenizer, model)
52    print("[SUCCESS] Model trained!")
53