Newer
Older
TelosDB / src-tauri / src / utils / tokenizer.rs
use anyhow::{Result, anyhow};
use vibrato::Dictionary;
use vibrato::Tokenizer;
use std::io::Read;

pub struct JapaneseTokenizer {
    tokenizer: Tokenizer,
}

impl JapaneseTokenizer {
    pub fn new() -> Result<Self> {
        // src-tauri/src/utils/tokenizer.rs から resources/ は ../../resources/
        // パスが見つからないエラーが出るため、絶対パスではなくソースコードの位置からの相対を再確認
        let dict_data = include_bytes!("../../resources/ipadic.vibrato.zst");
        
        // 辞書の読み込み
        // 8MB 程度あるなら解凍済みの可能性が高い。直接読み込みを試みる。
        let dict = match Dictionary::read(&dict_data[..]) {
            Ok(d) => d,
            Err(_) => {
                // 直接読み込みに失敗した場合は zstd 解凍を試みる
                let mut decoder = zstd::stream::read::Decoder::new(&dict_data[..])?;
                let mut decoded_data = Vec::new();
                decoder.read_to_end(&mut decoded_data)?;
                Dictionary::read(&decoded_data[..])
                    .map_err(|e| anyhow!("Failed to read Vibrato dictionary: {}", e))?
            }
        };
        
        let tokenizer = Tokenizer::new(dict);
        
        Ok(JapaneseTokenizer { tokenizer })
    }

    /// テキストを形態素解析し、わかち書き(スペース区切り)の文字列として返す
    pub fn tokenize_to_string(&self, text: &str) -> Result<String> {
        let mut worker = self.tokenizer.new_worker();
        worker.reset_sentence(text);
        worker.tokenize();
        
        let result: Vec<&str> = worker
            .token_iter()
            .map(|token| token.surface())
            .collect();
        
        Ok(result.join(" "))
    }

    /// 単語のリスト(ベクタ)として返す
    pub fn tokenize_to_vec(&self, text: &str) -> Result<Vec<String>> {
        let mut worker = self.tokenizer.new_worker();
        worker.reset_sentence(text);
        worker.tokenize();
        
        Ok(worker.token_iter().map(|t| t.surface().to_string()).collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_japanese_tokenization() {
        let tokenizer = JapaneseTokenizer::new().expect("Failed to create tokenizer");
        let text = "すもももももももものうち";
        let tokenized = tokenizer.tokenize_to_string(text).unwrap();
        println!("Tokenized: {}", tokenized);
        assert!(tokenized.contains("すもも"));
        assert!(tokenized.contains("もも"));
    }
}