SpamEmailNew / SpamEmail.py
Woubit's picture
Update SpamEmail.py
be27c73 verified
Raw
History Blame Contribute Delete
807 Bytes
import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
import joblib
# 1. Load Dataset (Ensure you have a CSV with 'text' and 'label' columns)
df = pd.read_csv('spam_dataset.csv')
# 2. Vectorize the email text
vectorizer = CountVectorizer()
X = vectorizer.fit_transform(df['text'])
y = df['label']
# 3. Train the SVM model
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
svm_model = SVC(kernel='linear')
svm_model.fit(X_train, y_train)
# 4. Save the Model and Vectorizer (Crucial for later use)
joblib.dump(svm_model, 'svm_spam_model.pkl')
joblib.dump(vectorizer, 'vectorizer.pkl')
print("Model and Vectorizer saved successfully!")