import pandas as pd | |
from sklearn.model_selection import train_test_split | |
from sklearn.feature_extraction.text import TfidfVectorizer | |
from sklearn.naive_bayes import MultinomialNB | |
from sklearn.pipeline import Pipeline | |
from sklearn import metrics | |
import joblib | |
url = "https://raw.githubusercontent.com/justmarkham/pycon-2016-tutorial/master/data/sms.tsv" | |
df = pd.read_csv(url, sep='\t', header=None, names=['label', 'message']) | |
# Encode labels | |
df['label_num'] = df.label.map({'ham': 0, 'spam': 1}) | |
# Train/test split | |
X_train, X_test, y_train, y_test = train_test_split(df['message'], df['label_num'], test_size=0.2, random_state=42) | |
# Build pipeline | |
model = Pipeline([ | |
('tfidf', TfidfVectorizer()), | |
('nb', MultinomialNB()) | |
]) | |
# Train | |
model.fit(X_train, y_train) | |
# Evaluate | |
preds = model.predict(X_test) | |
print("Accuracy:", metrics.accuracy_score(y_test, preds)) | |
# Save model | |
joblib.dump(model, "spam_classifier_model.joblib") | |