diff --git a/pythainlp/ancient/currency.py b/pythainlp/ancient/currency.py index 9159fc385..474549735 100644 --- a/pythainlp/ancient/currency.py +++ b/pythainlp/ancient/currency.py @@ -4,7 +4,7 @@ from __future__ import annotations -def convert_currency(value: float, from_unit: str) -> dict: +def convert_currency(value: float, from_unit: str) -> dict[str, float]: """Convert ancient Thai currency to other units * เบี้ย (Bia) @@ -23,7 +23,7 @@ def convert_currency(value: float, from_unit: str) -> dict: :param str from_unit: currency unit \ ('เบี้ย', 'อัฐ', 'ไพ', 'เฟื้อง', 'สลึง', 'บาท', 'ตำลึง', 'ชั่ง') :return: Thai currency - :rtype: dict + :rtype: dict[str, float] :Example: :: @@ -63,8 +63,8 @@ def convert_currency(value: float, from_unit: str) -> dict: # start from 'อัฐ' value_in_att = value * conversion_factors_to_att[from_unit] - # Calculate values ​​in other units - results = {} + # Calculate values in other units + results: dict[str, float] = {} for unit, factor in conversion_factors_to_att.items(): results[unit] = value_in_att / factor diff --git a/pythainlp/augment/word2vec/bpemb_wv.py b/pythainlp/augment/word2vec/bpemb_wv.py index 313f21a09..e461c022c 100644 --- a/pythainlp/augment/word2vec/bpemb_wv.py +++ b/pythainlp/augment/word2vec/bpemb_wv.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from pythainlp.augment.word2vec.core import Word2VecAug @@ -40,7 +40,7 @@ def tokenizer(self, text: str) -> list[str]: """:param str text: Thai text :rtype: List[str] """ - return self.bpemb_temp.encode(text) # type: ignore[no-any-return] + return cast(list[str], self.bpemb_temp.encode(text)) def load_w2v(self) -> None: """Load BPEmb model""" diff --git a/pythainlp/augment/word2vec/core.py b/pythainlp/augment/word2vec/core.py index b604a4d59..2dbf25b64 100644 --- a/pythainlp/augment/word2vec/core.py +++ b/pythainlp/augment/word2vec/core.py @@ -42,7 +42,9 @@ def __init__( _gensim_kv_logger.addFilter(_filter) try: if type == "file": - self.model = word2vec.KeyedVectors.load_word2vec_format(model) + self.model = word2vec.KeyedVectors.load_word2vec_format( + model + ) else: self.model = word2vec.KeyedVectors.load_word2vec_format( model, binary=True, unicode_errors="ignore" diff --git a/pythainlp/augment/wordnet.py b/pythainlp/augment/wordnet.py index 454d707a0..a5078bd3d 100644 --- a/pythainlp/augment/wordnet.py +++ b/pythainlp/augment/wordnet.py @@ -122,13 +122,13 @@ class WordNetAug: """Text Augment using wordnet""" synonyms: list[str] - list_synsets: list + list_synsets: list[Synset] p2w_pos: Optional[str] synset: Synset syn: str synonyms_without_duplicates: list[str] list_words: list[str] - list_synonym: list + list_synonym: list[list[str]] p_all: int list_pos: list[tuple[str, str]] temp: list[str] @@ -152,15 +152,15 @@ def find_synonyms( """ self.synonyms: list[str] = [] if pos is None: - self.list_synsets: list = wordnet.synsets(word) + self.list_synsets: list[Synset] = wordnet.synsets(word) else: self.p2w_pos: Optional[str] = postype2wordnet(pos, postag_corpus) if self.p2w_pos != "": - self.list_synsets: list = wordnet.synsets( + self.list_synsets: list[Synset] = wordnet.synsets( word, pos=self.p2w_pos ) else: - self.list_synsets: list = wordnet.synsets(word) + self.list_synsets: list[Synset] = wordnet.synsets(word) for self.synset in wordnet.synsets(word): for self.syn in self.synset.lemma_names(lang="tha"): @@ -206,7 +206,7 @@ def augment( """ new_sentences = [] self.list_words: list[str] = tokenize(sentence) - self.list_synonym: list = [] + self.list_synonym: list[list[str]] = [] self.p_all: int = 1 if postag: self.list_pos: list[tuple[str, str]] = pos_tag( diff --git a/pythainlp/benchmarks/word_tokenization.py b/pythainlp/benchmarks/word_tokenization.py index f74e1b1f1..01b67fee4 100644 --- a/pythainlp/benchmarks/word_tokenization.py +++ b/pythainlp/benchmarks/word_tokenization.py @@ -10,6 +10,7 @@ if TYPE_CHECKING: import numpy as np import pandas as pd + from numpy.typing import NDArray SEPARATOR: str = "|" @@ -43,7 +44,7 @@ def _f1(precision: float, recall: float) -> float: def _flatten_result( - my_dict: dict, sep: str = ":" + my_dict: dict[str, dict[str, Union[int, str]]], sep: str = ":" ) -> dict[str, Union[int, str]]: """Flatten two-dimension dictionary. @@ -54,7 +55,8 @@ def _flatten_result( { "a:b": 7 } - :param dict my_dict: dictionary containing stats + :param dict[str, dict[str, Union[int, str]]] my_dict: dictionary + containing stats :param str sep: separator between the two keys (default: ":") :return: a one-dimension dictionary with keys combined @@ -91,7 +93,7 @@ def benchmark(ref_samples: list[str], samples: list[str]) -> "pd.DataFrame": flat_stats["expected"] = r flat_stats["actual"] = s results.append(flat_stats) - except Exception: + except Exception as exc: reason = """ [Error] Reason: %s @@ -107,7 +109,7 @@ def benchmark(ref_samples: list[str], samples: list[str]) -> "pd.DataFrame": r, s, ) - raise SystemExit(reason) + raise SystemExit(reason) from exc return pd.DataFrame(results) @@ -209,7 +211,9 @@ def compute_stats( } -def _binary_representation(txt: str, verbose: bool = False) -> "np.ndarray": +def _binary_representation( + txt: str, verbose: bool = False +) -> "NDArray[np.int8]": """Transform text into {0, 1} sequence. where (1) indicates that the corresponding character is the beginning of @@ -219,7 +223,7 @@ def _binary_representation(txt: str, verbose: bool = False) -> "np.ndarray": :param bool verbose: for debugging purposes :return: {0, 1} sequence - :rtype: np.ndarray + :rtype: numpy.typing.NDArray[numpy.int8] """ import numpy as np @@ -228,7 +232,7 @@ def _binary_representation(txt: str, verbose: bool = False) -> "np.ndarray": boundary = np.argwhere(chars == SEPARATOR).reshape(-1) boundary = boundary - np.array(range(boundary.shape[0])) - bin_rept: np.ndarray = np.zeros(len(txt) - boundary.shape[0]) + bin_rept = np.zeros(len(txt) - boundary.shape[0], dtype=np.int8) bin_rept[list(boundary) + [0]] = 1 sample_wo_seps = list(txt.replace(SEPARATOR, "")) @@ -247,10 +251,13 @@ def _binary_representation(txt: str, verbose: bool = False) -> "np.ndarray": return bin_rept -def _find_word_boundaries(bin_reps: "np.ndarray") -> list[tuple[int, int]]: +def _find_word_boundaries( + bin_reps: "NDArray[np.int8]", +) -> list[tuple[int, int]]: """Find the starting and ending location of each word. - :param numpy.ndarray bin_reps: binary representation of a text + :param numpy.typing.NDArray[numpy.int8] bin_reps: binary representation + of a text :return: list of tuples (start, end) :rtype: list[tuple[int, int]] diff --git a/pythainlp/chunk/__init__.py b/pythainlp/chunk/__init__.py index 82c035e5b..f0f97dbb3 100644 --- a/pythainlp/chunk/__init__.py +++ b/pythainlp/chunk/__init__.py @@ -25,6 +25,7 @@ print(parser.parse(tokens_pos)) # output: ['B-NP', 'B-VP', 'I-VP'] """ + from __future__ import annotations __all__: list[str] = [ diff --git a/pythainlp/chunk/crfchunk.py b/pythainlp/chunk/crfchunk.py index 4a354d0cc..332aaa311 100644 --- a/pythainlp/chunk/crfchunk.py +++ b/pythainlp/chunk/crfchunk.py @@ -2,10 +2,11 @@ # SPDX-FileType: SOURCE # SPDX-License-Identifier: Apache-2.0 """CRF-based Thai phrase structure (chunk) parser.""" + from __future__ import annotations from importlib.resources import as_file, files -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union, cast if TYPE_CHECKING: import types @@ -130,7 +131,7 @@ def parse(self, token_pos: list[tuple[str, str]]) -> list[str]: :rtype: list[str] """ self.xseq = _extract_features(token_pos) - return self.tagger.tag(self.xseq) # type: ignore[no-any-return] + return cast(list[str], self.tagger.tag(self.xseq)) def __enter__(self) -> CRFChunkParser: """Context manager entry.""" diff --git a/pythainlp/coref/core.py b/pythainlp/coref/core.py index a197a8ea7..94459e453 100644 --- a/pythainlp/coref/core.py +++ b/pythainlp/coref/core.py @@ -3,16 +3,16 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import Any, Optional, Union +from typing import Any, Union, cast -_MODEL: Optional[Any] = None +_MODEL_CACHE: dict[tuple[str, str], Any] = {} def coreference_resolution( texts: Union[str, list[str]], model_name: str = "han-coref-v1.0", device: str = "cpu", -) -> list[dict]: +) -> list[dict[str, Any]]: """Coreference Resolution :param Union[str, list[str]] texts: list of texts to apply coreference resolution to @@ -20,7 +20,7 @@ def coreference_resolution( :param str device: device for running coreference resolution model on\ ("cpu", "cuda", and others) :return: List of texts with coreference resolution - :rtype: list[dict] + :rtype: list[dict[str, Any]] :Options for model_name: * *han-coref-v1.0* - (default) Han-Coref: Thai coreference resolution\ @@ -43,17 +43,18 @@ def coreference_resolution( # 'clusters': [[(0, 10), (50, 52)]]} # ] """ - global _MODEL if isinstance(texts, str): texts = [texts] - if _MODEL is None and model_name == "han-coref-v1.0": + model_key = (model_name, device) + if model_key not in _MODEL_CACHE and model_name == "han-coref-v1.0": from pythainlp.coref.han_coref import HanCoref - _MODEL = HanCoref(device=device) + _MODEL_CACHE[model_key] = HanCoref(device=device) - if _MODEL: - return _MODEL.predict(texts) # type: ignore[no-any-return] + model = _MODEL_CACHE.get(model_key) + if model is not None: + return cast(list[dict[str, Any]], model.predict(texts)) return [ {"text": text, "clusters_string": [], "clusters": []} for text in texts diff --git a/pythainlp/corpus/common.py b/pythainlp/corpus/common.py index 6bdfec631..d177cba3f 100644 --- a/pythainlp/corpus/common.py +++ b/pythainlp/corpus/common.py @@ -407,9 +407,12 @@ def thai_synonyms() -> dict[str, Union[list[str], list[list[str]]]]: pos = row.get("pos") synonym = row.get("synonym") if not ( - word and word.strip() - and pos and pos.strip() - and synonym and synonym.strip() + word + and word.strip() + and pos + and pos.strip() + and synonym + and synonym.strip() ): warnings.warn( f"Skipping thai_synonyms entry with missing or empty field(s): {dict(row)!r}", @@ -420,7 +423,11 @@ def thai_synonyms() -> dict[str, Union[list[str], list[list[str]]]]: words.append(word) pos_tags.append(pos) synonym_groups.append(synonym.split("|")) - _THAI_SYNONYMS = {"word": words, "pos": pos_tags, "synonym": synonym_groups} + _THAI_SYNONYMS = { + "word": words, + "pos": pos_tags, + "synonym": synonym_groups, + } return _THAI_SYNONYMS diff --git a/pythainlp/corpus/core.py b/pythainlp/corpus/core.py index c2964ea57..058408a3d 100644 --- a/pythainlp/corpus/core.py +++ b/pythainlp/corpus/core.py @@ -13,7 +13,7 @@ import zipfile from functools import lru_cache from importlib.resources import files -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any, Optional, cast from pythainlp import __version__ from pythainlp.corpus import corpus_db_path, corpus_db_url, corpus_path @@ -26,7 +26,6 @@ if TYPE_CHECKING: from http.client import HTTPMessage, HTTPResponse - from typing import Any, Optional _USER_AGENT: str = ( f"PyThaiNLP/{__version__} " @@ -50,7 +49,9 @@ def __init__(self, response: HTTPResponse) -> None: def json(self) -> dict[str, Any]: """Parse JSON content from response.""" try: - return json.loads(self._content.decode("utf-8")) # type: ignore[no-any-return] + return cast( + dict[str, Any], json.loads(self._content.decode("utf-8")) + ) except (json.JSONDecodeError, UnicodeDecodeError) as err: raise ValueError(f"Failed to parse JSON response: {err}") from err @@ -98,16 +99,15 @@ def get_corpus_db_detail(name: str, version: str = "") -> dict[str, Any]: if not version: for corpus in local_db["_default"].values(): if corpus["name"] == name: - return corpus # type: ignore[no-any-return] + return cast(dict[str, Any], corpus) else: for corpus in local_db["_default"].values(): if corpus["name"] == name and corpus["version"] == version: - return corpus # type: ignore[no-any-return] + return cast(dict[str, Any], corpus) return {} - @lru_cache(maxsize=None) def get_corpus(filename: str, comments: bool = True) -> frozenset[str]: """Read corpus data from file and return a frozenset. @@ -221,7 +221,7 @@ def _load_default_db() -> dict[str, Any]: corpus_files = files("pythainlp.corpus") default_db_file = corpus_files.joinpath("default_db.json") text = default_db_file.read_text(encoding="utf-8-sig") - return json.loads(text) # type: ignore[no-any-return] + return cast(dict[str, Any], json.loads(text)) def get_corpus_default_db(name: str, version: str = "") -> Optional[str]: @@ -848,7 +848,6 @@ def remove(name: str) -> bool: return False - def make_safe_directory_name(name: str) -> str: """Make safe directory name @@ -921,4 +920,4 @@ def get_hf_hub(repo_id: str, filename: str = "") -> str: output_path = snapshot_download( repo_id=repo_id, local_dir=root_project ) - return output_path # type: ignore[no-any-return] + return str(output_path) diff --git a/pythainlp/corpus/wordnet.py b/pythainlp/corpus/wordnet.py index c7583987b..732704105 100644 --- a/pythainlp/corpus/wordnet.py +++ b/pythainlp/corpus/wordnet.py @@ -12,7 +12,7 @@ from __future__ import annotations -from typing import IO, TYPE_CHECKING, Optional, Union +from typing import IO, TYPE_CHECKING, Optional, Union, cast if TYPE_CHECKING: from collections.abc import Iterable @@ -30,11 +30,12 @@ nltk.download("wordnet") from nltk.corpus import wordnet +from nltk.corpus.reader.wordnet import Lemma, Synset def synsets( word: str, pos: Optional[str] = None, lang: str = "tha" -) -> list[wordnet.Synset]: +) -> list[Synset]: """This function returns the synonym set for all lemmas of the given word with an optional argument to constrain the part of speech of the word. @@ -78,10 +79,10 @@ def synsets( >>> synsets("แรง", pos="a", lang="tha") [Synset('hard.s.10'), Synset('strong.s.02')] """ - return wordnet.synsets(lemma=word, pos=pos, lang=lang) # type: ignore[no-any-return] + return cast(list[Synset], wordnet.synsets(lemma=word, pos=pos, lang=lang)) -def synset(name_synsets: str) -> wordnet.Synset: +def synset(name_synsets: str) -> Synset: """This function returns the synonym set (synset) given the name of the synset (i.e. 'dog.n.01', 'chase.v.01'). @@ -144,10 +145,10 @@ def all_lemma_names(pos: Optional[str] = None, lang: str = "tha") -> list[str]: >>> len(all_lemma_names(pos="a")) 5277 """ - return wordnet.all_lemma_names(pos=pos, lang=lang) # type: ignore[no-any-return] + return cast(list[str], wordnet.all_lemma_names(pos=pos, lang=lang)) -def all_synsets(pos: Optional[str] = None) -> Iterable[wordnet.Synset]: +def all_synsets(pos: Optional[str] = None) -> Iterable[Synset]: """This function iterates over all synsets constrained by the given part of speech tag. @@ -174,7 +175,7 @@ def all_synsets(pos: Optional[str] = None) -> Iterable[wordnet.Synset]: >>> next(generator) Synset('unable.a.01') """ - return wordnet.all_synsets(pos=pos) # type: ignore[no-any-return] + return cast(Iterable[Synset], wordnet.all_synsets(pos=pos)) def langs() -> list[str]: @@ -192,12 +193,12 @@ def langs() -> list[str]: 'pol', 'por', 'qcn', 'slv', 'spa', 'swe', 'tha', 'zsm'] """ - return wordnet.langs() # type: ignore[no-any-return] + return cast(list[str], wordnet.langs()) def lemmas( word: str, pos: Optional[str] = None, lang: str = "tha" -) -> list[wordnet.Lemma]: +) -> list[Lemma]: """This function returns all lemmas given the word with an optional argument to constrain the part of speech of the word. @@ -237,10 +238,10 @@ def lemmas( >>> lemmas("ม้วน", pos="n") [Lemma('roll.n.11.ม้วน')] """ - return wordnet.lemmas(word, pos=pos, lang=lang) # type: ignore[no-any-return] + return cast(list[Lemma], wordnet.lemmas(word, pos=pos, lang=lang)) -def lemma(name_synsets: str) -> wordnet.Lemma: +def lemma(name_synsets: str) -> Lemma: """This function returns lemma object given the name. .. note:: @@ -267,7 +268,7 @@ def lemma(name_synsets: str) -> wordnet.Lemma: return wordnet.lemma(name_synsets) -def lemma_from_key(key: str) -> wordnet.Lemma: +def lemma_from_key(key: str) -> Lemma: """This function returns lemma object given the lemma key. This is similar to :func:`lemma` but it needs to be given the key of lemma instead of the name of lemma. @@ -293,9 +294,7 @@ def lemma_from_key(key: str) -> wordnet.Lemma: return wordnet.lemma_from_key(key) -def path_similarity( - synsets1: wordnet.Synset, synsets2: wordnet.Synset -) -> float: +def path_similarity(synsets1: Synset, synsets2: Synset) -> float: """This function returns similarity between two synsets based on the shortest path distance calculated using the equation below. @@ -331,12 +330,10 @@ def path_similarity( >>> path_similarity(obj, cat) 0.08333333333333333 """ - return wordnet.path_similarity(synsets1, synsets2) # type: ignore[no-any-return] + return cast(float, wordnet.path_similarity(synsets1, synsets2)) -def lch_similarity( - synsets1: wordnet.Synset, synsets2: wordnet.Synset -) -> float: +def lch_similarity(synsets1: Synset, synsets2: Synset) -> float: """This function returns Leacock Chodorow similarity (LCH) between two synsets, based on the shortest path distance and the maximum depth of the taxonomy. The equation to @@ -370,12 +367,10 @@ def lch_similarity( >>> lch_similarity(obj, cat) 1.1526795099383855 """ - return wordnet.lch_similarity(synsets1, synsets2) # type: ignore[no-any-return] + return cast(float, wordnet.lch_similarity(synsets1, synsets2)) -def wup_similarity( - synsets1: wordnet.Synset, synsets2: wordnet.Synset -) -> float: +def wup_similarity(synsets1: Synset, synsets2: Synset) -> float: """This function returns Wu-Palmer similarity (WUP) between two synsets, based on the depth of the two senses in the taxonomy and their Least Common Subsumer (most specific ancestor node). @@ -403,7 +398,7 @@ def wup_similarity( >>> wup_similarity(obj, cat) 0.35294117647058826 """ - return wordnet.wup_similarity(synsets1, synsets2) # type: ignore[no-any-return] + return cast(float, wordnet.wup_similarity(synsets1, synsets2)) def morphy(form: str, pos: Optional[str] = None) -> str: @@ -422,7 +417,7 @@ def morphy(form: str, pos: Optional[str] = None) -> str: >>> from pythainlp.corpus.wordnet import morphy >>> >>> morphy("dogs") - 'dogs' + 'dog' >>> >>> morphy("thieves") 'thief' @@ -433,7 +428,7 @@ def morphy(form: str, pos: Optional[str] = None) -> str: >>> morphy("calculated") 'calculate' """ - return wordnet.morphy(form, pos=None) # type: ignore[no-any-return] + return cast(str, wordnet.morphy(form, pos=pos)) def custom_lemmas(tab_file: Union[str, IO[str]], lang: str) -> None: @@ -444,4 +439,4 @@ def custom_lemmas(tab_file: Union[str, IO[str]], lang: str) -> None: :param tab_file: Tab file as a file or file-like object :param str lang: abbreviation of language (i.e. *eng*, *tha*). """ - return wordnet.custom_lemmas(tab_file, lang) # type: ignore[no-any-return] + wordnet.custom_lemmas(tab_file, lang) diff --git a/pythainlp/el/_multiel.py b/pythainlp/el/_multiel.py index d55908f9b..2c78b404f 100644 --- a/pythainlp/el/_multiel.py +++ b/pythainlp/el/_multiel.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Any, Union, cast if TYPE_CHECKING: from multiel import BELA @@ -30,7 +30,10 @@ def load_model(self) -> None: def process_batch( self, list_text: Union[list[str], str] - ) -> Union[list[dict], str]: + ) -> Union[list[dict[str, Any]], str]: if isinstance(list_text, str): list_text = [list_text] - return self._bela_run.process_batch(list_text) # type: ignore[no-any-return] + return cast( + Union[list[dict[str, Any]], str], + self._bela_run.process_batch(list_text), + ) diff --git a/pythainlp/el/core.py b/pythainlp/el/core.py index cdf1d438c..2034190b5 100644 --- a/pythainlp/el/core.py +++ b/pythainlp/el/core.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import Union +from typing import Any, Union class EntityLinker: @@ -41,12 +41,12 @@ def __init__( def get_el( self, list_text: Union[list[str], str] - ) -> Union[list[dict], str]: + ) -> Union[list[dict[str, Any]], str]: """Get Entity Linking from Thai Text - :param str Union[list[str], str]: list of Thai text or text + :param Union[list[str], str] list_text: list of Thai text or text :return: list of entity linking - :rtype: Union[list[dict], str] + :rtype: Union[list[dict[str, Any]], str] :Example: :: diff --git a/pythainlp/generate/wangchanglm.py b/pythainlp/generate/wangchanglm.py index ab5531c1b..1a4cf8c3c 100644 --- a/pythainlp/generate/wangchanglm.py +++ b/pythainlp/generate/wangchanglm.py @@ -4,7 +4,7 @@ from __future__ import annotations import re -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, cast if TYPE_CHECKING: import pandas as pd @@ -13,7 +13,7 @@ class WangChanGLM: - exclude_pattern: "re.Pattern" + exclude_pattern: "re.Pattern[str]" stop_token: str PROMPT_DICT: dict[str, str] device: str @@ -25,7 +25,7 @@ class WangChanGLM: exclude_ids: list[int] def __init__(self) -> None: - self.exclude_pattern: "re.Pattern" = re.compile(r"[^ก-๙]+") + self.exclude_pattern: "re.Pattern[str]" = re.compile(r"[^ก-๙]+") self.stop_token: str = "\n" # noqa: S105 self.PROMPT_DICT: dict[str, str] = { "prompt_input": ( @@ -137,9 +137,12 @@ def gen_instruct( typical_p=typical_p, temperature=temperature, # 0.9 ) - return self.tokenizer.decode( # type: ignore[no-any-return] - output_tokens[0][len(batch["input_ids"][0]) :], - skip_special_tokens=skip_special_tokens, + return cast( + str, + self.tokenizer.decode( + output_tokens[0][len(batch["input_ids"][0]) :], + skip_special_tokens=skip_special_tokens, + ), ) def instruct_generate( diff --git a/pythainlp/parse/ud_goeswith.py b/pythainlp/parse/ud_goeswith.py index 1038fbe1f..3e28d7412 100644 --- a/pythainlp/parse/ud_goeswith.py +++ b/pythainlp/parse/ud_goeswith.py @@ -45,11 +45,15 @@ def __call__( e = self.model(input_ids=torch.tensor(x)).logits.numpy()[ :, 1:-2, : ] - r = [ + root_adjustments: list[int] = [ 1 if i == 0 else -1 if j.endswith("|root") else 0 for i, j in sorted(self.model.config.id2label.items()) ] - e += np.where(np.add.outer(np.identity(e.shape[0]), r) == 0, 0, np.nan) + e += np.where( + np.add.outer(np.identity(e.shape[0]), root_adjustments) == 0, + 0, + np.nan, + ) g = self.model.config.label2id["X|_|goeswith"] r = np.tri(e.shape[0]) for i in range(e.shape[0]): diff --git a/pythainlp/phayathaibert/core.py b/pythainlp/phayathaibert/core.py index 3bbbb3011..036ebee25 100644 --- a/pythainlp/phayathaibert/core.py +++ b/pythainlp/phayathaibert/core.py @@ -6,7 +6,7 @@ import random import re import warnings -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, Union, cast if TYPE_CHECKING: from collections.abc import Callable @@ -193,14 +193,14 @@ def remove_space(self, toks: list[str]) -> list[str]: def preprocess( self, text: str, - pre_rules: list[Callable] = [ + pre_rules: list[Callable[..., str]] = [ rm_brackets, replace_newlines, rm_useless_spaces, replace_spaces, replace_rep_after, ], - tok_func: Callable = word_tokenize, + tok_func: Callable[..., list[str]] = word_tokenize, ) -> str: text = text.lower() for rule in pre_rules: @@ -460,4 +460,4 @@ def segment(sentence: str) -> list[str]: if not sentence or not isinstance(sentence, str): return [] - return _tokenizer.tokenize(sentence) # type: ignore[no-any-return] + return cast(list[str], _tokenizer.tokenize(sentence)) diff --git a/pythainlp/soundex/complete_soundex.py b/pythainlp/soundex/complete_soundex.py index be5c196de..e85bff934 100644 --- a/pythainlp/soundex/complete_soundex.py +++ b/pythainlp/soundex/complete_soundex.py @@ -208,7 +208,9 @@ def heuristic_split(self, text: str) -> list[tuple[str, Optional[str]]]: return [(text, None)] - def _process_leading_vowel(self, chars: list, idx: int) -> tuple: + def _process_leading_vowel( + self, chars: list[str], idx: int + ) -> tuple[str, int]: """Extract and process leading vowel.""" leading_vowel = "" if idx < len(chars) and chars[idx] in ["เ", "แ", "โ", "ไ", "ใ"]: @@ -217,8 +219,8 @@ def _process_leading_vowel(self, chars: list, idx: int) -> tuple: return leading_vowel, idx def _process_initial_consonant( - self, chars: list, idx: int, leading_vowel: str - ) -> tuple: + self, chars: list[str], idx: int, leading_vowel: str + ) -> tuple[str, str, str, int]: """Process initial consonant and cluster.""" init_char = "" init_code = "" @@ -248,7 +250,7 @@ def _process_initial_consonant( return init_char, init_code, cluster_char, idx def _detect_cluster( - self, chars: list, idx: int, leading_vowel: str + self, chars: list[str], idx: int, leading_vowel: str ) -> bool: """Detect if ร/ล/ว is a cluster.""" if idx + 1 < len(chars): @@ -269,7 +271,7 @@ def _detect_cluster( return True return False - def _map_leading_vowel_code(self, leading_vowel: str) -> tuple: + def _map_leading_vowel_code(self, leading_vowel: str) -> tuple[str, str]: """Map leading vowel to initial vowel and final code.""" vowel_code = "" final_code = "-" @@ -292,12 +294,12 @@ def _map_leading_vowel_code(self, leading_vowel: str) -> tuple: def _scan_vowels_tones_finals( self, - chars: list, + chars: list[str], idx: int, leading_vowel: str, vowel_code: str, final_code: str, - ) -> tuple: + ) -> tuple[str, str, str, list[str]]: """Scan remaining characters for vowels, tones, and finals.""" remaining = chars[idx:] final_candidates = [] @@ -317,7 +319,7 @@ def _scan_vowels_tones_finals( def _process_vowel_char( self, c: str, leading_vowel: str, vowel_code: str, final_code: str - ) -> tuple: + ) -> tuple[str, str]: """Process a single vowel character.""" # Complex Vowel Checks if leading_vowel == "เ" and c == "ื": @@ -354,8 +356,8 @@ def _process_final_consonant( syl: str, final_code: str, vowel_code: str, - final_candidates: list, - ) -> tuple: + final_candidates: list[str], + ) -> tuple[str, str, bool]: """Process final consonant and detect dropped ร.""" dropped_r = False @@ -388,7 +390,10 @@ def _process_final_consonant( return vowel_code, final_code, dropped_r def _check_special_format( - self, init_char: str, final_candidates: list, vowel_code: str + self, + init_char: str, + final_candidates: list[str], + vowel_code: str, ) -> bool: """Check if special format (tone before final) should be used.""" # Special format (tone before final) is used when: diff --git a/pythainlp/soundex/sound.py b/pythainlp/soundex/sound.py index f00dfdd38..b8e15453d 100644 --- a/pythainlp/soundex/sound.py +++ b/pythainlp/soundex/sound.py @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +from typing import cast + import panphon import panphon.distance @@ -71,7 +73,10 @@ def audio_vector(word: str) -> list[list[int]]: audio_vector("น้ำ") # output : [[-1, 1, 1, -1, -1, -1, ...]] """ - return _ft.word_to_vector_list(word2audio(word), numeric=True) # type: ignore[no-any-return] + return cast( + list[list[int]], + _ft.word_to_vector_list(word2audio(word), numeric=True), + ) def word_approximation(word: str, list_word: list[str]) -> list[float]: diff --git a/pythainlp/spell/phunspell.py b/pythainlp/spell/phunspell.py index caada60cb..6b5ee6fe4 100644 --- a/pythainlp/spell/phunspell.py +++ b/pythainlp/spell/phunspell.py @@ -12,7 +12,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast if TYPE_CHECKING: import phunspell @@ -32,4 +32,4 @@ def spell(text: str) -> list[str]: def correct(text: str) -> str: - return list(pspell.suggest(text))[0] # type: ignore[no-any-return] + return cast(str, list(pspell.suggest(text))[0]) diff --git a/pythainlp/spell/tltk.py b/pythainlp/spell/tltk.py index 29a17bfef..9fce54771 100644 --- a/pythainlp/spell/tltk.py +++ b/pythainlp/spell/tltk.py @@ -12,6 +12,8 @@ from __future__ import annotations +from typing import cast + try: from tltk.nlp import spell_candidates except ImportError: @@ -21,4 +23,4 @@ def spell(text: str) -> list[str]: - return spell_candidates(text) # type: ignore[no-any-return] + return cast(list[str], spell_candidates(text)) diff --git a/pythainlp/spell/words_spelling_correction.py b/pythainlp/spell/words_spelling_correction.py index c5f20c346..7f06bb2ac 100644 --- a/pythainlp/spell/words_spelling_correction.py +++ b/pythainlp/spell/words_spelling_correction.py @@ -4,15 +4,16 @@ from __future__ import annotations import os -from typing import TYPE_CHECKING, Any, Optional, Union +from importlib import import_module +from typing import TYPE_CHECKING, Union, cast + +from pythainlp.corpus import get_hf_hub if TYPE_CHECKING: import numpy as np from numpy.typing import NDArray from onnxruntime import InferenceSession -from pythainlp.corpus import get_hf_hub - class FastTextEncoder: """A class to load pre-trained FastText-like word embeddings, @@ -58,15 +59,13 @@ def __init__( """ try: - import numpy as np # noqa: F401 - except ModuleNotFoundError: + import_module("numpy") + import_module("onnxruntime") + except ModuleNotFoundError as exc: raise ModuleNotFoundError( - """ - Please installing the package via 'pip install numpy onnxruntime'. - """ - ) - except Exception as e: - raise RuntimeError(f"An unexpected error occurred: {e}") from e + "Please install the required packages via " + "'pip install numpy onnxruntime'." + ) from exc self.model_dir = model_dir self.nn_model_path = nn_model_path self.bucket = bucket @@ -81,7 +80,11 @@ def __init__( self.embedding_dim = self.embeddings.shape[1] def _load_embeddings(self) -> tuple[list[str], NDArray[np.float32]]: - """Loads embeddings matrix and vocabulary list.""" + """Load the embeddings matrix and vocabulary list. + + :return: vocabulary entries and their float32 embedding matrix + :rtype: tuple[list[str], numpy.typing.NDArray[numpy.float32]] + """ import numpy as np input_matrix = np.load( @@ -97,7 +100,12 @@ def _load_embeddings(self) -> tuple[list[str], NDArray[np.float32]]: def _load_suggestion_words( self, words_list: list[str] ) -> NDArray[np.str_]: - """Loads the list of words used for suggestions.""" + """Load suggestion words into a NumPy string array. + + :param list[str] words_list: words used for nearest-neighbor lookup + :return: words as a NumPy string array + :rtype: numpy.typing.NDArray[numpy.str_] + """ import numpy as np words = np.array(words_list) @@ -125,7 +133,12 @@ def _get_hash(self, subword: str) -> int: return h % self.bucket + self.nb_words def _get_subwords(self, word: str) -> tuple[list[str], NDArray[np.int_]]: - """Extracts subwords and their corresponding indices for a given word.""" + """Extract subwords and their corresponding integer indices. + + :param str word: input word + :return: extracted subwords and their NumPy index array + :rtype: tuple[list[str], numpy.typing.NDArray[numpy.int_]] + """ _word = "<" + word + ">" _subwords = [] _subword_ids = [] @@ -156,7 +169,12 @@ def _get_subwords(self, word: str) -> tuple[list[str], NDArray[np.int_]]: return _subwords, np.array(_subword_ids) def get_word_vector(self, word: str) -> NDArray[np.float32]: - """Computes the normalized vector for a single word.""" + """Compute the normalized vector for a single word. + + :param str word: input word + :return: normalized float32 embedding vector + :rtype: numpy.typing.NDArray[numpy.float32] + """ import numpy as np # subword_ids[1] contains the array of indices for the word and its subwords @@ -165,17 +183,21 @@ def get_word_vector(self, word: str) -> NDArray[np.float32]: # Check if the array of subword indices is empty if subword_ids.size == 0: # Return a 300-dimensional zero vector if no word/subword is found. - return np.zeros(self.embedding_dim) + return np.zeros(self.embedding_dim, dtype=np.float32) # Compute the mean of the embeddings for all subword indices - vector = np.mean([self.embeddings[s] for s in subword_ids], axis=0) + vector = np.mean( + [self.embeddings[s] for s in subword_ids], + axis=0, + dtype=np.float32, + ) # Normalize the vector norm = np.linalg.norm(vector) if norm > 0: vector /= norm - return vector + return cast("NDArray[np.float32]", vector) def _tokenize(self, sentence: str) -> list[str]: """Tokenizes a sentence based on whitespace.""" @@ -195,7 +217,12 @@ def _tokenize(self, sentence: str) -> list[str]: return tokens def get_sentence_vector(self, line: str) -> NDArray[np.float32]: - """Computes the mean vector for a sentence.""" + """Compute the mean embedding vector for a sentence. + + :param str line: input sentence + :return: float32 sentence embedding vector + :rtype: numpy.typing.NDArray[numpy.float32] + """ import numpy as np tokens = self._tokenize(line) @@ -207,9 +234,12 @@ def get_sentence_vector(self, line: str) -> NDArray[np.float32]: # If the sentence was empty and resulted in no vectors, return a zero vector if not vectors: - return np.zeros(self.embedding_dim) + return np.zeros(self.embedding_dim, dtype=np.float32) - return np.mean(vectors, axis=0) + return cast( + "NDArray[np.float32]", + np.mean(vectors, axis=0, dtype=np.float32), + ) # --- Nearest Neighbor Method --- @@ -274,13 +304,14 @@ def __init__(self) -> None: with open( get_hf_hub( self.model_name, "list_word-spelling-correction-char2vec.txt" - ) + ), + encoding="utf-8", ) as f: self.list_word = list(map(str.strip, f.readlines())) super().__init__(self.model_path, self.model_onnx, self.list_word) -_WSC: Optional[Any] = None +_WSC_CACHE: dict[str, Words_Spelling_Correction] = {} def get_words_spell_suggestion( @@ -309,7 +340,6 @@ def get_words_spell_suggestion( # output: [['คนดีผีคุ้ม', 'มีดคอม้า', 'คดี', 'มีดสองคม', 'มูลคดี'], # ['กระเพาะ', 'กระพา', 'กะเพรา', 'กระเพาะปลา', 'พระประธาน']] """ - global _WSC - if _WSC is None: - _WSC = Words_Spelling_Correction() - return _WSC.get_word_suggestion(list_words) + if "default" not in _WSC_CACHE: + _WSC_CACHE["default"] = Words_Spelling_Correction() + return _WSC_CACHE["default"].get_word_suggestion(list_words) diff --git a/pythainlp/summarize/freq.py b/pythainlp/summarize/freq.py index c488e1da8..017193acc 100644 --- a/pythainlp/summarize/freq.py +++ b/pythainlp/summarize/freq.py @@ -5,7 +5,7 @@ from __future__ import annotations -from collections import defaultdict +from collections import Counter from heapq import nlargest from string import punctuation from typing import cast @@ -20,36 +20,32 @@ class FrequencySummarizer: __min_cut: float __max_cut: float __stopwords: set[str] - __freq: "defaultdict[str, float]" + __freq: dict[str, float] def __init__(self, min_cut: float = 0.1, max_cut: float = 0.9) -> None: self.__min_cut: float = min_cut self.__max_cut: float = max_cut self.__stopwords: set[str] = set(punctuation).union(_STOPWORDS) - @staticmethod - def __rank(ranking: dict, n: int) -> list: - return nlargest(n, ranking, key=ranking.get) # type: ignore[arg-type] - def __compute_frequencies( self, word_tokenized_sents: list[list[str]] - ) -> defaultdict: - word_freqs: defaultdict[str, float] = defaultdict(int) + ) -> dict[str, float]: + counts: Counter[str] = Counter() for sent in word_tokenized_sents: for word in sent: if word not in self.__stopwords: - word_freqs[word] += 1 + counts[word] += 1 - max_freq = float(max(word_freqs.values())) - for w in list(word_freqs): - word_freqs[w] = word_freqs[w] / max_freq - if ( - word_freqs[w] >= self.__max_cut - or word_freqs[w] <= self.__min_cut - ): - del word_freqs[w] + if not counts: + return {} - return word_freqs + max_freq = float(max(counts.values())) + freqs = {w: (c / max_freq) for w, c in counts.items()} + return { + w: f + for w, f in freqs.items() + if self.__min_cut < f < self.__max_cut + } def summarize( self, text: str, n: int, tokenizer: str = "newmm" @@ -61,15 +57,9 @@ def summarize( word_tokenized_sents = [ word_tokenize(sent, engine=tokenizer) for sent in sents ] - self.__freq: "defaultdict[str, float]" = self.__compute_frequencies( - word_tokenized_sents - ) - ranking: defaultdict[int, float] = defaultdict(int) - + self.__freq = self.__compute_frequencies(word_tokenized_sents) + scores = [0.0] * len(word_tokenized_sents) for i, sent in enumerate(word_tokenized_sents): - for w in sent: - if w in self.__freq: - ranking[i] += self.__freq[w] - summaries_idx = self.__rank(ranking, n) - + scores[i] = sum(self.__freq.get(w, 0.0) for w in sent) + summaries_idx = nlargest(n, range(len(scores)), key=scores.__getitem__) return [sents[j] for j in summaries_idx] diff --git a/pythainlp/summarize/keybert.py b/pythainlp/summarize/keybert.py index a6007a6dd..8365d254c 100644 --- a/pythainlp/summarize/keybert.py +++ b/pythainlp/summarize/keybert.py @@ -13,7 +13,7 @@ from __future__ import annotations from collections import Counter -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Optional, Union, cast from pythainlp.corpus import thai_stopwords from pythainlp.tokenize import word_tokenize @@ -22,6 +22,7 @@ from collections.abc import Iterable import numpy as np + from numpy.typing import NDArray from transformers.pipelines.base import Pipeline @@ -112,11 +113,11 @@ def extract_keywords( """ try: text = text.strip() - except AttributeError: + except AttributeError as exc: raise AttributeError( f"Unable to process data of type {type(text)}. " f"Please provide input of string type." - ) + ) from exc if not text: return [] @@ -141,22 +142,29 @@ def extract_keywords( else: return [kw for kw, _ in keywords] - def embed(self, docs: Union[str, list[str]]) -> np.ndarray: - """Create an embedding of each input in `docs` by averaging vectors from the last hidden layer.""" + def embed(self, docs: Union[str, list[str]]) -> "NDArray[np.float32]": + """Create embeddings by averaging vectors from the last hidden layer. + + :param Union[str, list[str]] docs: input document or documents + :return: embeddings as a float32 array with one row per input document + :rtype: numpy.typing.NDArray[numpy.float32] + """ import numpy as np embs = self.ft_pipeline(docs) if isinstance(docs, str) or len(docs) == 1: # embed doc. return shape = [1, hidden_size] - emb_mean = np.array(embs).mean(axis=1) + emb_mean = np.array(embs, dtype=np.float32).mean( + axis=1, dtype=np.float32 + ) else: # mean of embedding of each word # return shape = [len(docs), hidden_size] emb_mean = np.stack( [np.array(emb[0]).mean(axis=0) for emb in embs] - ) + ).astype(np.float32) - return emb_mean + return cast("NDArray[np.float32]", emb_mean) def _generate_ngrams( @@ -210,8 +218,8 @@ def _join_ngram(ngrams: list[tuple[str, ...]]) -> list[str]: def _rank_keywords( - doc_vector: np.ndarray, - word_vectors: np.ndarray, + doc_vector: "NDArray[np.float32]", + word_vectors: "NDArray[np.float32]", keywords: list[str], max_keywords: int, ) -> list[tuple[str, float]]: diff --git a/pythainlp/tag/_tag_perceptron.py b/pythainlp/tag/_tag_perceptron.py index ac96ecc78..e9b97c410 100644 --- a/pythainlp/tag/_tag_perceptron.py +++ b/pythainlp/tag/_tag_perceptron.py @@ -21,7 +21,7 @@ import json from collections import defaultdict -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Any, Optional if TYPE_CHECKING: from collections.abc import Iterable @@ -202,7 +202,7 @@ def train( # save the model as JSON if save_loc is not None: - data: dict[str, Union[dict, list]] = {} + data: dict[str, Any] = {} data["weights"] = self.model.weights data["tagdict"] = self.tagdict data["classes"] = list(self.classes) diff --git a/pythainlp/tag/chunk.py b/pythainlp/tag/chunk.py index a8bfb971c..8c65ca4e8 100644 --- a/pythainlp/tag/chunk.py +++ b/pythainlp/tag/chunk.py @@ -6,6 +6,7 @@ .. deprecated:: 5.3.2 :func:`chunk_parse` has moved to :mod:`pythainlp.chunk`. """ + from __future__ import annotations from pythainlp.chunk import chunk_parse as _chunk_parse diff --git a/pythainlp/tag/crfchunk.py b/pythainlp/tag/crfchunk.py index 5f2faf107..772a15a78 100644 --- a/pythainlp/tag/crfchunk.py +++ b/pythainlp/tag/crfchunk.py @@ -7,6 +7,7 @@ This module has been superseded by :mod:`pythainlp.chunk`. Import :class:`pythainlp.chunk.CRFChunkParser` directly. """ + from __future__ import annotations from pythainlp.chunk.crfchunk import CRFChunkParser diff --git a/pythainlp/tag/thai_nner.py b/pythainlp/tag/thai_nner.py index d956f3333..ef2d683b1 100644 --- a/pythainlp/tag/thai_nner.py +++ b/pythainlp/tag/thai_nner.py @@ -9,7 +9,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union from pythainlp.corpus import get_corpus_path @@ -19,11 +19,13 @@ __all__: list[str] = ["ThaiNNER"] -def _is_contained_in(entity: dict, container: dict) -> bool: +def _is_contained_in( + entity: dict[str, Any], container: dict[str, Any] +) -> bool: """Check if an entity is strictly contained within a container entity. - :param dict entity: Entity to check - :param dict container: Potential container entity + :param dict[str, Any] entity: Entity to check + :param dict[str, Any] container: Potential container entity :return: True if entity is strictly contained in container :rtype: bool """ @@ -39,17 +41,20 @@ def _is_contained_in(entity: dict, container: dict) -> bool: ) -def get_top_level_entities(entities: list[dict]) -> list[dict]: +def get_top_level_entities( + entities: list[dict[str, Any]], +) -> list[dict[str, Any]]: """Extract only top-level (outermost) entities from nested NER results. In nested NER, entities can contain other entities. This function filters the results to return only the outermost entities that are not contained within other entity. - :param list[dict] entities: List of entity dictionaries with 'span', - 'text', and 'entity_type' keys + :param list[dict[str, Any]] entities: List of entity dictionaries with + 'span', 'text', and 'entity_type' + keys :return: List of top-level entities only - :rtype: list[dict] + :rtype: list[dict[str, Any]] :Example: :: @@ -76,7 +81,7 @@ def get_top_level_entities(entities: list[dict]) -> list[dict]: entities, key=lambda x: (x["span"][0], -x["span"][1]) ) - top_level: list[dict] = [] + top_level: list[dict[str, Any]] = [] for ent in sorted_entities: is_contained = False # Only check against entities already in top_level @@ -98,8 +103,9 @@ class ThaiNNER: The model recognizes nested named entities in Thai text, supporting 104 entity types across multiple levels of nesting. - :param str path_model: Path to the Thai-NNER model file. - If not specified, downloads from PyThaiNLP corpus. + :param Optional[str] path_model: Path to the Thai-NNER model file. + If not specified, uses the default + PyThaiNLP corpus path. :Example: :: diff --git a/pythainlp/tag/thainer.py b/pythainlp/tag/thainer.py index 0e27892f9..f8a10ad96 100644 --- a/pythainlp/tag/thainer.py +++ b/pythainlp/tag/thainer.py @@ -27,7 +27,10 @@ def _is_stopword(word: str) -> bool: # เช็คว่าเป็นคำ return word in thai_stopwords() -def _doc2features(doc: list, i: int) -> dict: +def _doc2features( + doc: list[tuple[str, str]], i: int +) -> dict[str, Union[str, bool]]: + features: dict[str, Union[str, bool]] word = doc[i][0] postag = doc[i][1] @@ -47,7 +50,7 @@ def _doc2features(doc: list, i: int) -> dict: if i > 0: prevword = doc[i - 1][0] prevpostag = doc[i - 1][1] - prev_features = { + prev_features: dict[str, Union[str, bool]] = { "word.prevword": prevword, "word.previsspace": prevword.isspace(), "word.previsthai": is_thai(prevword), @@ -63,7 +66,7 @@ def _doc2features(doc: list, i: int) -> dict: if i < len(doc) - 1: nextword = doc[i + 1][0] nextpostag = doc[i + 1][1] - next_features = { + next_features: dict[str, Union[str, bool]] = { "word.nextword": nextword, "word.nextisspace": nextword.isspace(), "word.nextpostag": nextpostag, diff --git a/pythainlp/tag/tltk.py b/pythainlp/tag/tltk.py index 3bbe0ec44..64caa265a 100644 --- a/pythainlp/tag/tltk.py +++ b/pythainlp/tag/tltk.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import Union +from typing import Union, cast try: from tltk import nlp @@ -19,8 +19,8 @@ def pos_tag(words: list[str], corpus: str = "tnc") -> list[tuple[str, str]]: if corpus != "tnc": - raise ValueError(f"tltk not support {0} corpus.") - return nlp.pos_tag_wordlist(words) # type: ignore[no-any-return] + raise ValueError(f"tltk not support {corpus!r} corpus.") + return cast(list[tuple[str, str]], nlp.pos_tag_wordlist(words)) def _post_process(text: str) -> str: diff --git a/pythainlp/tag/wangchanberta_onnx.py b/pythainlp/tag/wangchanberta_onnx.py index cf62104c0..5e28bb8b7 100644 --- a/pythainlp/tag/wangchanberta_onnx.py +++ b/pythainlp/tag/wangchanberta_onnx.py @@ -4,11 +4,12 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, Any, Optional, Union, cast if TYPE_CHECKING: import numpy as np import sentencepiece as spm + from numpy.typing import NDArray from onnxruntime import InferenceSession, SessionOptions from pythainlp.corpus import get_corpus_path @@ -26,14 +27,14 @@ class WngchanBerta_ONNX: sp: "spm.SentencePieceProcessor" _json: dict[str, Any] id2tag: dict[str, str] - _s: dict[str, "np.ndarray"] + _s: dict[str, "NDArray[np.int64]"] def __init__( self, model_name: str, model_version: str, file_onnx: str, - providers: list[str] = ["CPUExecutionProvider"], + providers: Optional[list[str]] = None, ) -> None: import sentencepiece as spm from onnxruntime import ( @@ -44,6 +45,8 @@ def __init__( self.model_name = model_name self.model_version = model_version + if providers is None: + providers = ["CPUExecutionProvider"] self.options = SessionOptions() self.options.graph_optimization_level = ( GraphOptimizationLevel.ORT_ENABLE_ALL @@ -58,9 +61,8 @@ def __init__( ) self.session.disable_fallback() self.outputs_name = self.session.get_outputs()[0].name - self.sp = spm.SentencePieceProcessor( - model_file=safe_path_join(_corpus_base, "sentencepiece.bpe.model") - ) + self.sp = spm.SentencePieceProcessor() + self.sp.Load(safe_path_join(_corpus_base, "sentencepiece.bpe.model")) with open( safe_path_join(_corpus_base, "config.json"), encoding="utf-8-sig", @@ -68,10 +70,17 @@ def __init__( self._json = json.load(fh) self.id2tag = self._json["id2label"] - def build_tokenizer(self, sent: str) -> dict[str, "np.ndarray"]: + def build_tokenizer(self, sent: str) -> dict[str, "NDArray[np.int64]"]: + """Build ONNX tokenizer inputs for a sentence. + + :param str sent: input sentence + :return: model inputs containing int64 ``input_ids`` and + ``attention_mask`` arrays + :rtype: dict[str, numpy.typing.NDArray[numpy.int64]] + """ import numpy as np - _t = [5] + [i + 4 for i in self.sp.encode(sent)] + [6] + _t = [5] + [i + 4 for i in self.sp.EncodeAsIds(sent)] + [6] model_inputs = {} model_inputs["input_ids"] = np.array([_t], dtype=np.int64) model_inputs["attention_mask"] = np.array( @@ -79,21 +88,32 @@ def build_tokenizer(self, sent: str) -> dict[str, "np.ndarray"]: ) return model_inputs - def postprocess(self, logits_data: "np.ndarray") -> "np.ndarray": + def postprocess( + self, logits_data: "NDArray[np.float32]" + ) -> "NDArray[np.float32]": + """Convert raw logits to probabilities. + + :param numpy.typing.NDArray[numpy.float32] logits_data: raw model + logits + :return: probability scores as a float32 array + :rtype: numpy.typing.NDArray[numpy.float32] + """ import numpy as np logits_t = logits_data[0] maxes = np.max(logits_t, axis=-1, keepdims=True) shifted_exp = np.exp(logits_t - maxes) scores = shifted_exp / shifted_exp.sum(axis=-1, keepdims=True) - return scores + return cast("NDArray[np.float32]", scores) def clean_output( self, list_text: list[tuple[str, str]] ) -> list[tuple[str, str]]: return list_text - def totag(self, post: np.ndarray, sent: str) -> list[tuple[str, str]]: + def totag( + self, post: "NDArray[np.float32]", sent: str + ) -> list[tuple[str, str]]: tag = [] _s = self.sp.EncodeAsPieces(sent) for i in range(len(_s)): @@ -115,10 +135,11 @@ def _config( def get_ner( self, text: str, tag: bool = False ) -> Union[str, list[tuple[str, str]]]: - self._s: dict[str, "np.ndarray"] = self.build_tokenizer(text) - logits = self.session.run( + self._s = self.build_tokenizer(text) + logits_raw = self.session.run( output_names=[self.outputs_name], input_feed=self._s )[0] + logits = cast("NDArray[np.float32]", logits_raw) _tag = self.clean_output(self.totag(self.postprocess(logits), text)) if tag: _tag = self._config(_tag) diff --git a/pythainlp/tokenize/core.py b/pythainlp/tokenize/core.py index e896507d8..ff133c68b 100644 --- a/pythainlp/tokenize/core.py +++ b/pythainlp/tokenize/core.py @@ -7,7 +7,7 @@ import re from collections import deque -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Optional, Union, cast if TYPE_CHECKING: from collections.abc import Iterable @@ -637,7 +637,7 @@ def paragraph_tokenize( It might be a typo; if not, please consult our document.""" ) - return segments # type: ignore[return-value] + return cast(list[list[str]], segments) def subword_tokenize( diff --git a/pythainlp/tokenize/han_solo.py b/pythainlp/tokenize/han_solo.py index 68bc1117c..dc1b18ecf 100644 --- a/pythainlp/tokenize/han_solo.py +++ b/pythainlp/tokenize/han_solo.py @@ -73,7 +73,7 @@ def featurize( padding: bool = True, indiv_char: bool = True, return_type: str = "list", - ) -> dict[str, list]: + ) -> dict[str, list[Any]]: if padding: sentence = self.pad(sentence) all_features_list: list[list[str]] = [] diff --git a/pythainlp/tokenize/newmm.py b/pythainlp/tokenize/newmm.py index d03e3c529..07542d1bc 100644 --- a/pythainlp/tokenize/newmm.py +++ b/pythainlp/tokenize/newmm.py @@ -61,7 +61,7 @@ def _bfs_paths_graph( - graph: defaultdict, start: int, goal: int + graph: defaultdict[int, list[int]], start: int, goal: int ) -> Generator[list[int], None, None]: # visited set prevents re-exploring nodes already reached via a shorter # path, converting worst-case BFS from exponential to O(V + E). diff --git a/pythainlp/translate/tokenization_small100.py b/pythainlp/translate/tokenization_small100.py index b5955d604..6625cf9a5 100644 --- a/pythainlp/translate/tokenization_small100.py +++ b/pythainlp/translate/tokenization_small100.py @@ -243,7 +243,7 @@ def tgt_lang(self, new_tgt_lang: str) -> None: self.set_lang_special_tokens(self._tgt_lang) def _tokenize(self, text: str) -> list[str]: - return self.sp_model.encode(text, out_type=str) # type: ignore[no-any-return] + return cast(list[str], self.sp_model.encode(text, out_type=str)) def _convert_token_to_id(self, token: str) -> int: if token in self.lang_token_to_id: @@ -256,12 +256,12 @@ def _convert_id_to_token(self, index: int) -> str: return self.id_to_lang_token[index] token = self.decoder.get(index, self.unk_token) if token is None: - return self.unk_token # type: ignore[no-any-return] + return cast(str, self.unk_token) return token def convert_tokens_to_string(self, tokens: list[str]) -> str: """Converts a sequence of tokens (strings for sub-words) in a single string.""" - return self.sp_model.decode(tokens) # type: ignore[no-any-return] + return cast(str, self.sp_model.decode(tokens)) def get_special_tokens_mask( self, @@ -289,10 +289,13 @@ def get_special_tokens_mask( """ if already_has_special_tokens: # External library method - return super().get_special_tokens_mask( # type: ignore[no-any-return] - token_ids_0=token_ids_0, - token_ids_1=token_ids_1, - already_has_special_tokens=True, + return cast( + list[int], + super().get_special_tokens_mask( + token_ids_0=token_ids_0, + token_ids_1=token_ids_1, + already_has_special_tokens=True, + ), ) prefix_ones = ( @@ -349,20 +352,20 @@ def build_inputs_with_special_tokens( + self.suffix_tokens ) - def get_vocab(self) -> dict: + def get_vocab(self) -> dict[str, int]: vocab = { self.convert_ids_to_tokens(i): i for i in range(self.vocab_size) } vocab.update(self.added_tokens_encoder) return vocab - def __getstate__(self) -> dict: + def __getstate__(self) -> dict[str, Any]: state = self.__dict__.copy() state["sp_model"] = None return state - def __setstate__(self, d: dict) -> None: - self.__dict__: dict = d + def __setstate__(self, d: dict[str, Any]) -> None: + self.__dict__: dict[str, Any] = d # for backward compatibility if not hasattr(self, "sp_model_kwargs"): @@ -425,7 +428,7 @@ def _build_translation_inputs( ) self.tgt_lang: str = tgt_lang inputs = self(raw_inputs, add_special_tokens=True, **extra_kwargs) - return inputs # type: ignore[no-any-return] + return cast(dict[str, Any], inputs) def _switch_to_input_mode(self) -> None: self.set_lang_special_tokens(self.tgt_lang) @@ -462,7 +465,7 @@ def load_spm( def load_json(path: str) -> Union[dict[str, str], list[str]]: with open(path) as f: - return json.load(f) # type: ignore[no-any-return] + return cast(Union[dict[str, str], list[str]], json.load(f)) def save_json( diff --git a/pythainlp/translate/word2word_translate.py b/pythainlp/translate/word2word_translate.py index 2dc8024e6..882ffc65d 100644 --- a/pythainlp/translate/word2word_translate.py +++ b/pythainlp/translate/word2word_translate.py @@ -3,7 +3,7 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import Optional +from typing import Optional, cast from word2word import Word2word @@ -87,4 +87,4 @@ def translate(word: str, src: str, target: str) -> Optional[list[str]]: elif src == target: return [word] _engine = Word2word(src, target) - return _engine(word) # type: ignore[no-any-return] + return cast(Optional[list[str]], _engine(word)) diff --git a/pythainlp/transliterate/ipa.py b/pythainlp/transliterate/ipa.py index bdffe8b08..487a8860f 100644 --- a/pythainlp/transliterate/ipa.py +++ b/pythainlp/transliterate/ipa.py @@ -11,18 +11,20 @@ from __future__ import annotations +from typing import cast + import epitran _EPI_THA: epitran.Epitran = epitran.Epitran("tha-Thai") def transliterate(text: str) -> str: - return _EPI_THA.transliterate(text) # type: ignore[no-any-return] + return cast(str, _EPI_THA.transliterate(text)) def trans_list(text: str) -> list[str]: - return _EPI_THA.trans_list(text) # type: ignore[no-any-return] + return cast(list[str], _EPI_THA.trans_list(text)) def xsampa_list(text: str) -> list[str]: - return _EPI_THA.xsampa_list(text) # type: ignore[no-any-return] + return cast(list[str], _EPI_THA.xsampa_list(text)) diff --git a/pythainlp/transliterate/lookup.py b/pythainlp/transliterate/lookup.py index 36178b39a..52103620e 100644 --- a/pythainlp/transliterate/lookup.py +++ b/pythainlp/transliterate/lookup.py @@ -10,7 +10,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, cast if TYPE_CHECKING: from collections.abc import Callable @@ -38,7 +38,7 @@ def follow_rtgs(text: str) -> Optional[bool]: except IndexError: return None else: - return follow # type: ignore[return-value] + return cast(bool, follow) def _romanize(text: str, fallback_func: Callable[[str], str]) -> str: @@ -52,7 +52,7 @@ def _romanize(text: str, fallback_func: Callable[[str], str]) -> str: except TypeError as e: raise TypeError(f"`fallback_engine` is not callable. {e}") from e else: - return lookup # type: ignore[return-value] + return cast(str, lookup) def romanize(text: str, fallback_func: Callable[[str], str]) -> str: diff --git a/pythainlp/transliterate/thai2rom_onnx.py b/pythainlp/transliterate/thai2rom_onnx.py index 83c8c030d..e11e614ce 100644 --- a/pythainlp/transliterate/thai2rom_onnx.py +++ b/pythainlp/transliterate/thai2rom_onnx.py @@ -6,7 +6,7 @@ from __future__ import annotations import json -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from onnxruntime import InferenceSession @@ -16,6 +16,7 @@ from typing import Dict, List import numpy as np + from numpy.typing import NDArray _MODEL_ENCODER_NAME: str = "thai2rom_encoder_onnx" _MODEL_DECODER_NAME: str = "thai2rom_decoder_onnx" @@ -58,7 +59,7 @@ def __init__(self) -> None: ) # loader = torch.load(self.__model_filename, map_location=device) - with open(str(self.__config_filename)) as f: + with open(str(self.__config_filename), encoding="utf-8") as f: loader = json.load(f) OUTPUT_DIM = loader["output_dim"] @@ -94,8 +95,13 @@ def __init__(self) -> None: target_vocab_size=OUTPUT_DIM, ) - def _prepare_sequence_in(self, text: str) -> "np.ndarray": - """Prepare input sequence for ONNX""" + def _prepare_sequence_in(self, text: str) -> "NDArray[np.int64]": + """Prepare an int64 input sequence for the ONNX encoder. + + :param str text: Thai text to encode + :return: encoded character ids ending with the ```` token + :rtype: numpy.typing.NDArray[numpy.int64] + """ import numpy as np idxs = [] @@ -105,7 +111,7 @@ def _prepare_sequence_in(self, text: str) -> "np.ndarray": else: idxs.append(self._char_to_ix[""]) idxs.append(self._char_to_ix[""]) - return np.array(idxs) + return np.array(idxs, dtype=np.int64) def romanize(self, text: str) -> str: """:param str text: Thai text to be romanized @@ -158,13 +164,31 @@ def __init__( self.target_vocab_size: int = target_vocab_size - def create_mask(self, source_seq: "np.ndarray") -> "np.ndarray": + def create_mask( + self, source_seq: "NDArray[np.int64]" + ) -> "NDArray[np.bool_]": + """Create a boolean mask for non-padding positions. + + :param numpy.typing.NDArray[numpy.int64] source_seq: encoded source + sequence + :return: boolean mask where True marks non-padding positions + :rtype: numpy.typing.NDArray[numpy.bool_] + """ mask = source_seq != self.pad_idx - return mask + return cast("NDArray[np.bool_]", mask) def run( - self, source_seq: "np.ndarray", source_seq_len: List[int] - ) -> "np.ndarray": + self, source_seq: "NDArray[np.int64]", source_seq_len: List[int] + ) -> "NDArray[np.float32]": + """Run ONNX seq2seq decoding and return logits. + + :param numpy.typing.NDArray[numpy.int64] source_seq: encoded source + sequence with shape ``(batch_size, sequence_length)`` + :param List[int] source_seq_len: unpadded source lengths + :return: decoder logits as a float32 array with shape + ``(decoded_length, batch_size, target_vocab_size)`` + :rtype: numpy.typing.NDArray[numpy.float32] + """ # source_seq: (batch_size, MAX_LENGTH) # source_seq_len: (batch_size, 1) # target_seq: (batch_size, MAX_LENGTH) @@ -176,7 +200,9 @@ def run( max_len = self.max_length # target_vocab_size = self.decoder.vocabulary_size - outputs = np.zeros((max_len, batch_size, self.target_vocab_size)) + outputs: "NDArray[np.float32]" = np.zeros( + (max_len, batch_size, self.target_vocab_size), dtype=np.float32 + ) expected_encoder_outputs = [ output.name for output in self.encoder.get_outputs() @@ -206,7 +232,7 @@ def run( mask = self.create_mask(source_seq[:, 0:max_source_len]) for di in range(max_len): - decoder_output, decoder_hidden = self.decoder.run( + decoder_output_raw, decoder_hidden = self.decoder.run( input_feed={ "decoder_input": decoder_input.astype("int32"), "decoder_hidden_1": decoder_hidden, @@ -218,11 +244,12 @@ def run( self.decoder.get_outputs()[1].name, ], ) + decoder_output = cast("NDArray[np.float32]", decoder_output_raw) topi = np.argmax(decoder_output, axis=1) outputs[di] = decoder_output - decoder_input = np.array([topi]) + decoder_input = np.array([topi], dtype=np.int64) if decoder_input.item() == end_token: return outputs[:di] diff --git a/pythainlp/transliterate/thaig2p.py b/pythainlp/transliterate/thaig2p.py index b4fa830ab..720006da1 100644 --- a/pythainlp/transliterate/thaig2p.py +++ b/pythainlp/transliterate/thaig2p.py @@ -8,7 +8,7 @@ from __future__ import annotations import random -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Any, Optional, Union import torch import torch.nn.functional as F @@ -163,7 +163,7 @@ def __init__( def forward( self, sequences: torch.Tensor, - sequences_lengths: Union[NDArray, list[int]], + sequences_lengths: Union[NDArray[Any], list[int]], ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: # sequences: (batch_size, sequence_length=MAX_LENGTH) # sequences_lengths: (batch_size) @@ -380,7 +380,7 @@ def create_mask(self, source_seq: torch.Tensor) -> torch.Tensor: def forward( self, source_seq: torch.Tensor, - source_seq_len: Union[NDArray, list[int]], + source_seq_len: Union[NDArray[Any], list[int]], target_seq: Optional[torch.Tensor], teacher_forcing_ratio: float = 0.5, ) -> torch.Tensor: diff --git a/pythainlp/transliterate/thaig2p_v2.py b/pythainlp/transliterate/thaig2p_v2.py index e90a52a36..c50558c4b 100644 --- a/pythainlp/transliterate/thaig2p_v2.py +++ b/pythainlp/transliterate/thaig2p_v2.py @@ -9,7 +9,7 @@ # Use a pipeline as a high-level helper from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, cast if TYPE_CHECKING: from transformers import Pipeline @@ -39,7 +39,8 @@ def __init__(self, device: str = "cpu") -> None: ) def g2p(self, text: str) -> str: - return self.pipe(text)[0]["generated_text"] # type: ignore[no-any-return] + outputs = cast(list[dict[str, str]], self.pipe(text)) + return outputs[0]["generated_text"] _THAI_G2P: Optional[ThaiG2P] = None diff --git a/pythainlp/transliterate/tltk.py b/pythainlp/transliterate/tltk.py index a1b79172f..852e4a912 100644 --- a/pythainlp/transliterate/tltk.py +++ b/pythainlp/transliterate/tltk.py @@ -3,6 +3,8 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +from typing import cast + try: from tltk.nlp import g2p, th2ipa, th2roman except ImportError: @@ -20,19 +22,24 @@ def romanize(text: str) -> str: """ # Replace ฅ with ค to avoid KeyError in tltk (out-of-vocabulary issue) text = text.replace("ฅ", "ค") - _temp = th2roman(text) - return _temp[: _temp.rfind(" ")].replace("", "") # type: ignore[no-any-return] + _temp = cast(str, th2roman(text)) + return _temp[: _temp.rfind(" ")].replace("", "") def tltk_g2p(text: str) -> str: # Replace ฅ with ค to avoid KeyError in tltk (out-of-vocabulary issue) text = text.replace("ฅ", "ค") - _temp = g2p(text).split("")[1].replace("|", "").replace("|", " ") - return _temp.replace("", "") # type: ignore[no-any-return] + _temp = ( + cast(str, g2p(text)) + .split("")[1] + .replace("|", "") + .replace("|", " ") + ) + return _temp.replace("", "") def tltk_ipa(text: str) -> str: # Replace ฅ with ค to avoid KeyError in tltk (out-of-vocabulary issue) text = text.replace("ฅ", "ค") - _temp = th2ipa(text) - return _temp[: _temp.rfind(" ")].replace("", "") # type: ignore[no-any-return] + _temp = cast(str, th2ipa(text)) + return _temp[: _temp.rfind(" ")].replace("", "") diff --git a/pythainlp/transliterate/umt5_thaig2p.py b/pythainlp/transliterate/umt5_thaig2p.py index 9cba585ff..85669104f 100644 --- a/pythainlp/transliterate/umt5_thaig2p.py +++ b/pythainlp/transliterate/umt5_thaig2p.py @@ -9,7 +9,7 @@ # Use a pipeline as a high-level helper from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, cast if TYPE_CHECKING: from transformers import Pipeline @@ -39,7 +39,8 @@ def __init__(self, device: str = "cpu") -> None: ) def g2p(self, text: str) -> str: - return self.pipe(text)[0]["generated_text"] # type: ignore[no-any-return] + outputs = cast(list[dict[str, str]], self.pipe(text)) + return outputs[0]["generated_text"] _THAI_G2P: Optional[Umt5ThaiG2P] = None diff --git a/pythainlp/transliterate/w2p.py b/pythainlp/transliterate/w2p.py index d2bb31f09..ccf0b018b 100644 --- a/pythainlp/transliterate/w2p.py +++ b/pythainlp/transliterate/w2p.py @@ -7,7 +7,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Optional, cast from pythainlp.corpus import get_corpus_path @@ -40,9 +40,9 @@ class _Hparams: hp: _Hparams = _Hparams() -def _load_vocab() -> ( - tuple[dict[str, int], dict[int, str], dict[str, int], dict[int, str]] -): +def _load_vocab() -> tuple[ + dict[str, int], dict[int, str], dict[str, int], dict[int, str] +]: g2idx = {g: idx for idx, g in enumerate(hp.graphemes)} idx2g = dict(enumerate(hp.graphemes)) @@ -60,18 +60,18 @@ class Thai_W2P: p2idx: dict[str, int] idx2p: dict[int, str] checkpoint: Optional[str] - enc_emb: "NDArray" - enc_w_ih: "NDArray" - enc_w_hh: "NDArray" - enc_b_ih: "NDArray" - enc_b_hh: "NDArray" - dec_emb: "NDArray" - dec_w_ih: "NDArray" - dec_w_hh: "NDArray" - dec_b_ih: "NDArray" - dec_b_hh: "NDArray" - fc_w: "NDArray" - fc_b: "NDArray" + enc_emb: "NDArray[np.float32]" + enc_w_ih: "NDArray[np.float32]" + enc_w_hh: "NDArray[np.float32]" + enc_b_ih: "NDArray[np.float32]" + enc_b_hh: "NDArray[np.float32]" + dec_emb: "NDArray[np.float32]" + dec_w_ih: "NDArray[np.float32]" + dec_w_hh: "NDArray[np.float32]" + dec_b_ih: "NDArray[np.float32]" + dec_b_hh: "NDArray[np.float32]" + fc_w: "NDArray[np.float32]" + fc_b: "NDArray[np.float32]" word: str def __init__(self) -> None: @@ -83,7 +83,9 @@ def __init__(self) -> None: self.p2idx: dict[str, int] self.idx2p: dict[int, str] self.g2idx, self.idx2g, self.p2idx, self.idx2p = _load_vocab() - self.checkpoint: Optional[str] = get_corpus_path(_MODEL_NAME, version="0.2") + self.checkpoint: Optional[str] = get_corpus_path( + _MODEL_NAME, version="0.2" + ) if not self.checkpoint: raise FileNotFoundError( f"corpus-not-found name={_MODEL_NAME!r}\n" @@ -95,69 +97,95 @@ def __init__(self) -> None: def _load_variables(self) -> None: import numpy as np + if self.checkpoint is None: raise RuntimeError("checkpoint path is not set") with np.load(self.checkpoint, allow_pickle=False) as variables: # (29, 64). (len(graphemes), emb) - self.enc_emb: "NDArray" = variables[ - "encoder_emb_weight" - ] + self.enc_emb = cast( + "NDArray[np.float32]", variables["encoder_emb_weight"] + ) # (3*128, 64) - self.enc_w_ih: "NDArray" = variables[ - "encoder_rnn_weight_ih_l0" - ] + self.enc_w_ih = cast( + "NDArray[np.float32]", + variables["encoder_rnn_weight_ih_l0"], + ) # (3*128, 128) - self.enc_w_hh: "NDArray" = variables[ - "encoder_rnn_weight_hh_l0" - ] + self.enc_w_hh = cast( + "NDArray[np.float32]", + variables["encoder_rnn_weight_hh_l0"], + ) # (3*128,) - self.enc_b_ih: "NDArray" = variables[ - "encoder_rnn_bias_ih_l0" - ] + self.enc_b_ih = cast( + "NDArray[np.float32]", variables["encoder_rnn_bias_ih_l0"] + ) # (3*128,) - self.enc_b_hh: "NDArray" = variables[ - "encoder_rnn_bias_hh_l0" - ] + self.enc_b_hh = cast( + "NDArray[np.float32]", variables["encoder_rnn_bias_hh_l0"] + ) # (74, 64). (len(phonemes), emb) - self.dec_emb: "NDArray" = variables[ - "decoder_emb_weight" - ] + self.dec_emb = cast( + "NDArray[np.float32]", variables["decoder_emb_weight"] + ) # (3*128, 64) - self.dec_w_ih: "NDArray" = variables[ - "decoder_rnn_weight_ih_l0" - ] + self.dec_w_ih = cast( + "NDArray[np.float32]", + variables["decoder_rnn_weight_ih_l0"], + ) # (3*128, 128) - self.dec_w_hh: "NDArray" = variables[ - "decoder_rnn_weight_hh_l0" - ] + self.dec_w_hh = cast( + "NDArray[np.float32]", + variables["decoder_rnn_weight_hh_l0"], + ) # (3*128,) - self.dec_b_ih: "NDArray" = variables[ - "decoder_rnn_bias_ih_l0" - ] + self.dec_b_ih = cast( + "NDArray[np.float32]", variables["decoder_rnn_bias_ih_l0"] + ) # (3*128,) - self.dec_b_hh: "NDArray" = variables[ - "decoder_rnn_bias_hh_l0" - ] + self.dec_b_hh = cast( + "NDArray[np.float32]", variables["decoder_rnn_bias_hh_l0"] + ) # (74, 128) - self.fc_w: "NDArray" = variables["decoder_fc_weight"] + self.fc_w = cast( + "NDArray[np.float32]", variables["decoder_fc_weight"] + ) # (74,) - self.fc_b: "NDArray" = variables["decoder_fc_bias"] + self.fc_b = cast( + "NDArray[np.float32]", variables["decoder_fc_bias"] + ) - def _sigmoid(self, x: "np.ndarray") -> "np.ndarray": + def _sigmoid(self, x: "NDArray[np.float32]") -> "NDArray[np.float32]": + """Apply the sigmoid function to a float32 array. + + :param numpy.typing.NDArray[numpy.float32] x: input array + :return: element-wise sigmoid values + :rtype: numpy.typing.NDArray[numpy.float32] + """ import numpy as np - return 1 / (1 + np.exp(-x)) + return cast("NDArray[np.float32]", 1 / (1 + np.exp(-x))) def _grucell( self, - x: "np.ndarray", - h: "np.ndarray", - w_ih: "np.ndarray", - w_hh: "np.ndarray", - b_ih: "np.ndarray", - b_hh: "np.ndarray", - ) -> "np.ndarray": + x: "NDArray[np.float32]", + h: "NDArray[np.float32]", + w_ih: "NDArray[np.float32]", + w_hh: "NDArray[np.float32]", + b_ih: "NDArray[np.float32]", + b_hh: "NDArray[np.float32]", + ) -> "NDArray[np.float32]": + """Run one GRU cell step on float32 inputs. + + :param numpy.typing.NDArray[numpy.float32] x: input features + :param numpy.typing.NDArray[numpy.float32] h: previous hidden state + :param numpy.typing.NDArray[numpy.float32] w_ih: input-hidden weights + :param numpy.typing.NDArray[numpy.float32] w_hh: hidden-hidden weights + :param numpy.typing.NDArray[numpy.float32] b_ih: input-hidden bias + :param numpy.typing.NDArray[numpy.float32] b_hh: hidden-hidden bias + :return: updated hidden state + :rtype: numpy.typing.NDArray[numpy.float32] + """ import numpy as np rzn_ih = np.matmul(x, w_ih.T) + b_ih @@ -182,14 +210,27 @@ def _grucell( def _gru( self, - x: "np.ndarray", + x: "NDArray[np.float32]", steps: int, - w_ih: "np.ndarray", - w_hh: "np.ndarray", - b_ih: "np.ndarray", - b_hh: "np.ndarray", - h0: Optional["np.ndarray"] = None, - ) -> "np.ndarray": + w_ih: "NDArray[np.float32]", + w_hh: "NDArray[np.float32]", + b_ih: "NDArray[np.float32]", + b_hh: "NDArray[np.float32]", + h0: Optional["NDArray[np.float32]"] = None, + ) -> "NDArray[np.float32]": + """Run a GRU over multiple time steps. + + :param numpy.typing.NDArray[numpy.float32] x: input sequence tensor + :param int steps: number of decoding steps + :param numpy.typing.NDArray[numpy.float32] w_ih: input-hidden weights + :param numpy.typing.NDArray[numpy.float32] w_hh: hidden-hidden weights + :param numpy.typing.NDArray[numpy.float32] b_ih: input-hidden bias + :param numpy.typing.NDArray[numpy.float32] b_hh: hidden-hidden bias + :param Optional[numpy.typing.NDArray[numpy.float32]] h0: initial + hidden state + :return: hidden states for all time steps + :rtype: numpy.typing.NDArray[numpy.float32] + """ import numpy as np if h0 is None: @@ -203,12 +244,20 @@ def _gru( return outputs - def _encode(self, word: str) -> "np.ndarray": + def _encode(self, word: str) -> "NDArray[np.float32]": + """Encode a word into its embedding sequence tensor. + + :param str word: input Thai word + :return: float32 embedding sequence for the encoder + :rtype: numpy.typing.NDArray[numpy.float32] + """ import numpy as np chars = list(word) + [""] - x = [self.g2idx.get(char, self.g2idx[""]) for char in chars] - x = np.take(self.enc_emb, np.expand_dims(x, 0), axis=0) + char_ids = [ + self.g2idx.get(char, self.g2idx[""]) for char in chars + ] + x = np.take(self.enc_emb, np.expand_dims(char_ids, 0), axis=0) return x @@ -244,7 +293,7 @@ def _predict(self, word: str) -> str: dec = np.take(self.dec_emb, [2], axis=0) # 2: h = last_hidden - preds = [] + preds: list[int] = [] for _ in range(20): h = self._grucell( dec, @@ -255,7 +304,7 @@ def _predict(self, word: str) -> str: self.dec_b_hh, ) # (b, h) logits = np.matmul(h, self.fc_w.T) + self.fc_b - pred = logits.argmax() + pred = int(logits.argmax()) if pred == 3: break preds.append(pred) diff --git a/pythainlp/transliterate/wunsen.py b/pythainlp/transliterate/wunsen.py index 0d79c7672..ddd0ec12f 100644 --- a/pythainlp/transliterate/wunsen.py +++ b/pythainlp/transliterate/wunsen.py @@ -12,7 +12,7 @@ from __future__ import annotations -from typing import Optional, Union +from typing import Optional, Union, cast from wunsen import ThapSap @@ -153,4 +153,4 @@ def transliterate( if self.thap_value is None: raise RuntimeError("ThapSap model not initialized") - return self.thap_value.thap(text) # type: ignore[no-any-return] + return cast(str, self.thap_value.thap(text)) diff --git a/pythainlp/ulmfit/core.py b/pythainlp/ulmfit/core.py index 21f60a091..f83343573 100644 --- a/pythainlp/ulmfit/core.py +++ b/pythainlp/ulmfit/core.py @@ -6,7 +6,7 @@ from __future__ import annotations import collections -from typing import TYPE_CHECKING, Any, Optional, Union +from typing import TYPE_CHECKING, Optional, cast import torch @@ -16,6 +16,7 @@ import numpy as np from fastai.basic_data import DataBunch from fastai.basic_train import Learner + from numpy.typing import NDArray from pythainlp.corpus import get_corpus_path from pythainlp.tokenize import thai2fit_tokenizer @@ -114,21 +115,25 @@ def get_thwiki_lstm() -> dict[str, str]: def process_thai( text: str, - pre_rules: Collection = pre_rules_th_sparse, - tok_func: Optional[Callable] = None, - post_rules: Collection = post_rules_th_sparse, -) -> Collection[str]: + pre_rules: Optional[Collection[Callable[[str], str]]] = None, + tok_func: Optional[Callable[[str], list[str]]] = None, + post_rules: Optional[Collection[Callable[[list[str]], list[str]]]] = None, +) -> list[str]: """Process Thai texts for models (with sparse features as default) :param str text: text to be cleaned - :param list[func] pre_rules: rules to apply before tokenization. - :param func tok_func: tokenization function (by default, **tok_func** is - :func:`pythainlp.tokenize.word_tokenize`) + :param Optional[Collection[Callable[[str], str]]] pre_rules: rules to + apply before tokenization. If None, use the default sparse pre-rules. + :param Optional[Callable[[str], list[str]]] tok_func: tokenization + function. By default, **tok_func** is + :func:`pythainlp.tokenize.word_tokenize`. - :param list[func] post_rules: rules to apply after tokenizations + :param Optional[Collection[Callable[[list[str]], list[str]]]] post_rules: + rules to apply after tokenization. If None, use the default sparse + post-rules. :return: a list of cleaned tokenized texts - :rtype: Collection[str] + :rtype: list[str] :Note: @@ -177,23 +182,30 @@ def process_thai( """ - res: Union[str, list[str]] = text + processed_text = text + if pre_rules is None: + pre_rules = pre_rules_th_sparse + if post_rules is None: + post_rules = cast( + Collection[Callable[[list[str]], list[str]]], + post_rules_th_sparse, + ) if tok_func is None: tok_func = thai2fit_tokenizer().word_tokenize - for rule in pre_rules: - res = rule(res) - res = tok_func(res) # type: ignore[arg-type] - for rule in post_rules: - res = rule(res) + for pre_rule in pre_rules: + processed_text = pre_rule(processed_text) + tokens = tok_func(processed_text) + for post_rule in post_rules: + tokens = post_rule(tokens) - return res + return tokens def document_vector( text: str, learn: "Learner", data: "DataBunch", agg: str = "mean" -) -> "np.ndarray": +) -> "NDArray[np.float32]": """This function vectorizes Thai input text into a 400 dimension vector using :class:`fastai` language model and data bunch. @@ -205,9 +217,9 @@ def document_vector( :param str agg: name of aggregation methods for word embeddings The available methods are "mean" and "sum" - :return: :class:`numpy.array` of document vector sized 400 based on - the encoder of the model - :rtype: :class:`numpy.ndarray((1, 400))` + :return: :class:`numpy.ndarray` of dtype ``numpy.float32`` containing + the document vector produced by the encoder + :rtype: numpy.typing.NDArray[numpy.float32] :Example: @@ -237,19 +249,22 @@ def document_vector( device ) m = learn.model[0].encoder.to(device) - res = m(t).cpu().detach().numpy() + res = m(t).cpu().detach().numpy().astype("float32", copy=False) if agg == "mean": - res = res.mean(0) + res = res.mean(0, dtype="float32") elif agg == "sum": - res = res.sum(0) + res = res.sum(0, dtype="float32") else: raise ValueError("Aggregate by mean or sum") - return res + return cast("NDArray[np.float32]", res) def merge_wgts( - em_sz: int, wgts: dict[str, Any], itos_pre: list[str], itos_new: list[str] + em_sz: int, + wgts: dict[str, torch.Tensor], + itos_pre: list[str], + itos_new: list[str], ) -> dict[str, torch.Tensor]: """This function is to insert new vocab into an existing model named `wgts` and update the model's weights for new vocab with the average embedding. @@ -287,10 +302,10 @@ def merge_wgts( # [0.5952, 0.4453, 0.0011]])} """ vocab_size = len(itos_new) - enc_wgts = wgts["0.encoder.weight"].numpy() + enc_wgts = wgts["0.encoder.weight"].numpy().astype("float32", copy=False) # Average weight of encoding - row_m = enc_wgts.mean(0) + row_m = enc_wgts.mean(0, dtype="float32") stoi_pre = collections.defaultdict( lambda: -1, {v: k for k, v in enumerate(itos_pre)} ) diff --git a/pythainlp/util/abbreviation.py b/pythainlp/util/abbreviation.py index 19e080371..85d3e443b 100644 --- a/pythainlp/util/abbreviation.py +++ b/pythainlp/util/abbreviation.py @@ -5,7 +5,7 @@ from __future__ import annotations -from typing import Optional +from typing import Optional, cast def abbreviation_to_full_text( @@ -44,4 +44,4 @@ def abbreviation_to_full_text( pip install pythainlp[abbreviation]. """ ) - return _replace(text, top_k=top_k) # type: ignore[no-any-return] + return cast(list[tuple[str, Optional[float]]], _replace(text, top_k=top_k)) diff --git a/pythainlp/util/collate.py b/pythainlp/util/collate.py index 3431dd73c..c0a821ee1 100644 --- a/pythainlp/util/collate.py +++ b/pythainlp/util/collate.py @@ -26,13 +26,13 @@ def _thkey(word: str) -> str: return cv + tone -def collate(data: Iterable, reverse: bool = False) -> list[str]: +def collate(data: Iterable[str], reverse: bool = False) -> list[str]: """Sorts strings (almost) according to Thai dictionary. Important notes: this implementation ignores tone marks and symbols - :param data: a list of words to be sorted - :type data: Iterable + :param data: an iterable of words to be sorted + :type data: Iterable[str] :param reverse: If `reverse` is set to **True** the result will be sorted in descending order. Otherwise, the result will be sorted in ascending order, defaults to False @@ -40,7 +40,7 @@ def collate(data: Iterable, reverse: bool = False) -> list[str]: :return: a list of strings, sorted alphabetically, (almost) according to Thai dictionary - :rtype: List[str] + :rtype: list[str] :Example: :: diff --git a/pythainlp/util/keywords.py b/pythainlp/util/keywords.py index 383dbb713..74d5ac2dc 100644 --- a/pythainlp/util/keywords.py +++ b/pythainlp/util/keywords.py @@ -13,7 +13,7 @@ def rank( words: list[str], exclude_stopwords: bool = False -) -> Optional[Counter]: +) -> Optional[Counter[str]]: """Count word frequencies given a list of Thai words with an option to exclude stopwords. @@ -23,8 +23,9 @@ def rank( Otherwise, the stopwords will be counted. By default, `exclude_stopwords`is set to **False** - :return: a Counter object representing word frequencies in the text - :rtype: :class:`collections.Counter` + :return: a Counter object representing word frequencies in the text, + or None if `words` is empty + :rtype: Optional[collections.Counter[str]] :Example: diff --git a/pythainlp/util/numtoword.py b/pythainlp/util/numtoword.py index 921eec0c2..95d098c3e 100644 --- a/pythainlp/util/numtoword.py +++ b/pythainlp/util/numtoword.py @@ -35,12 +35,14 @@ def bahttext(number: float) -> str: a suffix "บาท" (Baht). The precision will be fixed at two decimal places (0.00) to fit "สตางค์" (Satang) unit. - This function works similarly to the `BAHTTEXT` function in Microsoft Excel. + This function works similarly to the ``BAHTTEXT`` function in Microsoft Excel. :param float number: number to be converted into Thai Baht currency format :return: text representing the amount of money in the format of Thai currency :rtype: str + :raises TypeError: if *number* is not a numeric type + :Example: :: @@ -55,11 +57,14 @@ def bahttext(number: float) -> str: bahttext(200) # output: สองร้อยบาทถ้วน """ + if not isinstance(number, (int, float)): + raise TypeError( + f"number must be a numeric type, not {type(number).__name__!r}" + ) + ret = "" - if number is None: - pass - elif number == 0: + if number == 0: ret = "ศูนย์บาทถ้วน" else: num_int_str, num_dec_str = f"{number:.2f}".split(".") diff --git a/pythainlp/util/pronounce.py b/pythainlp/util/pronounce.py index 2377eee34..d3605674d 100644 --- a/pythainlp/util/pronounce.py +++ b/pythainlp/util/pronounce.py @@ -38,7 +38,9 @@ def rhyme(word: str) -> list[str]: # output: ['กลีบ', 'กีบ', 'ครีบ', ...] """ return sorted( - i for i in _single_syllable_thai_words() if kv.is_sumpus(word, i) and i != word + i + for i in _single_syllable_thai_words() + if kv.is_sumpus(word, i) and i != word ) diff --git a/pythainlp/util/thai.py b/pythainlp/util/thai.py index 37552bb27..9b97de97f 100644 --- a/pythainlp/util/thai.py +++ b/pythainlp/util/thai.py @@ -32,50 +32,52 @@ # A comprehensive map of Thai characters to their descriptive names. # MappingProxyType makes this constant read-only at runtime. -_THAI_CHAR_NAMES: MappingProxyType[str, str] = MappingProxyType({ - # Consonants - **{char: char for char in thai_consonants}, - # Vowels and Signs - "\u0e24": "ฤ", - "\u0e26": "ฦ", - "\u0e30": "สระ อะ", - "\u0e31": "ไม้หันอากาศ", - "\u0e32": "สระ อา", - "\u0e33": "สระ อำ", - "\u0e34": "สระ อิ", - "\u0e35": "สระ อี", - "\u0e36": "สระ อึ", - "\u0e37": "สระ อือ", - "\u0e38": "สระ อุ", - "\u0e39": "สระ อู", - "\u0e40": "สระ เอ", - "\u0e41": "สระ แอ", - "\u0e42": "สระ โอ", - "\u0e43": "สระ ใอ", - "\u0e44": "สระ ไอ", - "\u0e45": "ไม้ม้วน", - "\u0e4d": "นฤคหิต", - "\u0e47": "ไม้ไต่คู้", - # Tone Marks - "\u0e48": "ไม้เอก", - "\u0e49": "ไม้โท", - "\u0e4a": "ไม้ตรี", - "\u0e4b": "ไม้จัตวา", - # Other Signs - "\u0e2f": "ไปยาลน้อย", - "\u0e3a": "พินทุ", - "\u0e46": "ไม้ยมก", - "\u0e4c": "การันต์", - "\u0e4e": "ยามักการ", - # Punctuation - "\u0e4f": "ฟองมัน", - "\u0e5a": "อังคั่นคู่", - "\u0e5b": "โคมุต", - # Digits - **{char: char for char in thai_digits}, - # Symbol - "\u0e3f": "฿", -}) +_THAI_CHAR_NAMES: MappingProxyType[str, str] = MappingProxyType( + { + # Consonants + **{char: char for char in thai_consonants}, + # Vowels and Signs + "\u0e24": "ฤ", + "\u0e26": "ฦ", + "\u0e30": "สระ อะ", + "\u0e31": "ไม้หันอากาศ", + "\u0e32": "สระ อา", + "\u0e33": "สระ อำ", + "\u0e34": "สระ อิ", + "\u0e35": "สระ อี", + "\u0e36": "สระ อึ", + "\u0e37": "สระ อือ", + "\u0e38": "สระ อุ", + "\u0e39": "สระ อู", + "\u0e40": "สระ เอ", + "\u0e41": "สระ แอ", + "\u0e42": "สระ โอ", + "\u0e43": "สระ ใอ", + "\u0e44": "สระ ไอ", + "\u0e45": "ไม้ม้วน", + "\u0e4d": "นฤคหิต", + "\u0e47": "ไม้ไต่คู้", + # Tone Marks + "\u0e48": "ไม้เอก", + "\u0e49": "ไม้โท", + "\u0e4a": "ไม้ตรี", + "\u0e4b": "ไม้จัตวา", + # Other Signs + "\u0e2f": "ไปยาลน้อย", + "\u0e3a": "พินทุ", + "\u0e46": "ไม้ยมก", + "\u0e4c": "การันต์", + "\u0e4e": "ยามักการ", + # Punctuation + "\u0e4f": "ฟองมัน", + "\u0e5a": "อังคั่นคู่", + "\u0e5b": "โคมุต", + # Digits + **{char: char for char in thai_digits}, + # Symbol + "\u0e3f": "฿", + } +) def is_thai_char(ch: str) -> bool: @@ -322,7 +324,7 @@ def thai_word_tone_detector(word: Optional[str]) -> list[tuple[str, str]]: return [(i, tone_detector(i.replace("หฺ", "ห"))) for i in _pronunciate] -def count_thai_chars(text: str) -> dict: +def count_thai_chars(text: str) -> dict[str, int]: """Count Thai characters by type This function will give you numbers of Thai characters by type\ @@ -331,7 +333,7 @@ def count_thai_chars(text: str) -> dict: :param str text: Text :return: Dict with numbers of Thai characters by type - :rtype: dict + :rtype: dict[str, int] :Example: :: @@ -392,7 +394,7 @@ def count_thai_chars(text: str) -> dict: return _dict -def analyze_thai_text(text: str) -> dict: +def analyze_thai_text(text: str) -> dict[str, int]: """Analyzes a string of Thai text and returns a dictionaries, where each values represents a single classified character from the text. @@ -400,7 +402,7 @@ def analyze_thai_text(text: str) -> dict: character to its descriptive name or itself (for consonants and digits). :param str text: The Thai text string to be analyzed. - :rtype: dict + :rtype: dict[str, int] :return: A dictionaries, with each item containing a single character and a count of 1. diff --git a/pythainlp/wangchanberta/core.py b/pythainlp/wangchanberta/core.py index f5b028b3c..40a9aeec1 100644 --- a/pythainlp/wangchanberta/core.py +++ b/pythainlp/wangchanberta/core.py @@ -5,7 +5,7 @@ import re import warnings -from typing import TYPE_CHECKING, Optional, Union +from typing import TYPE_CHECKING, Optional, Union, cast if TYPE_CHECKING: from transformers import ( @@ -269,4 +269,4 @@ def segment(text: str) -> list[str]: if not text or not isinstance(text, str): return [] - return _get_tokenizer().tokenize(text) # type: ignore[no-any-return] + return cast(list[str], _get_tokenizer().tokenize(text)) diff --git a/pythainlp/word_vector/core.py b/pythainlp/word_vector/core.py index 560c3adb2..37c2f2794 100644 --- a/pythainlp/word_vector/core.py +++ b/pythainlp/word_vector/core.py @@ -5,14 +5,15 @@ import logging from collections.abc import Callable -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from pythainlp.corpus import get_corpus_path from pythainlp.tokenize import thai2fit_tokenizer, word_tokenize if TYPE_CHECKING: + import numpy as np from gensim.models.keyedvectors import Word2VecKeyedVectors - from numpy import ndarray + from numpy.typing import NDArray WV_DIM: int = 300 # word vector dimension @@ -135,7 +136,7 @@ def doesnt_match(self, words: list[str]) -> str: >>> wv.doesnt_match(words) 'เรือ' """ - return self.model.doesnt_match(words) # type: ignore[no-any-return] + return cast(str, self.model.doesnt_match(words)) def most_similar_cosmul( self, positive: list[str], negative: list[str] @@ -236,8 +237,11 @@ def most_similar_cosmul( >>> wv.most_similar_cosmul(list_positive, list_negative) KeyError: "word 'เมนูอาหารไทย' not in vocabulary" """ - return self.model.most_similar_cosmul( # type: ignore[no-any-return] - positive=positive, negative=negative + return cast( + list[tuple[str, float]], + self.model.most_similar_cosmul( + positive=positive, negative=negative + ), ) def similarity(self, word1: str, word2: str) -> float: @@ -276,9 +280,11 @@ def similarity(self, word1: str, word2: str) -> float: 0.04300258 """ - return self.model.similarity(word1, word2) # type: ignore[no-any-return] + return cast(float, self.model.similarity(word1, word2)) - def sentence_vectorizer(self, text: str, use_mean: bool = True) -> ndarray: + def sentence_vectorizer( + self, text: str, use_mean: bool = True + ) -> "NDArray[np.float32]": """Converts a Thai sentence into a vector. Specifically, it first tokenizes that text and maps each tokenized word with the word vectors from the model. @@ -290,9 +296,9 @@ def sentence_vectorizer(self, text: str, use_mean: bool = True) -> ndarray: word vectors. Otherwise, aggregate with summation of all word vectors - :return: 300-dimension vector representing the given sentence - in form of :mod:`numpy` array - :rtype: numpy.ndarray + :return: a :class:`numpy.ndarray` of dtype ``numpy.float32`` and + shape ``(1, 300)`` representing the given sentence + :rtype: numpy.typing.NDArray[numpy.float32] :Example: @@ -321,9 +327,9 @@ def sentence_vectorizer(self, text: str, use_mean: bool = True) -> ndarray: 0.40506999, 1.58591403, 0.63869202, -0.702155 , 1.62977601, 4.52269109, -0.70760502, 0.50952601, -0.914392 , 0.70673105]]) """ - from numpy import zeros + import numpy as np - vec = zeros((1, self.WV_DIM)) + vec = np.zeros((1, self.WV_DIM), dtype=np.float32) words = self.tokenize(text) len_words = len(words) diff --git a/pythainlp/wsd/core.py b/pythainlp/wsd/core.py index e64a47342..375d9e135 100644 --- a/pythainlp/wsd/core.py +++ b/pythainlp/wsd/core.py @@ -3,17 +3,14 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from typing import TYPE_CHECKING, Optional, Union, cast +from typing import Optional, Union, cast from pythainlp.corpus import thai_wsd_dict from pythainlp.tokenize import Tokenizer from pythainlp.util.trie import Trie -if TYPE_CHECKING: - from typing import Any - _wsd_dict: dict[str, Union[list[str], list[list[str]]]] = thai_wsd_dict() -_mean_all: dict[str, Any] = {} +_mean_all: dict[str, list[str]] = {} _all_word: set[str] = cast("set[str]", set(_mean_all.keys())) _TRIE: Trie = Trie(_all_word) @@ -21,12 +18,12 @@ words: list[str] = cast("list[str]", _wsd_dict["word"]) meanings: list[list[str]] = cast("list[list[str]]", _wsd_dict["meaning"]) -i: str -j: list[str] -for i, j in zip(words, meanings): - _mean_all[i] = j +i_word: str +i_meanings: list[str] +for i_word, i_meanings in zip(words, meanings): + _mean_all[i_word] = i_meanings -_MODEL: Optional[Any] = None +_MODEL_CACHE: dict[str, _SentenceTransformersModel] = {} class _SentenceTransformersModel: @@ -54,14 +51,16 @@ def get_score(self, sentences1: str, sentences2: str) -> float: embedding_1 = self.model.encode(sentences1, convert_to_tensor=True) embedding_2 = self.model.encode(sentences2, convert_to_tensor=True) - return 1 - util.pytorch_cos_sim(embedding_1, embedding_2)[0][0].item() # type: ignore[no-any-return] + return float( + 1 - util.pytorch_cos_sim(embedding_1, embedding_2)[0][0].item() + ) def get_sense( sentence: str, word: str, device: str = "cpu", - custom_dict: Optional[dict] = None, + custom_dict: Optional[dict[str, list[str]]] = None, custom_tokenizer: Tokenizer = _word_cut, ) -> list[tuple[str, float]]: """Get word sense from the sentence. @@ -70,7 +69,8 @@ def get_sense( :param str sentence: Thai sentence :param str word: Thai word :param str device: device for running model on. - :param dict custom_dict: Thai dictionary {"word":["definition",..]} + :param Optional[dict[str, list[str]]] custom_dict: Thai dictionary in the + form {"word": ["definition", ...]} :param Tokenizer custom_tokenizer: Tokenizer used to tokenize words in \ sentence. :return: a list of definitions and distances (1 - cos_sim) or \ @@ -107,7 +107,6 @@ def get_sense( # ('ชื่อขนมชนิดหนึ่งจำพวกขนมเค้ก แต่ทำเป็นชิ้นเล็ก ๆ แบน ๆ แล้วอบให้กรอบ', # 0.12473666667938232)] """ - global _MODEL if not custom_dict: custom_dict = _mean_all @@ -115,24 +114,26 @@ def get_sense( if word not in set(custom_dict.keys()) or word not in sentence: return [] - if not _MODEL: - _MODEL = _SentenceTransformersModel(device=device) - if _MODEL.device != device: - _MODEL.change_device(device=device) + if device not in _MODEL_CACHE: + _MODEL_CACHE[device] = _SentenceTransformersModel(device=device) + + model = _MODEL_CACHE[device] temp_mean = custom_dict[word] - temp = [] - for i in temp_mean: - _temp_2 = [] - for j in w: - if j == word: - j = ( + temp: list[tuple[str, float]] = [] + for meaning in temp_mean: + tokens_with_sense: list[str] = [] + for token in w: + if token == word: + token = ( word + f" ({word} ความหมาย '" - + i.replace("(", "").replace(")", "") + + meaning.replace("(", "").replace(")", "") + "') " ) - _temp_2.append(j) - temp.append((i, _MODEL.get_score(sentence, "".join(_temp_2)))) + tokens_with_sense.append(token) + temp.append( + (meaning, model.get_score(sentence, "".join(tokens_with_sense))) + ) return temp diff --git a/tests/core/test_util.py b/tests/core/test_util.py index 35f34b055..a5ba83938 100644 --- a/tests/core/test_util.py +++ b/tests/core/test_util.py @@ -123,7 +123,8 @@ def test_number(self): ) self.assertEqual(bahttext(116), "หนึ่งร้อยสิบหกบาทถ้วน") self.assertEqual(bahttext(0), "ศูนย์บาทถ้วน") - self.assertEqual(bahttext(None), "") + with self.assertRaises(TypeError): + bahttext(None) # type: ignore[arg-type] # Edge cases: negative number self.assertEqual(bahttext(-100), "ลบหนึ่งร้อยบาทถ้วน") self.assertEqual(bahttext(-50.50), "ลบห้าสิบบาทห้าสิบสตางค์") diff --git a/tests/extra/testx_corpus.py b/tests/extra/testx_corpus.py index adac2c622..e36fd4418 100644 --- a/tests/extra/testx_corpus.py +++ b/tests/extra/testx_corpus.py @@ -27,6 +27,9 @@ def test_wordnet(self): self.assertIsNotNone(wordnet.lemma("cat.n.01.cat")) self.assertEqual(wordnet.morphy("dogs"), "dog") + # morphy with pos=None (default) and with explicit pos + self.assertEqual(wordnet.morphy("dogs", pos=None), "dog") + self.assertEqual(wordnet.morphy("dogs", pos="n"), "dog") bird = wordnet.synset("bird.n.01") mouse = wordnet.synset("mouse.n.01")