e2k 0.1.0__py3-none-any.whl

e2k/__init__.py

@@ -0,0 +1,8 @@
+from .inference import P2K, C2K
+from .constants import en_phones, kanas, ascii_entries
+import os.path
+
+# make internal modules invisible
+__path__ = [os.path.dirname(__file__)]
+# re-export the public API
+__all__ = ["P2K", "C2K", "en_phones", "kanas", "ascii_entries"]

e2k/constants.py

@@ -0,0 +1,171 @@
+from string import ascii_letters
+
+
+PAD_IDX = 0
+SOS_IDX = 1
+EOS_IDX = 2
+
+kanas = [
+    "ァ",
+    "ア",
+    "ィ",
+    "イ",
+    "ゥ",
+    "ウ",
+    "ェ",
+    "エ",
+    "ォ",
+    "オ",
+    "カ",
+    "ガ",
+    "キ",
+    "ギ",
+    "ク",
+    "グ",
+    "ケ",
+    "ゲ",
+    "コ",
+    "ゴ",
+    "サ",
+    "ザ",
+    "シ",
+    "ジ",
+    "ス",
+    "ズ",
+    "セ",
+    "ゼ",
+    "ソ",
+    "ゾ",
+    "タ",
+    "ダ",
+    "チ",
+    "ヂ",
+    "ッ",
+    "ツ",
+    "ヅ",
+    "テ",
+    "デ",
+    "ト",
+    "ド",
+    "ナ",
+    "ニ",
+    "ヌ",
+    "ネ",
+    "ノ",
+    "ハ",
+    "バ",
+    "パ",
+    "ヒ",
+    "ビ",
+    "ピ",
+    "フ",
+    "ブ",
+    "プ",
+    "ヘ",
+    "ベ",
+    "ペ",
+    "ホ",
+    "ボ",
+    "ポ",
+    "マ",
+    "ミ",
+    "ム",
+    "メ",
+    "モ",
+    "ャ",
+    "ヤ",
+    "ュ",
+    "ユ",
+    "ョ",
+    "ヨ",
+    "ラ",
+    "リ",
+    "ル",
+    "レ",
+    "ロ",
+    "ヮ",
+    "ワ",
+    "ヰ",
+    "ヱ",
+    "ヲ",
+    "ン",
+    "ヴ",
+    "ー",
+]
+
+en_phones = [
+    "AA0",
+    "AA1",
+    "AA2",
+    "AE0",
+    "AE1",
+    "AE2",
+    "AH0",
+    "AH1",
+    "AH2",
+    "AO0",
+    "AO1",
+    "AO2",
+    "AW0",
+    "AW1",
+    "AW2",
+    "AY0",
+    "AY1",
+    "AY2",
+    "B",
+    "CH",
+    "D",
+    "DH",
+    "EH0",
+    "EH1",
+    "EH2",
+    "ER0",
+    "ER1",
+    "ER2",
+    "EY0",
+    "EY1",
+    "EY2",
+    "F",
+    "G",
+    "HH",
+    "IH0",
+    "IH1",
+    "IH2",
+    "IY0",
+    "IY1",
+    "IY2",
+    "JH",
+    "K",
+    "L",
+    "M",
+    "N",
+    "NG",
+    "OW0",
+    "OW1",
+    "OW2",
+    "OY0",
+    "OY1",
+    "OY2",
+    "P",
+    "R",
+    "S",
+    "SH",
+    "T",
+    "TH",
+    "UH0",
+    "UH1",
+    "UH2",
+    "UW0",
+    "UW1",
+    "UW2",
+    "V",
+    "W",
+    "Y",
+    "Z",
+    "ZH",
+]
+
+ascii_entries = ["<pad>", "<sos>", "<eos>"] + list(ascii_letters) + [" "]
+
+kanas = ["<pad>", "<sos>", "<eos>"] + kanas
+en_phones = ["<pad>", "<sos>", "<eos>"] + en_phones + [" "]

e2k/inference.py

