76 lines
2.4 KiB
Python
76 lines
2.4 KiB
Python
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()
|