g2ps-torch 0.1.0__py3-none-any.whl

g2ps_torch/__init__.py

@@ -0,0 +1 @@
+from .runtime import G2pS

g2ps_torch/_resolve.py

@@ -0,0 +1,35 @@
+"""Tier auto-resolution: packaged weights first, then on-demand download."""
+import os
+import shutil
+import urllib.request
+
+MODELS_DIR = os.path.join(os.path.dirname(__file__), 'models')
+CACHE_DIR = os.environ.get('G2PS_CACHE', os.path.join(os.path.expanduser('~'), '.cache', 'g2ps'))
+
+
+def resolve_model_dir(model_dir, arch_keys, weights_name, urls):
+    """model_dir may be a path, or a tier name ('nano'/'tiny'/'small')."""
+    if os.path.isdir(model_dir) and os.path.exists(os.path.join(model_dir, weights_name)):
+        return model_dir
+    tier = model_dir if model_dir in arch_keys else os.path.basename(os.path.normpath(model_dir))
+    if tier not in arch_keys:
+        raise ValueError(f"model_dir must be a path or one of {list(arch_keys)}, got {model_dir!r}")
+    pkg = os.path.join(MODELS_DIR, tier)
+    if os.path.exists(os.path.join(pkg, weights_name)):
+        return pkg
+    dst = os.path.join(CACHE_DIR, tier)
+    os.makedirs(dst, exist_ok=True)
+    if not os.path.exists(os.path.join(dst, 'meta.pkl')):
+        shutil.copy(os.path.join(pkg, 'meta.pkl'), os.path.join(dst, 'meta.pkl'))
+    wp = os.path.join(dst, weights_name)
+    if not os.path.exists(wp):
+        err = []
+        for url in urls:
+            try:
+                print(f'g2pS: downloading {url.rsplit("/", 1)[-1]} ({tier}) ...')
+                urllib.request.urlretrieve(url, wp)
+                return dst
+            except Exception as e:  # noqa: BLE001
+                err.append(str(e)[:80])
+        raise RuntimeError('could not download weights from any mirror:\n' + '\n'.join(err))
+    return dst

g2ps_torch/modules.py