@@ -0,0 +1,272 @@
+# Description: Inference functions for the E2K model in numpy
+import numpy as np
+from typing import Optional, Dict
+import importlib.resources
+from e2k.constants import SOS_IDX, EOS_IDX, en_phones, kanas, ascii_entries
+
+
+class Linear:
+    def __init__(self, weight, bias):
+        self.weight = weight
+        self.bias = bias
+
+    def forward(self, x):
+        return np.matmul(x, self.weight.T) + self.bias
+
+
+class Embedding:
+    def __init__(self, weight):
+        self.weight = weight
+
+    def forward(self, x):
+        return self.weight[x]
+
+
+def sigmoid(x):
+    return 1 / (1 + np.exp(-x))
+
+
+def tanh(x):
+    return np.tanh(x)
+
+
+class GRUCell:
+    def __init__(self, weight_ih, weight_hh, bias_ih, bias_hh):
+        self.ih = Linear(weight_ih, bias_ih)
+        self.hh = Linear(weight_hh, bias_hh)
+
+    def forward(self, x: np.ndarray, h: Optional[np.ndarray] = None):
+        """
+        x: [D]
+        h: [D]
+        """
+        if h is None:
+            h = np.zeros(
+                [
+                    self.hh.weight.shape[-1],
+                ]
+            )
+        rzn_ih = self.ih.forward(x)
+        rzn_hh = self.hh.forward(h)
+
+        rz_ih, n_ih = (
+            rzn_ih[: rzn_ih.shape[-1] * 2 // 3],
+            rzn_ih[rzn_ih.shape[-1] * 2 // 3 :],
+        )
+        rz_hh, n_hh = (
+            rzn_hh[: rzn_hh.shape[-1] * 2 // 3],
+            rzn_hh[rzn_hh.shape[-1] * 2 // 3 :],
+        )
+
+        rz = sigmoid(rz_ih + rz_hh)
+        r, z = np.split(rz, 2, axis=-1)
+
+        n = np.tanh(n_ih + r * n_hh)
+        h = (1 - z) * n + z * h
+
+        return h
+
+
+class GRU:
+    def __init__(self, cell: GRUCell, reverse: bool = False):
+        self.cell = cell
+        self.reverse = reverse
+
+    def forward(self, x, h: Optional[np.ndarray] = None):
+        """
+        x: [T, D], unbatched
+        """
+        if self.reverse:
+            x = np.flip(x, axis=0)
+        outputs = []
+        for i in range(x.shape[0]):
+            h = self.cell.forward(x[i], h)
+            outputs.append(h)
+        outputs = np.stack(outputs)
+        if self.reverse:
+            outputs = np.flip(outputs, axis=0)
+        return outputs, h
+
+
+class MHA:
+    """
+    Multi-head attention
+    """
+
+    def __init__(
+        self,
+        in_proj_weight,
+        in_proj_bias,
+        out_proj_weight,
+        out_proj_bias,
+        num_heads: int,
+    ):
+        [q_w, k_w, v_w] = np.split(in_proj_weight, 3, axis=0)
+        [q_b, k_b, v_b] = np.split(in_proj_bias, 3, axis=0)
+        self.dim = q_w.shape[-1]
+        self.q_proj = Linear(q_w, q_b)
+        self.k_proj = Linear(k_w, k_b)
+        self.v_proj = Linear(v_w, v_b)
+        self.o_proj = Linear(out_proj_weight, out_proj_bias)
+        self.num_heads = num_heads
+        self.d_heads = self.dim // num_heads
+        self.scale = np.sqrt(self.dim)
+
+    def forward(self, q: np.ndarray, k: np.ndarray, v: np.ndarray):
+        """
+        q: [T, D]
+        k: [T, D]
+        v: [T, D]
+        """
+        q = self.q_proj.forward(q)
+        k = self.k_proj.forward(k)
+        v = self.v_proj.forward(v)
+        q = np.split(q, self.num_heads, axis=-1)
+        q = np.stack(q, axis=0)
+        k = np.split(k, self.num_heads, axis=-1)
+        k = np.stack(k, axis=0)
+        v = np.split(v, self.num_heads, axis=-1)
+        v = np.stack(v, axis=0)
+        attn = np.matmul(q, np.transpose(k, (0, 2, 1)))
+        attn = attn / self.scale
+        attn = np.exp(attn)
+        attn = attn / attn.sum(axis=-1, keepdims=True)
+        o = np.matmul(attn, v)
+        o = np.transpose(o, (1, 0, 2))
+        o = o.reshape([o.shape[0], -1])
+        return self.o_proj.forward(o)
+
+
+class S2S:
+    def __init__(self, weights: Dict[str, np.ndarray], max_len: int = 16):
+        self.e_emb = Embedding(weights["e_emb.weight"])
+        self.k_emb = Embedding(weights["k_emb.weight"])
+        self.encoder = GRU(
+            GRUCell(
+                weights["encoder.weight_ih_l0"],
+                weights["encoder.weight_hh_l0"],
+                weights["encoder.bias_ih_l0"],
+                weights["encoder.bias_hh_l0"],
+            )
+        )
+        self.encoder_reverse = GRU(
+            GRUCell(
+                weights["encoder.weight_ih_l0_reverse"],
+                weights["encoder.weight_hh_l0_reverse"],
+                weights["encoder.bias_ih_l0_reverse"],
+                weights["encoder.bias_hh_l0_reverse"],
+            ),
+            reverse=True,
+        )
+        self.encoder_fc = Linear(
+            weights["encoder_fc.0.weight"], weights["encoder_fc.0.bias"]
+        )
+        self.pre_decoder = GRU(
+            GRUCell(
+                weights["pre_decoder.weight_ih_l0"],
+                weights["pre_decoder.weight_hh_l0"],
+                weights["pre_decoder.bias_ih_l0"],
+                weights["pre_decoder.bias_hh_l0"],
+            )
+        )
+        self.post_decoder = GRU(
+            GRUCell(
+                weights["post_decoder.weight_ih_l0"],
+                weights["post_decoder.weight_hh_l0"],
+                weights["post_decoder.bias_ih_l0"],
+                weights["post_decoder.bias_hh_l0"],
+            )
+        )
+        self.attn = MHA(
+            weights["attn.in_proj_weight"],
+            weights["attn.in_proj_bias"],
+            weights["attn.out_proj.weight"],
+            weights["attn.out_proj.bias"],
+            4,
+        )
+        self.fc = Linear(weights["fc.weight"], weights["fc.bias"])
+        self.max_len = max_len
+
+    def greedy(self, step_dec: np.ndarray):
+        """
+        step_dec: [N]
+        """
+        return np.argmax(step_dec, axis=-1)
+
+    def forward(self, src):
+        """
+        In numpy, only inference is supported.
+        """
+        e_emb = self.e_emb.forward(src)
+        enc_out, _ = self.encoder.forward(e_emb)
+        enc_out_rev, _ = self.encoder_reverse.forward(e_emb)
+        enc_out = np.concatenate([enc_out, enc_out_rev], axis=-1)
+        enc_out = self.encoder_fc.forward(enc_out)
+        enc_out = tanh(enc_out)
+        res = [SOS_IDX]
+        h1 = None
+        h2 = None
+        for _ in range(self.max_len):
+            dec_emb = self.k_emb.forward(np.array([res[-1]]))
+            dec_out, h1 = self.pre_decoder.forward(dec_emb, h1)
+            attn_out = self.attn.forward(dec_out, enc_out, enc_out)
+            x = np.concatenate([dec_out, attn_out], axis=-1)
+            x, h2 = self.post_decoder.forward(x, h2)
+            x = self.fc.forward(x)
+            x = x[0]
+            res.append(self.greedy(x))
+            if res[-1] == EOS_IDX:
+                break
+        return res
+
+
+class BaseE2K:
+    def __init__(self):
+        self.s2s = None
+        self.in_table = None
+        self.out_table = None
+        raise NotImplementedError
+
+    def __call__(self, src: str) -> str:
+        src = list(filter(lambda x: x in self.in_table, src))
+        src = [self.in_table[c] for c in src]
+        src = [SOS_IDX] + src + [EOS_IDX]
+        src = np.array(src)
+        tgt = self.s2s.forward(src)
+        tgt = [int(v) for v in tgt[1:-1]]
+        tgt = [self.out_table[c] for c in tgt]
+        return "".join(tgt)
+
+def get_weight_path(filename) -> str:
+    return importlib.resources.files("e2k.models").joinpath(filename)
+
+class P2K(BaseE2K):
+    def __init__(self, max_len: int = 16):
+        weights = np.load(get_weight_path("model-p2k.npz"), allow_pickle=True)
+        self.s2s = S2S(weights, max_len)
+        self.in_table = {c: i for i, c in enumerate(en_phones)}
+        self.out_table = {i: c for i, c in enumerate(kanas)}
+
+
+class C2K(BaseE2K):
+    def __init__(self, max_len: int = 16):
+        weights = np.load(get_weight_path("model-c2k.npz"), allow_pickle=True)
+        self.s2s = S2S(weights, max_len)
+        self.in_table = {c: i for i, c in enumerate(ascii_entries)}
+        self.out_table = {i: c for i, c in enumerate(kanas)}
+
+
+if __name__ == "__main__":
+    import argparse
+    from g2p_en import G2p
+
+    g2p = G2p()
+    parser = argparse.ArgumentParser()
+    args = parser.parse_args()
+    p2k = P2K()
+    c2k = C2K()
+    word = "voicevox"
+    phonemes = g2p(word)
+    print(word)
+    print("P2K: ", p2k(phonemes)) # ザワールド
+    print("C2K: ", c2k(word)) # ザワールド

e2k-0.1.0.dist-info/METADATA

@@ -0,0 +1,147 @@
+Metadata-Version: 2.2
+Name: e2k
+Version: 0.1.0
+Summary: A tool for automatic English to Katakana conversion
+Author-email: Patchethium <asguftuikh@gmail.com>
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=2.2.2
+
+# English to Katakana Dictionary
+
+It's an extraction of the Wikitionary's Japanese dictionary and `JMdict / EDICT` to create a dictionary of English words and their Katakana representation.
+
+We also provide a GRU model to convert English to Katakana without a dictionary.
+
+```json
+{"word": "word", "katakana": "ワード"}
+{"word": "another", "katakana": "アナザー"}
+{"word": "example", "katakana": "エグザンプル"}
+```
+
+It's for the purpose of converting English words to Katakana and further use in Japanese Text-to-Speech applications or Japanese learning.
+
+## Re-creating the dictionary
+
+The script has zero dependencies, as long as you have a Python 3 interpreter it should work.
+
+### Download the dictionary
+
+Download the raw dump of the Japanese Wikitionary from https://kaikki.org/dictionary/rawdata.html, they kindly provide the parsed data in a JSONL format.
+
+Look for the `Japanese ja-extract.jsonl.gz (compressed 37.5MB)` entry and download it. If you prefer command line, use
+
+```bash
+curl -O https://kaikki.org/dictionary/downloads/ja/ja-extract.jsonl.gz
+```
+
+### Extract it
+
+Extract it into `/vendor` folder.
+
+On Linux, you can use
+
+```bash
+gzip -d ja-extract.jsonl.gz
+```
+
+### Run the script
+
+```bash
+python extract.py
+# if you have another name for the file
+python extract.py --path /path/to/your_file.jsonl
+```
+
+And a `katakana_dict.jsonl` file will be created in the root folder.
+
+## E2K model
+
+We also provide a seq2seq GRU model to convert English automatically to Katakana. It's trained on the aforementioned dictionary.
+
+We provide 2 types of models, one converts phoneme to Katakana and one that converts character to Katakana. Choose the one that fits your use case.
+
+```bash
+pip install e2k
+```
+
+The weights are included in PyPI, you can simply use the model like this:
+
+```python
+from e2k import P2K, C2K
+from g2p_en import G2p # or `phonemizer`, as long as it outputs CMUdict phoneme
+
+# cmudict phoneme to katakana
+p2k = P2K()
+
+g2p = G2p()
+
+katakana = p2k(g2p("word"))
+
+print(katakana) # "ワード"
+
+# character directly to katakana
+c2k = C2K()
+
+# since it's an autoregressive model, you can set the decoding strategy
+# to greedy, top-k or nucleus sampling (top-p)
+# see https://huggingface.co/docs/transformers/en/generation_strategies
+c2k.set_decoding_strategy("top-p", p=0.3)
+
+katakana = c2k("word")
+
+print(katakana) # "ワード"
+```
+
+We rewrite the inference of GRU model in `numpy`, minimizing the dependencies to `numpy` only.
+
+### Benchmark
+
+```bash
+# --p2k for phoneme to katakana, if not provided, it will be character to katakana
+python eval.py --data ./vendor/katakana_dict.jsonl --model /path/to/your/model.pth --p2k
+```
+
+| Model                 | BLEU Score |
+| --------------------- | ---------- |
+| Phoneme to Katakana   | 0.85       |
+| Character to Katakana | 0.90       |
+
+### Training the model
+
+You'll need `torch` and `g2p_en`. After that, you can run the `train.py` script to train the model.
+
+```bash
+python train.py --data ./vendor/kanji_dict.jsonl
+```
+
+It takes around 10 minutes to train the model on a desktop CPU. The model will be saved as `model.pth` in the root folder.
+
+Also, you'll need to either download the `kanji_dict.jsonl` from the releases or create it yourself using the `extract.py` script.
+
+Be noted that the training script is not included in the PyPI package, you'll need to clone the repository to train the model.
+
+## Development
+
+### Install the dependencies
+
+I use [`uv`](https://docs.astral.sh/uv/) to manage the dependencies and publish the package.
+
+```bash
+uv sync
+```
+
+## License
+
+The code is released under WTFPL.
+
+The dictionary should follow the [Wikimedia's license](https://dumps.wikimedia.org/legal.html) and the [JMdict / EDICT's Copyright](https://www.edrdg.org/) license.
+
+In short, as long as you provide the attribution, you can use it for both commercial and non-commercial purposes.
+
+The model weights are released under the same license as the dictionary.
+
+## Credits
+
+- [Wikitionary](https://www.wiktionary.org/)
+- [JMdict / EDICT](http://www.edrdg.org/jmdict/edict.html)

e2k-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+e2k/__init__.py,sha256=KxTSqeTPYlyYhjozlcruFWtgn9iODsLYin7cY9QKPBI,267
+e2k/constants.py,sha256=Lajx--JL1kduvNN3BIbr9soQ_A9T-jPEk-jdJFthKQw,1935
+e2k/inference.py,sha256=g1OG68x4I26u2uii-NNzISG4_1la_1eLmf4y9UtVMzs,8163
+e2k/models/model-c2k.npz,sha256=6Fi5I4USkgYJ-V4QiEbWX2_b_ZTLn3O7qvOKiZgsBZE,4466282
+e2k/models/model-p2k.npz,sha256=jqUjk_MsGX7kc882Pkwf6GaBm2vdHnhFtu6NKrYPe-o,4474986
+e2k-0.1.0.dist-info/METADATA,sha256=mq0f2ofBy2dZKbtgS-GMdhG_u_uzQUJwJB_5rBY9E20,4330
+e2k-0.1.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+e2k-0.1.0.dist-info/top_level.txt,sha256=C1ChHRXMXLY8K9ZIwiRtchhn1MW-fW9gY1HlFjRLCYo,4
+e2k-0.1.0.dist-info/RECORD,,

e2k-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

e2k-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+e2k