Switch to unified view

a b/src/features/token_length_feature.py
1
# Base Dependencies
2
# ----------------
3
import numpy
4
5
# Local Dependencies
6
# ------------------
7
from models import RelationCollection
8
9
# 3rd-Party Dependencies
10
# ----------------------
11
from sklearn.base import BaseEstimator
12
13
14
class TokenLengthFeature(BaseEstimator):
15
    """
16
    TokenLengthFeature
17
18
    Computes the number of tokens of each relation. This is used in evaluation.
19
    """
20
21
    def __init__(self):
22
        pass
23
24
    def get_feature_names(self, input_features=None):
25
        return ["token_length"]
26
27
    def create_token_length_feature(
28
        self, collection: RelationCollection
29
    ) -> numpy.array:
30
        features = []
31
        for doc in collection.tokens:
32
            features.append([len(doc)])
33
34
        return numpy.array(features)
35
36
    def fit(self, x: RelationCollection, y=None):
37
        return self
38
39
    def transform(self, x: RelationCollection, y=None) -> numpy.array:
40
        return self.create_token_length_feature(x)
41
42
    def fit_transform(self, x: RelationCollection, y=None) -> numpy.array:
43
        return self.create_token_length_feature(x)