@@ -0,0 +1,192 @@
+"""Student model: modern small char-level encoder.
+
+6 x 384 bidirectional transformer with RoPE / RMSNorm / SwiGLU (2024-style recipe),
+candidate-masked readout at every labeled polyphone position.
+"""
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        ms = x.float().pow(2).mean(-1, keepdim=True)          # fp32 accumulation (fp16-safe)
+        return (self.w * (x.float() * torch.rsqrt(ms + self.eps))).to(x.dtype)
+
+
+def rope_table(max_len, dim, base=10000.0, device='cpu', dtype=torch.float32):
+    inv = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float64) / dim))
+    t = torch.arange(max_len, device=device, dtype=torch.float64)
+    angles = torch.outer(t, inv)
+    return torch.stack([torch.cos(angles), torch.sin(angles)], dim=-1).to(dtype)  # (T, D/2, 2)
+
+
+def apply_rope(x, rope):
+    """x: (B, T, H, Dh); rope: (T, Dh/2, 2)"""
+    T = x.shape[1]
+    r = rope[:T].to(x.dtype)  # (T, Dh/2, 2)
+    x1, x2 = x[..., 0::2], x[..., 1::2]
+    cos = r[..., 0].unsqueeze(0).unsqueeze(2)  # (1, T, 1, Dh/2)
+    sin = r[..., 1].unsqueeze(0).unsqueeze(2)
+    o1 = x1 * cos - x2 * sin
+    o2 = x1 * sin + x2 * cos
+    return torch.stack((o1, o2), dim=-1).flatten(-2)
+
+
+class Attention(nn.Module):
+    def __init__(self, d, n_heads):
+        super().__init__()
+        self.n_heads = n_heads
+        self.dh = d // n_heads
+        self.qkv = nn.Linear(d, 3 * d, bias=False)
+        self.o = nn.Linear(d, d, bias=False)
+
+    def forward(self, x, mask, rope):
+        B, T, D = x.shape
+        q, k, v = self.qkv(x).split(D, dim=2)
+        q = q.view(B, T, self.n_heads, self.dh).transpose(1, 2)
+        k = k.view(B, T, self.n_heads, self.dh).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.dh).transpose(1, 2)
+        q, k = apply_rope(q, rope), apply_rope(k, rope)
+        # mask: (B, T) bool True=valid -> (B,1,1,T) additive
+        att_mask = ((~mask).unsqueeze(1).unsqueeze(2) * -1e4).to(q.dtype)  # (B,1,1,T)
+        y = F.scaled_dot_product_attention(q, k, v, attn_mask=att_mask)
+        y = y.transpose(1, 2).contiguous().view(B, T, D)
+        return self.o(y)
+
+
+class ConvStem(nn.Module):
+    """depthwise conv k=3 with dilations 1/2/3 -> fused local n-gram features. ~3.5K params."""
+    def __init__(self, d):
+        super().__init__()
+        self.convs = nn.ModuleList([nn.Conv1d(d, d, 3, padding=dl, dilation=dl, groups=d) for dl in (1, 2, 3)])
+        self.norm = RMSNorm(d)
+
+    def forward(self, x):
+        # x: (B, T, D); padding-aware: zero pad chars already are 0-emb? mask conv beyond seq via padding zeros
+        h = x.transpose(1, 2)  # (B, D, T)
+        hs = [c(h) for c in self.convs]
+        out = x + sum(hs).transpose(1, 2)
+        return self.norm(out)
+
+class SwiGLU(nn.Module):
+    def __init__(self, d, hidden):
+        super().__init__()
+        self.w1 = nn.Linear(d, hidden, bias=False)
+        self.w2 = nn.Linear(d, hidden, bias=False)
+        self.w3 = nn.Linear(hidden, d, bias=False)
+
+    def forward(self, x):
+        return self.w3(F.silu(self.w1(x)) * self.w2(x))
+
+
+class Block(nn.Module):
+    def __init__(self, d, n_heads, ffn_hidden):
+        super().__init__()
+        self.n1 = RMSNorm(d)
+        self.attn = Attention(d, n_heads)
+        self.n2 = RMSNorm(d)
+        self.ffn = SwiGLU(d, ffn_hidden)
+
+    def forward(self, x, mask, rope):
+        x = x + self.attn(self.n1(x), mask, rope)
+        x = x + self.ffn(self.n2(x))
+        return x
+
+
+class StudentG2P(nn.Module):
+    def __init__(self, vocab_size, n_labels, n_poly_chars, d=384, n_layers=6,
+                 n_heads=6, ffn_hidden=1024, max_len=128, dropout=0.1, use_conv_stem=True,
+                 emb_dim=None, bias_rank=0, n_passes=1):
+        self.n_passes = n_passes
+        super().__init__()
+        rope = rope_table(max_len, d // n_heads)
+        self.register_buffer('rope', rope, persistent=False)
+        emb_dim = emb_dim or d
+        self.emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
+        self.emb_proj = nn.Linear(emb_dim, d, bias=False) if emb_dim != d else nn.Identity()
+        self.stem = ConvStem(d) if use_conv_stem else nn.Identity()
+        self.drop = nn.Dropout(dropout)
+        self.blocks = nn.ModuleList([Block(d, n_heads, ffn_hidden)
+                                     for _ in range(n_layers)])
+        self.n_final = RMSNorm(d)
+        self.readout = nn.Linear(d, n_labels, bias=False)
+        # tone-factored auxiliary heads (set via set_tone_heads)
+        self.n_syl = 0
+        # per (poly char, label) bias; low-rank option for tiny models
+        if bias_rank > 0:
+            self.char_bias = nn.Sequential(
+                nn.Embedding(n_poly_chars, bias_rank),
+                nn.Linear(bias_rank, n_labels, bias=False),
+            )
+        else:
+            self.char_bias = nn.Embedding(n_poly_chars, n_labels)
+        nn.init.zeros_(self.char_bias[-1].weight if bias_rank > 0 else self.char_bias.weight)
+        self.pos_head = nn.Linear(d, 11)
+        self.apply(self._init)
+
+    def _init(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+
+    def set_tone_heads(self, n_syl, n_tone=6):
+        import torch.nn as nn
+        self.n_syl = n_syl
+        dev = self.readout.weight.device
+        self.syl_head = nn.Linear(self.readout.in_features, n_syl).to(dev)
+        self.tone_head = nn.Linear(self.readout.in_features, n_tone).to(dev)
+
+    def init_bias_prior(self, counts):
+        """counts: list (aligned with poly chars) of dict label_id -> count"""
+        import numpy as np
+        if isinstance(self.char_bias, nn.Sequential):
+            n = self.char_bias[0].weight.shape[0]
+            L = self.char_bias[-1].weight.shape[0]
+        else:
+            n = self.char_bias.weight.shape[0]
+            L = self.char_bias.weight.shape[1]
+        M = np.zeros((n, L), dtype=np.float32)
+        for i, d in enumerate(counts):
+            tot = sum(d.values()) + 1e-9
+            for j, c in d.items():
+                M[i, j] = math.log((c + 1.0) / (tot + len(d)))
+        if isinstance(self.char_bias, nn.Sequential):
+            U, S, Vh = np.linalg.svd(M, full_matrices=False)
+            r = self.char_bias[0].weight.shape[1]
+            emb = U[:, :r] * np.sqrt(S[:r])
+            lin = (np.sqrt(S[:r])[:, None] * Vh[:r]).T
+            with torch.no_grad():
+                self.char_bias[0].weight.copy_(torch.from_numpy(emb))
+                self.char_bias[-1].weight.copy_(torch.from_numpy(lin))
+        else:
+            with torch.no_grad():
+                self.char_bias.weight.copy_(torch.from_numpy(M))
+
+    def forward_with_ids(self, char_ids, poly_ids, positions, pos_mask, cand_mask):
+        pad = char_ids.ne(0)
+        x = self.emb_proj(self.emb(char_ids))
+        x = self.drop(x)
+        x = self.stem(x)
+        for blk in self.blocks:
+            x = blk(x, pad, self.rope)
+        h = self.n_final(x)
+        B, T, D = h.shape
+        P = positions.shape[1]
+        idx = positions.clamp(min=0).unsqueeze(-1).expand(B, P, D)
+        ph = torch.gather(h, 1, idx)
+        poly_ids_s = poly_ids.clamp(min=0)
+        logits = self.readout(ph) + self.char_bias(poly_ids_s)
+        masked = logits.masked_fill(~cand_mask, -1e4)
+        logp = F.log_softmax(masked.float(), dim=-1)
+        pos_logits = self.pos_head(ph)
+        extra = (self.syl_head(ph), self.tone_head(ph)) if self.n_syl else None
+        return logp, pos_logits, ph, masked, extra

g2ps_torch/runtime.py

@@ -0,0 +1,107 @@
+"""g2pS PyTorch runtime — same API as the numpy runtime, with GPU (CUDA) support.
+
+Deps: torch, numpy.
+GPU: moves model+inputs to CUDA automatically when available. Set fp16=True for
+half-precision inference (recommended on tensor-core GPUs).
+"""
+import os
+import pickle
+
+import numpy as np
+import torch
+
+from ._resolve import resolve_model_dir
+
+from .modules import StudentG2P
+
+MAX_LEN = 128
+
+ARCH = {
+    'small': dict(d=384, n_layers=6, n_heads=6, ffn_hidden=1024, emb_dim=None, bias_rank=0),
+    'tiny': dict(d=256, n_layers=5, n_heads=4, ffn_hidden=640, emb_dim=128, bias_rank=64),
+    'nano': dict(d=96, n_layers=6, n_heads=4, ffn_hidden=256, emb_dim=48, bias_rank=16, n_passes=2),
+}
+
+
+class G2pS:
+    _URLS = ['https://github.com/baicai-1145/g2pS-torch/raw/main/models/{t}/{t}.pt',
+             'https://huggingface.co/baicai1145/g2pS/resolve/main/torch_weights/{t}/{t}.pt']
+
+    def __init__(self, model_dir, tier=None, device=None, fp16=None):
+        tier = tier or (model_dir if model_dir in ARCH else None) or os.path.basename(os.path.normpath(model_dir))
+        model_dir = resolve_model_dir(model_dir, ARCH, f'{tier}.pt',
+                                      [u.format(t=tier) for u in self._URLS])
+        self.tier = tier
+        self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        self.fp16 = (fp16 if fp16 is not None else self.device == 'cuda') and self.device == 'cuda'
+
+        meta = pickle.load(open(os.path.join(model_dir, 'meta.pkl'), 'rb'))
+        self.vocab = meta['vocab']
+        self.labels = meta['labels']
+        self.poly_id = {c: i for i, c in enumerate(meta['poly_chars'])}
+        self.char_cands = meta['char_cands']
+        self.n_labels = len(self.labels)
+
+        kw = dict(ARCH[tier])
+        n_passes = kw.pop('n_passes', 1)
+        self.model = StudentG2P(len(self.vocab), len(self.labels), len(meta['poly_chars']),
+                                max_len=MAX_LEN, use_conv_stem=True, n_passes=n_passes, **kw)
+        sd = torch.load(os.path.join(model_dir, tier + '.pt'), map_location='cpu')
+        sd = {k[10:] if k.startswith('_orig_mod.') else k: v for k, v in sd.items()}
+        sd = {k: v for k, v in sd.items() if not k.startswith(('syl_head', 'tone_head'))}
+        self.model.load_state_dict(sd)
+        self.model.to(self.device)
+        if self.fp16:
+            self.model.half()
+        self.model.eval()
+        self.dtype = torch.float16 if self.fp16 else torch.float32
+
+    def _prepare(self, text):
+        text = text.lower()
+        positions = [i for i, c in enumerate(text) if c in self.poly_id]
+        if not positions:
+            return None
+        if len(text) > MAX_LEN:
+            first = positions[0]
+            start = max(0, min(first - MAX_LEN // 2, len(text) - MAX_LEN))
+            text = text[start:start + MAX_LEN]
+            positions = [p - start for p in positions if start <= p < start + MAX_LEN]
+            if not positions:
+                return None
+        char_ids = np.array([[self.vocab.get(c, 1) for c in text]], dtype=np.int64)
+        P = len(positions)
+        poly_ids = np.array([[self.poly_id[text[p]] for p in positions]], dtype=np.int64)
+        pos_arr = np.array([positions], dtype=np.int64)
+        cand = np.zeros((1, P, self.n_labels), dtype=bool)
+        for k, p in enumerate(positions):
+            cand[0, k, list(self.char_cands[text[p]])] = True
+        return text, positions, char_ids, poly_ids, pos_arr, cand
+
+    def predict(self, text):
+        prep = self._prepare(text)
+        if prep is None:
+            return {}, text
+        text, positions, char_ids, poly_ids, pos_arr, cand = prep
+        dev = self.device
+        with torch.no_grad():
+            logp, _, _, _, _ = self.model.forward_with_ids(
+                torch.as_tensor(char_ids, device=dev),
+                torch.as_tensor(poly_ids, device=dev),
+                torch.as_tensor(pos_arr, device=dev),
+                torch.ones(1, len(positions), dtype=torch.bool, device=dev),
+                torch.as_tensor(cand, device=dev))
+        pred = logp.argmax(-1)[0].cpu().numpy()
+        return {p: self.labels[int(pred[k])] for k, p in enumerate(positions)}, text
+
+    def convert(self, text):
+        res, t = self.predict(text)
+        return t, [(c, res.get(i)) for i, c in enumerate(t)]
+
+
+if __name__ == '__main__':
+    import sys
+    g = G2pS(sys.argv[1] if len(sys.argv) > 1 else 'models/nano')
+    print('device:', g.device, '| fp16:', g.fp16)
+    for s in ['他在银行工作，行踪不定。', '能把小孩哄睡是个技术活。']:
+        _, pairs = g.convert(s)
+        print(s, '->', [(c, p) for c, p in pairs if p])

g2ps_torch-0.1.0.dist-info/METADATA

@@ -0,0 +1,6 @@
+Metadata-Version: 2.4
+Name: g2ps-torch
+Version: 0.1.0
+Summary: g2pS PyTorch backend (GPU + fp16)
+Requires-Python: >=3.8
+Requires-Dist: numpy

g2ps_torch-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+g2ps_torch/__init__.py,sha256=ab4S7abmoWlWylOKA8ajspt8ir8tcUlGb1aiWt8lLKg,26
+g2ps_torch/_resolve.py,sha256=_UtLV6jlFUsrsbPFs14tACY6HL6XJYsxo7erAYDEc3k,1595
+g2ps_torch/modules.py,sha256=2TvSJF2fPqdQOtsNyONceiQ7ODleBam-RPUuPSgb9JM,7724
+g2ps_torch/runtime.py,sha256=ooy-1weo2GVJKZ4DbiaspVi0l5Va3W7UwOmyezOUbGw,4660
+g2ps_torch/models/nano/meta.pkl,sha256=oxziLhr2oZ_1aL1-5sGaT2VALiuCwOWdurRck3JDZ4U,77622
+g2ps_torch/models/nano/nano.pt,sha256=5TlMLwDFY67JWC2LkXvBEHh46rBO4kwsxYK5q2fq4IE,4304251
+g2ps_torch/models/small/meta.pkl,sha256=krEDH4PE5clMaGpNXLePhhoBQD_DuDsUvsdvuN1qlko,123832
+g2ps_torch/models/tiny/meta.pkl,sha256=oxziLhr2oZ_1aL1-5sGaT2VALiuCwOWdurRck3JDZ4U,77622
+g2ps_torch-0.1.0.dist-info/METADATA,sha256=kuAHN6Wf7p_GKKR3X94r4kzY6_5nsP4HeNThu7HvdRc,141
+g2ps_torch-0.1.0.dist-info/WHEEL,sha256=mOdXR9C3wMVov7cSY8PY-XCzo3SmmWCHvGQr5csJoao,91
+g2ps_torch-0.1.0.dist-info/top_level.txt,sha256=HZRCxz0yNEiD4E8LzWTFkFzqZqTEcAEAGtoIaVzwpts,11
+g2ps_torch-0.1.0.dist-info/RECORD,,

g2ps_torch-0.1.0.dist-info/WHEEL

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

g2ps_torch-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+g2ps_torch