Skip to main content

lindera_ruby/
tokenizer.rs

1//! Tokenizer implementation for morphological analysis.
2//!
3//! This module provides a builder pattern for creating tokenizers and the tokenizer itself.
4//! The build-flow orchestration is delegated to
5//! [`lindera_binding_core::CoreTokenizerBuilder`] / [`lindera_binding_core::CoreTokenizer`];
6//! this module only adds the magnus wrappers and the Ruby-hash conversion.
7
8use std::cell::RefCell;
9use std::path::Path;
10
11use magnus::prelude::*;
12use magnus::{Error, RArray, RHash, Ruby, function, method};
13
14use lindera_binding_core::{CoreTokenizer, CoreTokenizerBuilder};
15
16use crate::dictionary::{RbDictionary, RbUserDictionary};
17use crate::error::to_magnus_error;
18use crate::token::RbToken;
19use crate::util::rb_hash_to_json;
20
21/// Builder for creating a `Tokenizer` with custom configuration.
22///
23/// The builder pattern allows for fluent configuration of tokenizer parameters.
24/// Uses `RefCell` for interior mutability since Magnus `method!` requires `&self`.
25#[magnus::wrap(class = "Lindera::TokenizerBuilder", free_immediately, size)]
26pub struct RbTokenizerBuilder {
27    /// Inner binding-core tokenizer builder (wrapped in RefCell for interior mutability).
28    inner: RefCell<CoreTokenizerBuilder>,
29}
30
31impl RbTokenizerBuilder {
32    /// Creates a new `RbTokenizerBuilder` with default configuration.
33    ///
34    /// # Returns
35    ///
36    /// A new `RbTokenizerBuilder` instance.
37    fn new() -> Result<Self, Error> {
38        let ruby = Ruby::get().expect("Ruby runtime not initialized");
39        let inner =
40            CoreTokenizerBuilder::new().map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
41        Ok(Self {
42            inner: RefCell::new(inner),
43        })
44    }
45
46    /// Loads configuration from a file.
47    ///
48    /// # Arguments
49    ///
50    /// * `file_path` - Path to the configuration file.
51    ///
52    /// # Returns
53    ///
54    /// A new `RbTokenizerBuilder` with the loaded configuration.
55    fn from_file(file_path: String) -> Result<Self, Error> {
56        let ruby = Ruby::get().expect("Ruby runtime not initialized");
57        let inner = CoreTokenizerBuilder::from_file(Path::new(&file_path))
58            .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
59        Ok(Self {
60            inner: RefCell::new(inner),
61        })
62    }
63
64    /// Sets the tokenization mode.
65    ///
66    /// # Arguments
67    ///
68    /// * `mode` - Mode string ("normal" or "decompose").
69    fn set_mode(&self, mode: String) -> Result<(), Error> {
70        let ruby = Ruby::get().expect("Ruby runtime not initialized");
71        let mut builder = self.inner.borrow_mut();
72        builder
73            .set_mode(&mode)
74            .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
75        Ok(())
76    }
77
78    /// Sets the dictionary path.
79    ///
80    /// # Arguments
81    ///
82    /// * `path` - Path to the dictionary directory.
83    fn set_dictionary(&self, path: String) {
84        self.inner.borrow_mut().set_dictionary(&path);
85    }
86
87    /// Sets the user dictionary URI.
88    ///
89    /// # Arguments
90    ///
91    /// * `uri` - URI to the user dictionary.
92    fn set_user_dictionary(&self, uri: String) {
93        self.inner.borrow_mut().set_user_dictionary(&uri);
94    }
95
96    /// Sets whether to keep whitespace in tokenization results.
97    ///
98    /// # Arguments
99    ///
100    /// * `keep_whitespace` - If true, whitespace tokens will be included.
101    fn set_keep_whitespace(&self, keep_whitespace: bool) {
102        self.inner.borrow_mut().set_keep_whitespace(keep_whitespace);
103    }
104
105    /// Appends a character filter to the filter pipeline.
106    ///
107    /// # Arguments
108    ///
109    /// * `kind` - Type of character filter to add.
110    /// * `args` - Optional hash of filter arguments.
111    fn append_character_filter(&self, kind: String, args: Option<RHash>) -> Result<(), Error> {
112        let ruby = Ruby::get().expect("Ruby runtime not initialized");
113        let filter_args = if let Some(hash) = args {
114            rb_hash_to_json(&ruby, hash)?
115        } else {
116            serde_json::Value::Object(serde_json::Map::new())
117        };
118        self.inner
119            .borrow_mut()
120            .append_character_filter(&kind, &filter_args);
121        Ok(())
122    }
123
124    /// Appends a token filter to the filter pipeline.
125    ///
126    /// # Arguments
127    ///
128    /// * `kind` - Type of token filter to add.
129    /// * `args` - Optional hash of filter arguments.
130    fn append_token_filter(&self, kind: String, args: Option<RHash>) -> Result<(), Error> {
131        let ruby = Ruby::get().expect("Ruby runtime not initialized");
132        let filter_args = if let Some(hash) = args {
133            rb_hash_to_json(&ruby, hash)?
134        } else {
135            serde_json::Value::Object(serde_json::Map::new())
136        };
137        self.inner
138            .borrow_mut()
139            .append_token_filter(&kind, &filter_args);
140        Ok(())
141    }
142
143    /// Builds the tokenizer with the configured settings.
144    ///
145    /// # Returns
146    ///
147    /// A configured `RbTokenizer` instance.
148    fn build(&self) -> Result<RbTokenizer, Error> {
149        let ruby = Ruby::get().expect("Ruby runtime not initialized");
150        let inner = self
151            .inner
152            .borrow()
153            .build()
154            .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
155        Ok(RbTokenizer { inner })
156    }
157}
158
159/// Tokenizer for performing morphological analysis.
160///
161/// The tokenizer processes text and returns tokens with their morphological features.
162#[magnus::wrap(class = "Lindera::Tokenizer", free_immediately, size)]
163pub struct RbTokenizer {
164    /// Inner binding-core tokenizer.
165    inner: CoreTokenizer,
166}
167
168/// Creates a new tokenizer with the given dictionary and mode.
169///
170/// # Arguments
171///
172/// * `dictionary` - Dictionary to use.
173/// * `mode` - Tokenization mode ("normal" or "decompose"). Defaults to "normal".
174/// * `user_dictionary` - Optional user dictionary.
175///
176/// # Returns
177///
178/// A new `RbTokenizer` instance.
179fn tokenizer_new(
180    dictionary: &RbDictionary,
181    mode: Option<String>,
182    user_dictionary: Option<&RbUserDictionary>,
183) -> Result<RbTokenizer, Error> {
184    let ruby = Ruby::get().expect("Ruby runtime not initialized");
185    let mode_str = mode.as_deref().unwrap_or("normal");
186
187    let dict = dictionary.inner.clone();
188    let user_dict = user_dictionary.map(|d| d.inner.clone());
189
190    let inner = CoreTokenizer::from_segmenter(mode_str, dict, user_dict)
191        .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
192
193    Ok(RbTokenizer { inner })
194}
195
196impl RbTokenizer {
197    /// Tokenizes the given text.
198    ///
199    /// # Arguments
200    ///
201    /// * `text` - Text to tokenize.
202    ///
203    /// # Returns
204    ///
205    /// An array of Token objects.
206    fn tokenize(&self, text: String) -> Result<RArray, Error> {
207        let ruby = Ruby::get().expect("Ruby runtime not initialized");
208        let views = self
209            .inner
210            .tokenize(&text)
211            .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
212
213        let arr = ruby.ary_new_capa(views.len());
214        for view in views {
215            arr.push(ruby.into_value(RbToken::from_view(view)))?;
216        }
217        Ok(arr)
218    }
219
220    /// Tokenizes the given text and returns N-best results.
221    ///
222    /// # Arguments
223    ///
224    /// * `text` - Text to tokenize.
225    /// * `n` - Number of N-best results.
226    /// * `unique` - Whether to deduplicate results (default: false).
227    /// * `cost_threshold` - Optional cost threshold.
228    ///
229    /// # Returns
230    ///
231    /// An array of [tokens, cost] pairs.
232    fn tokenize_nbest(
233        &self,
234        text: String,
235        n: usize,
236        unique: Option<bool>,
237        cost_threshold: Option<i64>,
238    ) -> Result<RArray, Error> {
239        let ruby = Ruby::get().expect("Ruby runtime not initialized");
240        let results = self
241            .inner
242            .tokenize_nbest(&text, n, unique.unwrap_or(false), cost_threshold)
243            .map_err(|err| to_magnus_error(&ruby, err.to_string()))?;
244
245        let rb_results = ruby.ary_new_capa(results.len());
246        for (views, cost) in results {
247            let token_arr = ruby.ary_new_capa(views.len());
248            for view in views {
249                token_arr.push(ruby.into_value(RbToken::from_view(view)))?;
250            }
251            let pair = ruby.ary_new_capa(2);
252            pair.push(token_arr)?;
253            pair.push(cost)?;
254            rb_results.push(pair)?;
255        }
256
257        Ok(rb_results)
258    }
259}
260
261/// Defines TokenizerBuilder and Tokenizer classes in the given Ruby module.
262///
263/// # Arguments
264///
265/// * `ruby` - Ruby runtime handle.
266/// * `module` - Parent Ruby module.
267///
268/// # Returns
269///
270/// `Ok(())` on success, or a Magnus `Error` on failure.
271pub fn define(ruby: &Ruby, module: &magnus::RModule) -> Result<(), Error> {
272    let builder_class = module.define_class("TokenizerBuilder", ruby.class_object())?;
273    builder_class.define_singleton_method("new", function!(RbTokenizerBuilder::new, 0))?;
274    builder_class
275        .define_singleton_method("from_file", function!(RbTokenizerBuilder::from_file, 1))?;
276    builder_class.define_method("set_mode", method!(RbTokenizerBuilder::set_mode, 1))?;
277    builder_class.define_method(
278        "set_dictionary",
279        method!(RbTokenizerBuilder::set_dictionary, 1),
280    )?;
281    builder_class.define_method(
282        "set_user_dictionary",
283        method!(RbTokenizerBuilder::set_user_dictionary, 1),
284    )?;
285    builder_class.define_method(
286        "set_keep_whitespace",
287        method!(RbTokenizerBuilder::set_keep_whitespace, 1),
288    )?;
289    builder_class.define_method(
290        "append_character_filter",
291        method!(RbTokenizerBuilder::append_character_filter, 2),
292    )?;
293    builder_class.define_method(
294        "append_token_filter",
295        method!(RbTokenizerBuilder::append_token_filter, 2),
296    )?;
297    builder_class.define_method("build", method!(RbTokenizerBuilder::build, 0))?;
298
299    let tokenizer_class = module.define_class("Tokenizer", ruby.class_object())?;
300    tokenizer_class.define_singleton_method("new", function!(tokenizer_new, 3))?;
301    tokenizer_class.define_method("tokenize", method!(RbTokenizer::tokenize, 1))?;
302    tokenizer_class.define_method("tokenize_nbest", method!(RbTokenizer::tokenize_nbest, 4))?;
303
304    Ok(())
305}