|
1 | 1 | # SPDX-FileCopyrightText: 2016-2026 PyThaiNLP Project |
2 | 2 | # SPDX-FileType: SOURCE |
3 | 3 | # SPDX-License-Identifier: Apache-2.0 |
4 | | -"""Wrapper for deepcut Thai word segmentation. deepcut is a |
5 | | -Thai word segmentation library using 1D Convolution Neural Network. |
| 4 | +"""DeepCut Thai word segmentation using ONNX runtime. |
6 | 5 |
|
7 | | -User need to install deepcut (and its dependency: tensorflow) by themselves. |
| 6 | +DeepCut is a Thai word segmentation library using 1D Convolution Neural |
| 7 | +Network. This module provides ONNX-based inference, removing the need for |
| 8 | +TensorFlow. |
| 9 | +
|
| 10 | +The ONNX model is ported from the original DeepCut TensorFlow model, |
| 11 | +available from the LEKCut project. |
8 | 12 |
|
9 | 13 | :See Also: |
10 | | - * `GitHub repository <https://github.com/rkcosmos/deepcut>`_ |
| 14 | + * `DeepCut GitHub <https://github.com/rkcosmos/deepcut>`_ |
| 15 | + * `LEKCut GitHub <https://github.com/PyThaiNLP/LEKCut>`_ |
| 16 | +
|
| 17 | +:References: |
| 18 | + Rakpong Kittinaradorn, Titipat Achakulvisut, Korakot Chaovavanich, |
| 19 | + Kittinan Srithaworn, Pattarawat Chormai, Chanwit Kaewkasi, |
| 20 | + Tulakan Ruangrong, Krichkorn Oparad. |
| 21 | + (2019, September 23). DeepCut: A Thai word tokenization library using |
| 22 | + Deep Neural Network. Zenodo. https://doi.org/10.5281/zenodo.3457707 |
11 | 23 | """ |
12 | 24 |
|
13 | 25 | from __future__ import annotations |
14 | 26 |
|
15 | | -from typing import Union, cast |
| 27 | +from typing import TYPE_CHECKING, Optional |
| 28 | + |
| 29 | +import numpy as np |
| 30 | +from onnxruntime import InferenceSession |
| 31 | + |
| 32 | +from pythainlp.corpus import get_corpus_path |
| 33 | + |
| 34 | +if TYPE_CHECKING: |
| 35 | + from numpy.typing import NDArray |
| 36 | + |
| 37 | +_MODEL_NAME: str = "deepcut_onnx" |
| 38 | +_N_PAD: int = 21 |
| 39 | +_THRESHOLD: float = 0.5 |
| 40 | + |
| 41 | +# Character type mapping from the original DeepCut model |
| 42 | +_CHAR_TYPE: dict[str, str] = { |
| 43 | + "กขฃคฆงจชซญฎฏฐฑฒณดตถทธนบปพฟภมยรลวศษสฬอ": "c", |
| 44 | + "ฅฉผฟฌหฮ": "n", |
| 45 | + "ะาำิีืึุู": "v", |
| 46 | + "เแโใไ": "w", |
| 47 | + "่้๊๋": "t", |
| 48 | + "์ๆฯ.": "s", |
| 49 | + "0123456789๑๒๓๔๕๖๗๘๙": "d", |
| 50 | + '"': "q", |
| 51 | + "'": "q", |
| 52 | + "\u2018": "q", |
| 53 | + "\u2019": "q", |
| 54 | + " ": "p", |
| 55 | + "abcdefghijklmnopqrstuvwxyz": "s_e", |
| 56 | + "ABCDEFGHIJKLMNOPQRSTUVWXYZ": "b_e", |
| 57 | +} |
| 58 | + |
| 59 | +_CHAR_TYPE_FLAT: dict[str, str] = {} |
| 60 | +for _ks, _ct in _CHAR_TYPE.items(): |
| 61 | + for _k in _ks: |
| 62 | + _CHAR_TYPE_FLAT[_k] = _ct |
| 63 | + |
| 64 | +_CHARS: list[str] = [ |
| 65 | + "\n", " ", "!", '"', "#", "$", "%", "&", "'", "(", ")", "*", "+", |
| 66 | + ",", "-", ".", "/", "0", "1", "2", "3", "4", "5", "6", "7", "8", |
| 67 | + "9", ":", ";", "<", "=", ">", "?", "@", "A", "B", "C", "D", "E", |
| 68 | + "F", "G", "H", "I", "J", "K", "L", "M", "N", "O", "P", "Q", "R", |
| 69 | + "S", "T", "U", "V", "W", "X", "Y", "Z", "[", "\\", "]", "^", "_", |
| 70 | + "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", |
| 71 | + "n", "o", "other", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", |
| 72 | + "z", "}", "~", "ก", "ข", "ฃ", "ค", "ฅ", "ฆ", "ง", "จ", "ฉ", "ช", |
| 73 | + "ซ", "ฌ", "ญ", "ฎ", "ฏ", "ฐ", "ฑ", "ฒ", "ณ", "ด", "ต", "ถ", "ท", |
| 74 | + "ธ", "น", "บ", "ป", "ผ", "ฝ", "พ", "ฟ", "ภ", "ม", "ย", "ร", "ฤ", |
| 75 | + "ล", "ว", "ศ", "ษ", "ส", "ห", "ฬ", "อ", "ฮ", "ฯ", "ะ", "ั", "า", |
| 76 | + "ำ", "ิ", "ี", "ึ", "ื", "ุ", "ู", "ฺ", "เ", "แ", "โ", "ใ", "ไ", |
| 77 | + "ๅ", "ๆ", "็", "่", "้", "๊", "๋", "์", "ํ", "๐", "๑", "๒", "๓", |
| 78 | + "๔", "๕", "๖", "๗", "๘", "๙", "\u2018", "\u2019", "\ufeff", |
| 79 | +] |
| 80 | +_CHARS_MAP: dict[str, int] = {v: k for k, v in enumerate(_CHARS)} |
| 81 | + |
| 82 | +_CHAR_TYPES: list[str] = [ |
| 83 | + "b_e", "c", "d", "n", "o", "p", "q", "s", "s_e", "t", "v", "w", |
| 84 | +] |
| 85 | +_CHAR_TYPES_MAP: dict[str, int] = {v: k for k, v in enumerate(_CHAR_TYPES)} |
| 86 | + |
| 87 | +# Default index for unknown characters and types |
| 88 | +_OTHER_CHAR_INDEX: int = _CHARS_MAP.get("other", 80) |
| 89 | +_OTHER_TYPE_INDEX: int = _CHAR_TYPES_MAP.get("o", 4) |
16 | 90 |
|
17 | | -try: |
18 | | - from deepcut import tokenize |
19 | | -except ImportError as e: |
20 | | - raise ImportError( |
21 | | - "deepcut is not installed. Install it with: pip install deepcut" |
22 | | - ) from e |
23 | | -from pythainlp.util import Trie |
| 91 | +_session: Optional[InferenceSession] = None |
| 92 | + |
| 93 | + |
| 94 | +def _get_session() -> InferenceSession: |
| 95 | + """Return a cached ONNX inference session, loading it on first call.""" |
| 96 | + global _session |
| 97 | + if _session is None: |
| 98 | + model_path = get_corpus_path(_MODEL_NAME) |
| 99 | + if not model_path: |
| 100 | + raise FileNotFoundError( |
| 101 | + f"corpus-not-found name={_MODEL_NAME!r}\n" |
| 102 | + " DeepCut ONNX model file not found in the package.\n" |
| 103 | + " Try reinstalling PyThaiNLP:\n" |
| 104 | + " pip install --force-reinstall pythainlp" |
| 105 | + ) |
| 106 | + _session = InferenceSession(model_path) |
| 107 | + return _session |
| 108 | + |
| 109 | + |
| 110 | +def _create_feature_array( |
| 111 | + text: str, n_pad: int = _N_PAD |
| 112 | +) -> tuple["NDArray[np.float32]", "NDArray[np.float32]"]: |
| 113 | + """Create character and type feature arrays for ONNX model input. |
| 114 | +
|
| 115 | + :param str text: input text |
| 116 | + :param int n_pad: window size for padding (default: 21) |
| 117 | + :return: character and type feature arrays of shape (n, n_pad) |
| 118 | + :rtype: tuple[numpy.ndarray, numpy.ndarray] |
| 119 | + """ |
| 120 | + n = len(text) |
| 121 | + n_pad_2 = (n_pad - 1) // 2 |
| 122 | + text_pad = [" "] * n_pad_2 + list(text) + [" "] * n_pad_2 |
| 123 | + x_char: list[list[int]] = [] |
| 124 | + x_type: list[list[int]] = [] |
| 125 | + for i in range(n_pad_2, n_pad_2 + n): |
| 126 | + char_list = ( |
| 127 | + text_pad[i + 1 : i + n_pad_2 + 1] |
| 128 | + + list(reversed(text_pad[i - n_pad_2 : i])) |
| 129 | + + [text_pad[i]] |
| 130 | + ) |
| 131 | + x_char.append([_CHARS_MAP.get(c, _OTHER_CHAR_INDEX) for c in char_list]) |
| 132 | + x_type.append( |
| 133 | + [ |
| 134 | + _CHAR_TYPES_MAP.get(_CHAR_TYPE_FLAT.get(c, "o"), _OTHER_TYPE_INDEX) |
| 135 | + for c in char_list |
| 136 | + ] |
| 137 | + ) |
| 138 | + return ( |
| 139 | + np.array(x_char, dtype=np.float32), |
| 140 | + np.array(x_type, dtype=np.float32), |
| 141 | + ) |
24 | 142 |
|
25 | 143 |
|
26 | 144 | def segment( |
27 | | - text: str, custom_dict: Union[Trie, list[str], str] = [] |
| 145 | + text: str, |
28 | 146 | ) -> list[str]: |
| 147 | + """Segment Thai text using the DeepCut ONNX model. |
| 148 | +
|
| 149 | + :param str text: text to segment |
| 150 | + :return: list of word tokens |
| 151 | + :rtype: list[str] |
| 152 | +
|
| 153 | + :Example: |
| 154 | + :: |
| 155 | +
|
| 156 | + from pythainlp.tokenize import deepcut |
| 157 | +
|
| 158 | + deepcut.segment("ทดสอบตัดคำ") |
| 159 | + # output: ['ทดสอบ', 'ตัด', 'คำ'] |
| 160 | + """ |
29 | 161 | if not text or not isinstance(text, str): |
30 | 162 | return [] |
31 | 163 |
|
32 | | - if custom_dict: |
33 | | - if isinstance(custom_dict, Trie): |
34 | | - custom_dict = list(custom_dict) |
35 | | - |
36 | | - return cast("list[str]", tokenize(text, custom_dict)) |
| 164 | + session = _get_session() |
| 165 | + x_char, x_type = _create_feature_array(text) |
| 166 | + outputs = session.run(None, {"input_1": x_char, "input_2": x_type}) |
| 167 | + y_predict = (outputs[0].ravel() > _THRESHOLD).astype(int) |
| 168 | + word_end = y_predict[1:].tolist() + [1] |
37 | 169 |
|
38 | | - return cast("list[str]", tokenize(text)) |
| 170 | + tokens: list[str] = [] |
| 171 | + word = "" |
| 172 | + for char, is_end in zip(text, word_end): |
| 173 | + word += char |
| 174 | + if is_end: |
| 175 | + tokens.append(word) |
| 176 | + word = "" |
| 177 | + return tokens |
0 commit comments