File size: 1,830 Bytes
497f2f3 | 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 | from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments
from datasets import load_dataset
# 1. Preparar o Dataset
def load_data_from_csv(csv_file):
dataset = load_dataset("csv", data_files=csv_file)
return dataset['train']
# 2. Configurar o Tokenizer e Modelo
def get_model_and_tokenizer():
model_name = "microsoft/codebert-base"
tokenizer = RobertaTokenizer.from_pretrained(model_name)
model = RobertaForSequenceClassification.from_pretrained(model_name, num_labels=2) # 2 classes: safe/unsafe
return model, tokenizer
# 3. Tokenizar os Dados
def tokenize_function(example, tokenizer):
return tokenizer(example['content'], truncation=True, padding="max_length", max_length=512)
# 4. Treinar o Modelo
def train_model(dataset, tokenizer, model):
tokenized_data = dataset.map(lambda x: tokenize_function(x, tokenizer), batched=True)
training_args = TrainingArguments(
output_dir="./results",
evaluation_strategy="epoch",
save_strategy="epoch",
learning_rate=2e-5,
num_train_epochs=3,
per_device_train_batch_size=16,
per_device_eval_batch_size=16,
warmup_steps=500,
weight_decay=0.01,
logging_dir="./logs",
logging_steps=10,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_data,
eval_dataset=tokenized_data,
tokenizer=tokenizer,
)
trainer.train()
if __name__ == "__main__":
# Carregar Dados e Modelo
dataset = load_data_from_csv("code_analysis_dataset.csv")
model, tokenizer = get_model_and_tokenizer()
# Treinar Modelo
train_model(dataset, tokenizer, model)
print("[SUCCESS] Model trained!")
|