Initial commit

cleaning/transformers.py

@@ -0,0 +1,45 @@
+import re
+
+import numpy as np
+from sklearn.base import BaseEstimator, TransformerMixin
+
+
+STOPWORDS = {
+    "a", "an", "the", "is", "it", "in", "on", "at", "to", "for",
+    "of", "and", "or", "but", "not", "with", "as", "by", "from",
+    "this", "that", "was", "are", "be", "been", "have", "has",
+    "had", "do", "did", "will", "would", "could", "should",
+    "i", "me", "my", "you", "your", "he", "she", "we", "they",
+    "his", "her", "our", "their", "its", "what", "which",
+}
+
+
+class LowercaseTransformer(BaseEstimator, TransformerMixin):
+    def fit(self, X, y=None):
+        self.fitted_ = True
+        return self
+
+    def transform(self, X):
+        return np.array([msg.lower() for msg in X])
+
+
+class StopwordRemover(BaseEstimator, TransformerMixin):
+    def fit(self, X, y=None):
+        self.fitted_ = True
+        return self
+
+    def transform(self, X):
+        return np.array([self._remove(msg) for msg in X])
+
+    def _remove(self, message):
+        words = message.split()
+        return " ".join(w for w in words if w.lower() not in STOPWORDS)
+
+
+class PunctuationRemover(BaseEstimator, TransformerMixin):
+    def fit(self, X, y=None):
+        self.fitted_ = True
+        return self
+
+    def transform(self, X):
+        return np.array([re.sub(r"[^\w\s]", " ", msg) for msg in X])