1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
|
// Copyright 2022 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
#ifndef COMPONENTS_OMNIBOX_BROWSER_ON_DEVICE_TAIL_TOKENIZER_H_
#define COMPONENTS_OMNIBOX_BROWSER_ON_DEVICE_TAIL_TOKENIZER_H_
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "base/containers/flat_map.h"
#include "base/containers/flat_set.h"
#include "base/files/file_path.h"
// The tokenizer performs tokenization for on device tail machine learning
// model. It basically maps raw strings to/from tokens, and tokens to/from IDs
// accepted by the ML model or vice versa based on the given vocabulary file.
// This tokenizer has not supported CJK yet.
class OnDeviceTailTokenizer {
public:
using TokenId = int32_t;
using TokenIds = std::vector<TokenId>;
// Data structure to store tokenization information.
struct Tokenization {
// Unambiguous token IDs. This should at least include the begin query
// token.
TokenIds unambiguous_ids;
// Human-readable unambiguous part of the prefix.
std::string unambiguous_prefix;
// The constraint prefix for the next forward RNN step, if the last typed
// token was ambiguous. For example, given prefix [(pa)(t)] and if the
// trailing (t) could match multiple tokens, constraint prefix will be set
// as "t" and only outputs matching this prefix from the next step will be
// kept.
std::string constraint_prefix;
Tokenization();
~Tokenization();
};
~OnDeviceTailTokenizer();
OnDeviceTailTokenizer();
// Loads the vocabulary file and initializes the tokenizer.
bool Init(const base::FilePath& vocabulary_filepath);
// Determines whether the instance is successfully initialized.
bool IsReady() const;
// Fills the Tokenization struct for the given prefix.
void CreatePrefixTokenization(const std::string& prefix,
Tokenization* tokenization) const;
// Tokenizes the previous query greedily.
void TokenizePrevQuery(const std::string& prev_query,
TokenIds* prev_query_token_ids) const;
// Resets tokens <-> IDs maps.
void Reset();
// Maps token to ID and vice versa.
std::string IdToToken(const TokenId token_id) const;
TokenId TokenToId(const std::string& token) const;
// Returns the size of the vocabulary.
size_t vocab_size() const { return token_to_id_.size(); }
// Special query token related helpers.
bool IsEndQueryTokenId(TokenId token_id) const;
bool IsBeginQueryTokenId(TokenId token_id) const;
TokenId GetEndQueryTokenId() const;
// Determines if the token related to the given ID can be properly printed.
bool IsTokenPrintable(TokenId token_id) const;
private:
// Determines if the given token is ambiguous.
bool IsAmbiguousToken(const std::string& token) const;
// Initializes the ambiguous tokens map.
void InitAmbiguousMap();
// Encodes the given raw string to its corresponding token and ID pairs.
// Note we always use the longest tokens in the vocabulary first and gradually
// switch to shorter tokens until a match for the prefix of the remaining
// string is found. Then jump to the start of the unmatched part of the string
// and do this again until we match all characters of the string.
// For example, given vocabulary:
// [1:a], [2:b], [3:c], [4:abc], [5:ab]
// Encoding:
// string:abcabc -> tokens:[abc][abc] -> IDs:[4][4]
// string:abcab -> tokens:[abc][ab] -> IDs:[4][5]
// string:abcaba -> tokens:[abc][ab][a] -> IDs:[4][5][1]
// string:cbacba -> tokens:[c][b][a][c][b][a] -> IDs:[3][2][1][3][2][1]
void EncodeRawString(
const std::string& raw_string,
std::vector<std::pair<std::string, TokenId>>* token_and_ids) const;
// Insert token and its ID to tokens <-> IDs maps.
void InsertTokenToMaps(const std::string& token);
// Maps for tokens <-> IDs.
base::flat_map<std::string, TokenId> token_to_id_;
std::vector<std::string> id_to_token_;
// The max length of tokens in the vocabulary.
size_t max_token_length_;
// The list of ambiguous tokens.
base::flat_set<std::string> ambiguous_tokens_;
};
#endif // COMPONENTS_OMNIBOX_BROWSER_ON_DEVICE_TAIL_TOKENIZER_H_
|
