init commit
This commit is contained in:
75
lab4/train_spam_model.py
Normal file
75
lab4/train_spam_model.py
Normal file
@@ -0,0 +1,75 @@
|
||||
import joblib
|
||||
import csv
|
||||
|
||||
import numpy as np
|
||||
from sklearn.feature_extraction.text import TfidfVectorizer
|
||||
from sklearn.linear_model import LogisticRegression, Ridge
|
||||
from sklearn.metrics import (
|
||||
accuracy_score,
|
||||
classification_report,
|
||||
mean_absolute_error,
|
||||
r2_score,
|
||||
)
|
||||
from sklearn.model_selection import train_test_split
|
||||
|
||||
|
||||
def load_data(path="data.csv"):
|
||||
texts, labels, scores = [], [], []
|
||||
with open(path, newline="", encoding="utf-8") as f:
|
||||
for row in csv.DictReader(f):
|
||||
texts.append(row["text"])
|
||||
labels.append(int(row["label"]))
|
||||
scores.append(float(row["spam_score"]))
|
||||
return texts, np.array(labels), np.array(scores)
|
||||
|
||||
|
||||
def main():
|
||||
texts, labels, scores = load_data()
|
||||
|
||||
X_train_txt, X_test_txt, y_train, y_test, s_train, s_test = train_test_split(
|
||||
texts, labels, scores, test_size=0.25, random_state=42, stratify=labels
|
||||
)
|
||||
|
||||
vectorizer = TfidfVectorizer(lowercase=True, ngram_range=(1, 2), max_features=2000)
|
||||
X_train = vectorizer.fit_transform(X_train_txt)
|
||||
X_test = vectorizer.transform(X_test_txt)
|
||||
|
||||
# Classification Model (Spam/Ham)
|
||||
clf = LogisticRegression(max_iter=1000)
|
||||
clf.fit(X_train, y_train)
|
||||
y_pred = clf.predict(X_test)
|
||||
|
||||
print("Classification: Spam vs Ham")
|
||||
print(f"Accuracy: {accuracy_score(y_test, y_pred):.3f}")
|
||||
print(classification_report(y_test, y_pred, target_names=["ham", "spam"]))
|
||||
|
||||
# Regression model
|
||||
reg = Ridge(alpha=1.0)
|
||||
reg.fit(X_train, s_train)
|
||||
s_pred = reg.predict(X_test)
|
||||
s_pred = np.clip(s_pred, 0, 1)
|
||||
|
||||
print("Regression: Spam Score (0-1)")
|
||||
print(f"MAE: {mean_absolute_error(s_test, s_pred):.3f}")
|
||||
print(f"R^2: {r2_score(s_test, s_pred):.3f}")
|
||||
|
||||
# Demo model
|
||||
print("\nTest Samples")
|
||||
samples = [
|
||||
"Congratulations! You've won a free vacation, click here now!",
|
||||
"Hey, can you send me the meeting notes from earlier?",
|
||||
]
|
||||
X_demo = vectorizer.transform(samples)
|
||||
preds = clf.predict(X_demo)
|
||||
scores_demo = np.clip(reg.predict(X_demo), 0, 1)
|
||||
for text, label, score in zip(samples, preds, scores_demo):
|
||||
verdict = "SPAM" if label == 1 else "HAM"
|
||||
print(f"[{verdict}] score={score:.2f} :: {text}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# persist our model
|
||||
joblib.dump(vectorizer, "vectorizer.joblib")
|
||||
joblib.dump(clf, "classifier.joblib")
|
||||
print("\nSaved model artifacts.")
|
||||
main()
|
||||
Reference in New Issue
Block a user