47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
from collections import Counter
|
|
|
|
import numpy as np
|
|
from sklearn.base import BaseEstimator, ClassifierMixin
|
|
from sklearn.feature_extraction import DictVectorizer
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.pipeline import Pipeline
|
|
|
|
from classifiers.feature_classifier import FeatureExtractor
|
|
from cleaning.transformers import LowercaseTransformer, PunctuationRemover
|
|
|
|
|
|
class BagOfWordsClassifier(BaseEstimator, ClassifierMixin):
|
|
def __init__(self):
|
|
self.cleaning = Pipeline([
|
|
("lowercase", LowercaseTransformer()),
|
|
("punctuation", PunctuationRemover()),
|
|
])
|
|
|
|
def fit(self, X, y):
|
|
X_clean = self.cleaning.fit_transform(X)
|
|
self._pipeline = Pipeline([
|
|
("features", FeatureExtractor(self.extract_features)),
|
|
("vectorizer", DictVectorizer()),
|
|
("classifier", LogisticRegression(max_iter=1000)),
|
|
])
|
|
y_binary = (np.array(y) == "spam").astype(int)
|
|
self._pipeline.fit(X_clean, y_binary)
|
|
return self
|
|
|
|
def predict(self, X):
|
|
X_clean = self.cleaning.transform(X)
|
|
y_binary = self._pipeline.predict(X_clean)
|
|
return np.where(y_binary == 1, "spam", "ham")
|
|
|
|
def extract_features(self, message):
|
|
return dict(Counter(message.split()))
|
|
|
|
def feature_weights(self, top_n=10):
|
|
vectorizer = self._pipeline.named_steps["vectorizer"]
|
|
classifier = self._pipeline.named_steps["classifier"]
|
|
names = vectorizer.get_feature_names_out()
|
|
weights = classifier.coef_[0]
|
|
pairs = sorted(zip(names, weights), key=lambda x: x[1])
|
|
half = top_n // 2
|
|
return pairs[-half:][::-1] + pairs[:half]
|