Spaces:
Sleeping
Sleeping
File size: 3,410 Bytes
7730772 8ebacce 7730772 8ebacce |
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 |
from datasets import load_dataset, Dataset
from sklearn.model_selection import train_test_split
from transformers import (
BertTokenizer,
AutoModelForSequenceClassification,
Trainer,
TrainingArguments
)
import torch
from sklearn.metrics import accuracy_score, precision_recall_fscore_support
import numpy as np
def compute_metrics(eval_pred):
logits, labels = eval_pred
preds = np.argmax(logits, axis=-1)
precision, recall, f1, _ = precision_recall_fscore_support(labels, preds, average='binary')
acc = accuracy_score(labels, preds)
return {
'accuracy': acc,
'f1': f1,
'precision': precision,
'recall': recall
}
def main():
# Check for GPU availability
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Load and prepare dataset
print("Loading dataset...")
dataset = load_dataset("ealvaradob/phishing-dataset", "combined_reduced", trust_remote_code=True)
df = dataset['train'].to_pandas()
# Split dataset
train_df, test_df = train_test_split(df, test_size=0.2, random_state=42)
train_dataset = Dataset.from_pandas(train_df, preserve_index=False)
test_dataset = Dataset.from_pandas(test_df, preserve_index=False)
# Initialize tokenizer and model
print("Initializing model...")
tokenizer = BertTokenizer.from_pretrained('bert-large-uncased')
model = AutoModelForSequenceClassification.from_pretrained(
'bert-large-uncased',
num_labels=2
).to(device)
def tokenize_function(examples):
return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)
# Tokenize datasets
print("Tokenizing datasets...")
train_dataset = train_dataset.map(tokenize_function, batched=True)
test_dataset = test_dataset.map(tokenize_function, batched=True)
# Convert to PyTorch datasets
train_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
test_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])
# Set up training arguments
epochs = 3
batch_size = 64
training_args = TrainingArguments(
output_dir="./results",
evaluation_strategy="epoch",
save_strategy="epoch",
learning_rate=5e-5,
per_device_train_batch_size=batch_size,
per_device_eval_batch_size=batch_size,
num_train_epochs=epochs,
weight_decay=0.01,
logging_dir='./logs',
logging_steps=50,
load_best_model_at_end=True,
metric_for_best_model="accuracy"
)
# Define Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=test_dataset,
tokenizer=tokenizer,
compute_metrics=compute_metrics
)
# Train model
print("Starting training...")
trainer.train()
# Evaluate the model
print("Evaluating model...")
eval_results = trainer.evaluate()
print(eval_results)
# Save the model and tokenizer
print("Saving model...")
model_path = "./phishing_model"
model.save_pretrained(model_path)
tokenizer.save_pretrained(model_path)
print(f"Model and tokenizer saved to {model_path}")
print("Training completed and model saved!")
if __name__ == "__main__":
main()
|