blob: 8502199ae965b4e95b43ebc2a7f18f6c33dc33c1 [file] [log] [blame]
Matt Sharifibda09f12017-03-10 12:29:15 +01001/*
2 * Copyright (C) 2017 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#ifndef LIBTEXTCLASSIFIER_SMARTSELECT_TOKEN_FEATURE_EXTRACTOR_H_
18#define LIBTEXTCLASSIFIER_SMARTSELECT_TOKEN_FEATURE_EXTRACTOR_H_
19
Lukas Zilkad3bc59a2017-04-03 17:32:27 +020020#include <memory>
Matt Sharifibda09f12017-03-10 12:29:15 +010021#include <vector>
22
23#include "base.h"
24#include "smartselect/types.h"
Lukas Zilka26e8c2e2017-04-06 15:54:24 +020025#include "util/strings/stringpiece.h"
Lukas Zilkad3bc59a2017-04-03 17:32:27 +020026#include "unicode/regex.h"
Matt Sharifibda09f12017-03-10 12:29:15 +010027
28namespace libtextclassifier {
29
30struct TokenFeatureExtractorOptions {
31 // Number of buckets used for hashing charactergrams.
Lukas Zilkad3bc59a2017-04-03 17:32:27 +020032 int num_buckets = 0;
Matt Sharifibda09f12017-03-10 12:29:15 +010033
34 // Orders of charactergrams to extract. E.g., 2 means character bigrams, 3
35 // character trigrams etc.
36 std::vector<int> chargram_orders;
37
38 // Whether to extract the token case feature.
Lukas Zilkad3bc59a2017-04-03 17:32:27 +020039 bool extract_case_feature = false;
40
41 // If true, will use the unicode-aware functionality for extracting features.
42 bool unicode_aware_features = false;
Matt Sharifibda09f12017-03-10 12:29:15 +010043
44 // Whether to extract the selection mask feature.
Lukas Zilkad3bc59a2017-04-03 17:32:27 +020045 bool extract_selection_mask_feature = false;
46
47 // Regexp features to extract.
48 std::vector<std::string> regexp_features;
49
50 // Whether to remap digits to a single number.
51 bool remap_digits = false;
52
53 // Maximum length of a word.
54 int max_word_length = 20;
Matt Sharifibda09f12017-03-10 12:29:15 +010055};
56
57class TokenFeatureExtractor {
58 public:
Lukas Zilkad3bc59a2017-04-03 17:32:27 +020059 explicit TokenFeatureExtractor(const TokenFeatureExtractorOptions& options);
Matt Sharifibda09f12017-03-10 12:29:15 +010060
61 // Extracts features from a token.
Lukas Zilka6bb39a82017-04-07 19:55:11 +020062 // - is_in_span is a bool indicator whether the token is a part of the
63 // selection span (true) or not (false).
Matt Sharifibda09f12017-03-10 12:29:15 +010064 // - sparse_features are indices into a sparse feature vector of size
65 // options.num_buckets which are set to 1.0 (others are implicitly 0.0).
66 // - dense_features are values of a dense feature vector of size 0-2
67 // (depending on the options) for the token
Lukas Zilka6bb39a82017-04-07 19:55:11 +020068 bool Extract(const Token& token, bool is_in_span,
69 std::vector<int>* sparse_features,
Matt Sharifibda09f12017-03-10 12:29:15 +010070 std::vector<float>* dense_features) const;
71
Lukas Zilka6bb39a82017-04-07 19:55:11 +020072 int DenseFeaturesCount() const {
73 return options_.extract_case_feature +
74 options_.extract_selection_mask_feature + regex_patterns_.size();
75 }
Matt Sharifibda09f12017-03-10 12:29:15 +010076
77 protected:
78 // Hashes given token to given number of buckets.
Lukas Zilka26e8c2e2017-04-06 15:54:24 +020079 int HashToken(StringPiece token) const;
Matt Sharifibda09f12017-03-10 12:29:15 +010080
81 // Extracts the charactergram features from the token.
82 std::vector<int> ExtractCharactergramFeatures(const Token& token) const;
83
Lukas Zilkad3bc59a2017-04-03 17:32:27 +020084 // Extracts the charactergram features from the token in a non-unicode-aware
85 // way.
86 std::vector<int> ExtractCharactergramFeaturesAscii(const Token& token) const;
87
88 // Extracts the charactergram features from the token in a unicode-aware way.
89 std::vector<int> ExtractCharactergramFeaturesUnicode(
90 const Token& token) const;
91
Matt Sharifibda09f12017-03-10 12:29:15 +010092 private:
93 TokenFeatureExtractorOptions options_;
Lukas Zilkad3bc59a2017-04-03 17:32:27 +020094
95 std::vector<std::unique_ptr<icu::RegexPattern>> regex_patterns_;
Matt Sharifibda09f12017-03-10 12:29:15 +010096};
97
98} // namespace libtextclassifier
99
100#endif // LIBTEXTCLASSIFIER_SMARTSELECT_TOKEN_FEATURE_EXTRACTOR_H_