FeatureTransformer.py

from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import LabelEncoder
import numpy as np
from MemoryTagger import MemoryTagger


class FeatureTransformer(BaseEstimator, TransformerMixin):
    def __init__(self):
        self.memory_tagger = MemoryTagger()
        self.tag_encoder = LabelEncoder()
        self.pos_encoder = LabelEncoder()

    def fit(self, X, y):
        words = X["word"].values.tolist()
        self.pos = X["POS"].values.tolist()
        tags = X["tag"].values.tolist()
        self.memory_tagger.fit(words, tags)
        self.tag_encoder.fit(tags)
        self.pos_encoder.fit(self.pos)
        return self

    def transform(self, X, y=None):
        def pos_default(p):
            if p in self.pos:
                return self.pos_encoder.transform([p])[0]

            else:
                return -1

        pos = X["POS"].values.tolist()
        words = X["word"].values.tolist()
        out = []

        for i in range(len(words)):
            w = words[i]
            p = pos[i]
            if i < len(words) - 1:
                wp = self.tag_encoder.transform(self.memory_tagger.predict([words[i + 1]]))[0]
                posp = pos_default(pos[i + 1])
            else:
                wp = self.tag_encoder.transform(['EXTRA'])[0]
                posp = pos_default(".")
            if i > 0:
                if words[i - 1] != ".":
                    wm = self.tag_encoder.transform(self.memory_tagger.predict([words[i - 1]]))[0]
                    posm = pos_default(pos[i - 1])
                else:
                    wm = self.tag_encoder.transform(['EXTRA'])[0]
                    posm = pos_default(".")
            else:
                posm = pos_default(".")
                wm = self.tag_encoder.transform(['EXTRA'])[0]

            out.append(np.array([w.istitle(), w.islower(), w.isupper(), len(w), w.isdigit(), w.isalpha(),
            self.tag_encoder.transform(self.memory_tagger.predict([w]))[0],
            pos_default(p), wp, wm, posp, posm]))
        return out