lindera_analysis/token_filter/
mapping.rs1use std::borrow::Cow;
2use std::collections::HashMap;
3
4use daachorse::DoubleArrayAhoCorasick;
5use daachorse::DoubleArrayAhoCorasickBuilder;
6use daachorse::MatchKind;
7use serde_json::Value;
8
9use crate::token_filter::TokenFilter;
10use lindera::LinderaResult;
11use lindera::error::LinderaErrorKind;
12use lindera::token::Token;
13
14pub const MAPPING_TOKEN_FILTER_NAME: &str = "mapping";
15
16pub type MappingTokenFilterConfig = Value;
17
18#[derive(Clone)]
21pub struct MappingTokenFilter {
22 mapping: HashMap<String, String>,
23 trie: DoubleArrayAhoCorasick<u32>,
24}
25
26impl MappingTokenFilter {
27 pub fn new(mapping: HashMap<String, String>) -> LinderaResult<Self> {
41 if mapping.keys().any(|key| key.is_empty()) {
42 return Err(LinderaErrorKind::Args
43 .with_error(anyhow::anyhow!("mapping key must not be empty.")));
44 }
45
46 let mut keyset: Vec<(&[u8], u32)> = Vec::new();
47 let mut keys = mapping.keys().collect::<Vec<_>>();
48 keys.sort();
49 for (value, key) in keys.into_iter().enumerate() {
50 keyset.push((key.as_bytes(), value as u32));
51 }
52
53 let trie = DoubleArrayAhoCorasickBuilder::new()
54 .match_kind(MatchKind::LeftmostLongest)
55 .build_with_values(keyset)
56 .map_err(|err| LinderaErrorKind::Build.with_error(anyhow::anyhow!(err)))?;
57
58 Ok(Self { mapping, trie })
59 }
60
61 pub fn from_config(config: &MappingTokenFilterConfig) -> LinderaResult<Self> {
62 let mapping = config
63 .get("mapping")
64 .and_then(Value::as_object)
65 .ok_or_else(|| {
66 LinderaErrorKind::Parse.with_error(anyhow::anyhow!("mapping must be an object."))
67 })?
68 .iter()
69 .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
70 .collect::<HashMap<String, String>>();
71
72 Self::new(mapping)
73 }
74}
75
76impl TokenFilter for MappingTokenFilter {
77 fn name(&self) -> &'static str {
78 MAPPING_TOKEN_FILTER_NAME
79 }
80
81 fn apply(&self, tokens: &mut Vec<Token<'_>>) -> LinderaResult<()> {
93 for token in tokens.iter_mut() {
94 let mut result = String::with_capacity(token.surface.len());
95
96 {
97 let source = token.surface.as_ref();
98 let mut cursor = 0_usize;
99
100 for m in self.trie.leftmost_find_iter(source) {
101 debug_assert!(m.start() >= cursor && m.end() > m.start());
105
106 result.push_str(&source[cursor..m.start()]);
107 result.push_str(&self.mapping[&source[m.start()..m.end()]]);
108 cursor = m.end();
109 }
110
111 result.push_str(&source[cursor..]);
112 }
113
114 token.surface = Cow::Owned(result);
115 }
116
117 Ok(())
118 }
119}
120
121#[cfg(test)]
122mod tests {
123 use std::collections::HashMap;
124
125 use crate::token_filter::mapping::{MappingTokenFilter, MappingTokenFilterConfig};
126
127 #[test]
128 fn test_mapping_token_filter_empty_key_rejected() {
129 let mut mapping = HashMap::new();
130 mapping.insert(String::new(), "x".to_string());
131 assert!(MappingTokenFilter::new(mapping).is_err());
132
133 let mut mapping = HashMap::new();
134 mapping.insert("a".to_string(), "b".to_string());
135 assert!(MappingTokenFilter::new(mapping).is_ok());
136 }
137
138 #[test]
139 fn test_mapping_token_filter_config() {
140 let config_str = r#"
141 {
142 "mapping": {
143 "ア": "ア",
144 "イ": "イ",
145 "ウ": "ウ",
146 "エ": "エ",
147 "オ": "オ"
148 }
149 }
150 "#;
151 let result: Result<MappingTokenFilterConfig, _> = serde_json::from_str(config_str);
152 assert!(result.is_ok());
153 }
154
155 #[test]
156 fn test_mapping_token_filter() {
157 let config_str = r#"
158 {
159 "mapping": {
160 "ア": "ア",
161 "イ": "イ",
162 "ウ": "ウ",
163 "エ": "エ",
164 "オ": "オ"
165 }
166 }
167 "#;
168 let config = serde_json::from_str::<MappingTokenFilterConfig>(config_str).unwrap();
169
170 let result = MappingTokenFilter::from_config(&config);
171 assert!(result.is_ok());
172 }
173
174 #[test]
175 #[cfg(feature = "embed-ipadic")]
176 fn test_mapping_token_filter_apply_ipadic() {
177 use std::borrow::Cow;
178
179 use crate::token_filter::TokenFilter;
180 use lindera::dictionary::{DictionaryKind, WordId, load_embedded_dictionary};
181 use lindera::token::Token;
182 use lindera_dictionary::viterbi::LexType;
183
184 let config_str = r#"
185 {
186 "mapping": {
187 "籠": "篭"
188 }
189 }
190 "#;
191 let config = serde_json::from_str::<MappingTokenFilterConfig>(config_str).unwrap();
192
193 let filter = MappingTokenFilter::from_config(&config).unwrap();
194
195 let dictionary = load_embedded_dictionary(DictionaryKind::IPADIC).unwrap();
196
197 let mut tokens: Vec<Token> = vec![
198 Token {
199 surface: Cow::Borrowed("籠原"),
200 byte_start: 0,
201 byte_end: 6,
202 position: 0,
203 position_length: 1,
204 word_id: WordId::new(LexType::System, 312630),
205 dictionary: &dictionary,
206 user_dictionary: None,
207 details: Some(vec![
208 Cow::Borrowed("名詞"),
209 Cow::Borrowed("固有名詞"),
210 Cow::Borrowed("一般"),
211 Cow::Borrowed("*"),
212 Cow::Borrowed("*"),
213 Cow::Borrowed("*"),
214 Cow::Borrowed("籠原"),
215 Cow::Borrowed("カゴハラ"),
216 Cow::Borrowed("カゴハラ"),
217 ]),
218 },
219 Token {
220 surface: Cow::Borrowed("駅"),
221 byte_start: 6,
222 byte_end: 9,
223 position: 1,
224 position_length: 1,
225 word_id: WordId::new(LexType::System, 383791),
226 dictionary: &dictionary,
227 user_dictionary: None,
228 details: Some(vec![
229 Cow::Borrowed("名詞"),
230 Cow::Borrowed("接尾"),
231 Cow::Borrowed("地域"),
232 Cow::Borrowed("*"),
233 Cow::Borrowed("*"),
234 Cow::Borrowed("*"),
235 Cow::Borrowed("駅"),
236 Cow::Borrowed("エキ"),
237 Cow::Borrowed("エキ"),
238 ]),
239 },
240 ];
241
242 filter.apply(&mut tokens).unwrap();
243
244 assert_eq!(tokens.len(), 2);
245 assert_eq!(&tokens[0].surface, "篭原");
246 assert_eq!(&tokens[1].surface, "駅");
247 }
248
249 #[test]
250 #[cfg(feature = "embed-ipadic")]
251 fn test_mapping_token_filter_apply_longest_match() {
252 use std::borrow::Cow;
253
254 use crate::token_filter::TokenFilter;
255 use lindera::dictionary::{DictionaryKind, WordId, load_embedded_dictionary};
256 use lindera::token::Token;
257 use lindera_dictionary::viterbi::LexType;
258
259 let config_str = r#"
260 {
261 "mapping": {
262 "ab": "X",
263 "abc": "YY"
264 }
265 }
266 "#;
267 let config = serde_json::from_str::<MappingTokenFilterConfig>(config_str).unwrap();
268
269 let filter = MappingTokenFilter::from_config(&config).unwrap();
270
271 let dictionary = load_embedded_dictionary(DictionaryKind::IPADIC).unwrap();
272
273 let mut tokens: Vec<Token> = vec![Token {
274 surface: Cow::Borrowed("abcab"),
275 byte_start: 0,
276 byte_end: 5,
277 position: 0,
278 position_length: 1,
279 word_id: WordId::new(LexType::System, 0),
280 dictionary: &dictionary,
281 user_dictionary: None,
282 details: None,
283 }];
284
285 filter.apply(&mut tokens).unwrap();
286
287 assert_eq!(tokens.len(), 1);
289 assert_eq!(&tokens[0].surface, "YYX");
290 }
291}