Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.axmodel filter=lfs diff=lfs merge=lfs -text

model.axmodel

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bdc5e8051c89961b4212929ddf129b9e8976c04ec55000a7c9b72c66678e4886
+size 283103944

python/example.py

@@ -0,0 +1,99 @@
+#!/usr/bin/env python3
+"""Sherpa Punctuation Prediction on AX650 NPU.
+
+Usage:
+    python example.py                                    # run demo texts
+    python example.py "今天天气真好我们出去散步吧"         # custom text
+    python example.py -m /path/to/model.axmodel          # specify model
+"""
+
+import argparse
+import os
+import sys
+
+from sherpa_punct_sdk import PunctuationPipeline
+
+
+DEMO_TEXTS = [
+    "你好吗how are you我很好谢谢",
+    "今天天气真不错我们出去走走吧",
+    "这个方案有三个优点第一成本低第二效率高第三维护简单",
+    "请确认以下事项一合同已签署二款项已到账三交付日期已确定",
+    # Long text (>64 tokens): auto-sliding-window test
+    "人工智能技术正在改变我们的生活方式"
+    "明天下午三点在公司会议室开会请准时参加"
+    "他是一名优秀的工程师工作认真负责"
+    "北京是中国的首都拥有悠久的历史文化"
+    "随着科技的发展人们的生活越来越便利",
+]
+
+
+def find_file(*candidates):
+    """Return the first existing file from candidates."""
+    for path in candidates:
+        if os.path.exists(path):
+            return path
+    return None
+
+
+def resolve_paths(args):
+    """Resolve model and tokens paths with smart defaults."""
+    base = os.path.dirname(os.path.abspath(__file__))
+    project = os.path.dirname(base)
+
+    model = args.model or find_file(
+        os.path.join(project, "model.axmodel"),
+        os.path.join(project, "models", "model.axmodel"),
+        os.path.join(project, "model_convert", "compile", "model.axmodel"),
+    )
+    tokens = args.tokens or find_file(
+        os.path.join(project, "tokens.json"),
+        os.path.join(project, "models", "tokens.json"),
+        os.path.join(project, "model_convert", "export", "tokens.json"),
+    )
+
+    if not model:
+        print("ERROR: model.axmodel not found. Use -m to specify path.")
+        sys.exit(1)
+    if not tokens:
+        print("ERROR: tokens.json not found. Use -t to specify path.")
+        sys.exit(1)
+
+    return model, tokens
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Sherpa Punctuation Prediction on AX650 NPU",
+    )
+    parser.add_argument(
+        "text", nargs="*",
+        help="Text to punctuate. If omitted, runs demo texts.",
+    )
+    parser.add_argument(
+        "-m", "--model",
+        help="Path to model.axmodel (default: auto-detect)",
+    )
+    parser.add_argument(
+        "-t", "--tokens",
+        help="Path to tokens.json (default: auto-detect)",
+    )
+    args = parser.parse_args()
+
+    model_path, tokens_path = resolve_paths(args)
+
+    print(f"Model:  {model_path}")
+    print(f"Tokens: {tokens_path}")
+
+    pipeline = PunctuationPipeline(model_path, tokens_path)
+
+    texts = args.text if args.text else DEMO_TEXTS
+
+    for text in texts:
+        result = pipeline(text)
+        print(f"\nInput:  {text}")
+        print(f"Output: {result}")
+
+
+if __name__ == "__main__":
+    main()

python/requirements.txt

@@ -0,0 +1,5 @@
+numpy>=1.21
+# On-device AX650 NPU inference:
+#   git clone https://github.com/AXERA-TECH/pyaxengine.git
+#   cd pyaxengine && pip install .
+# pyaxengine

python/sherpa_punct_sdk/__init__.py

@@ -0,0 +1,17 @@
+# Sherpa Punctuation SDK for AX650
+#
+# Converts raw Chinese text to punctuation-annotated text using
+# the compiled AXMODEL on AXera NPU (or CPU fallback).
+
+from .pipeline import PunctuationPipeline
+from .preprocess import CharTokenizer
+from .postprocess import decode_punctuation
+from .inference import PunctInference
+
+__all__ = [
+    "PunctuationPipeline",
+    "CharTokenizer",
+    "decode_punctuation",
+    "PunctInference",
+]
+__version__ = "1.0.0"

python/sherpa_punct_sdk/inference.py

@@ -0,0 +1,62 @@
+# Sherpa Punctuation Inference Engine
+#
+# Uses pyaxengine's AxEngineExecutionProvider to run the compiled AXMODEL
+# on AX650 NPU.
+
+import os
+from typing import Optional
+
+import numpy as np
+
+
+class PunctInference:
+    """Inference wrapper for sherpa punct CT Transformer AXMODEL."""
+
+    def __init__(
+        self,
+        model_path: str,
+        provider: Optional[str] = None,
+    ):
+        """Initialize the inference engine.
+
+        Args:
+            model_path: Path to compiled model.axmodel.
+            provider: Execution provider (default: AxEngineExecutionProvider).
+        """
+        if not os.path.exists(model_path):
+            raise FileNotFoundError(f"Model not found: {model_path}")
+
+        self.model_path = model_path
+        self.provider = provider or "AxEngineExecutionProvider"
+        self._session = None
+
+    def _create_session(self):
+        """Create AX Engine inference session."""
+        import axengine
+
+        available = axengine.get_available_providers()
+        if self.provider in available:
+            return axengine.InferenceSession(
+                self.model_path,
+                providers=[self.provider],
+            )
+        return axengine.InferenceSession(
+            self.model_path,
+            providers=available,
+        )
+
+    def __call__(self, inputs: np.ndarray) -> np.ndarray:
+        """Run inference.
+
+        Args:
+            inputs: (1, 64) int32 numpy array of token IDs.
+
+        Returns:
+            logits: (1, 64, 6) float32 numpy array.
+        """
+        if self._session is None:
+            self._session = self._create_session()
+
+        input_name = self._session.get_inputs()[0].name
+        results = self._session.run(None, {input_name: inputs})
+        return results[0]

python/sherpa_punct_sdk/pipeline.py

@@ -0,0 +1,93 @@
+# Sherpa Onnx Punctuation Pipeline
+#
+# End-to-end pipeline: text → tokens → inference → punctuation-annotated text.
+# Long text is automatically split into overlapping windows for the model's
+# fixed 64-token input.
+
+from typing import List
+
+import numpy as np
+
+from .preprocess import CharTokenizer
+from .inference import PunctInference
+from .postprocess import decode_punctuation
+
+
+INPUT_LENGTH = 64
+WINDOW_STRIDE = 60  # step size; overlap = INPUT_LENGTH - WINDOW_STRIDE = 4
+
+
+class PunctuationPipeline:
+    """End-to-end punctuation prediction pipeline.
+
+    Usage:
+        pipeline = PunctuationPipeline("model.axmodel", "tokens.json")
+        result = pipeline("你好吗how are you我很好谢谢")
+        print(result)  # 你好吗，how are you，我很好谢谢。
+    """
+
+    def __init__(
+        self,
+        model_path: str,
+        tokens_path: str,
+        provider: str = "AxEngineExecutionProvider",
+    ):
+        self.tokenizer = CharTokenizer(tokens_path)
+        if not hasattr(self.tokenizer, "id2token") or not self.tokenizer.id2token:
+            raise RuntimeError("Failed to load tokens.json")
+        self.id2token = self.tokenizer.id2token
+        self.inference = PunctInference(model_path, provider)
+
+    def _run_window(self, tokens: list[int]) -> np.ndarray:
+        """Run inference on a single window, return logits for valid tokens."""
+        n = len(tokens)
+        padded = np.zeros((1, INPUT_LENGTH), dtype=np.int32)
+        padded[0, :n] = tokens
+        logits = self.inference(padded)  # (1, INPUT_LENGTH, 6)
+        return logits[0, :n, :]  # only valid token positions
+
+    def __call__(self, text: str) -> str:
+        """Add punctuation to input text.
+
+        Long text (>64 tokens) is processed in overlapping windows:
+        window_size=64, stride=60, overlap=4.
+
+        Args:
+            text: Raw Chinese text (may include English words).
+
+        Returns:
+            Punctuation-annotated text.
+        """
+        token_ids = self.tokenizer.tokenize(text)
+        if not token_ids:
+            return text
+
+        # Short text: single inference
+        if len(token_ids) <= INPUT_LENGTH:
+            logits = self._run_window(token_ids)
+            return decode_punctuation(
+                logits[np.newaxis], token_ids, self.id2token, len(token_ids),
+            )
+
+        # Long text: sliding window
+        all_logits = []
+        for start in range(0, len(token_ids), WINDOW_STRIDE):
+            end = min(start + INPUT_LENGTH, len(token_ids))
+            window_tokens = token_ids[start:end]
+
+            logits = self._run_window(window_tokens)  # (end-start, 6)
+
+            if start == 0:
+                all_logits.append(logits)
+            else:
+                # Discard overlap: previous window already covered those tokens
+                overlap = INPUT_LENGTH - WINDOW_STRIDE
+                new_tokens_start = overlap
+                all_logits.append(logits[new_tokens_start:])
+
+        combined = np.concatenate(all_logits, axis=0)[:len(token_ids)]
+        combined = combined[np.newaxis, :, :]  # (1, N, 6)
+
+        return decode_punctuation(
+            combined, token_ids, self.id2token, len(token_ids),
+        )

python/sherpa_punct_sdk/postprocess.py

@@ -0,0 +1,102 @@
+# Sherpa Onnx Punctuation Postprocessor
+#
+# Converts model logits output to punctuation-annotated text.
+# 6 classes: <unk>(0), _(1), ，(2), 。(3), ？(4), 、(5)
+
+import numpy as np
+from typing import List
+
+
+PUNCT_CLASSES = [0, 1, 2, 3, 4, 5]
+PUNCT_MARKS = ["", "", "，", "。", "？", "、"]
+IGNORE_ID = 1  # underscore class = no punctuation
+
+
+def decode_punctuation(
+    logits: np.ndarray,
+    token_ids: List[int],
+    id2token: List[str],
+    original_length: int,
+    dot_id: int = 3,
+    comma_id: int = 2,
+    quest_id: int = 4,
+    pause_id: int = 5,
+) -> str:
+    """Decode the model output to punctuation-annotated text.
+
+    Args:
+        logits: (1, 64, 6) float32 array from model
+        token_ids: list of original (unpadded) token IDs
+        id2token: vocab list mapping ID → token string
+        original_length: length before padding (<= 64)
+        dot_id, comma_id, quest_id, pause_id: class IDs for punctuation
+
+    Returns:
+        Annotated text string with punctuation inserted
+    """
+    # Take only valid portion
+    if original_length > logits.shape[1]:
+        original_length = logits.shape[1]
+    if original_length > len(token_ids):
+        original_length = len(token_ids)
+
+    logits = logits[0, :original_length, :]
+    ids = token_ids[:original_length]
+
+    # Argmax over classes
+    out = np.argmax(logits, axis=-1).tolist()
+
+    # Segment with sentence-boundary heuristics
+    # (simplified from original sherpa code)
+    max_len = 200
+    segment_size = 20
+    num_segments = (len(ids) + segment_size - 1) // segment_size
+
+    punctuations = []
+    last = -1
+    for i in range(num_segments):
+        this_start = i * segment_size
+        this_end = min(this_start + segment_size, len(ids))
+        if last != -1:
+            this_start = last
+
+        seg_out = out[this_start:this_end]
+
+        dot_index = -1
+        comma_index = -1
+        for k in range(len(seg_out) - 1, 1, -1):
+            if seg_out[k] in (dot_id, quest_id):
+                dot_index = k
+                break
+            if comma_index == -1 and seg_out[k] == comma_id:
+                comma_index = k
+
+        if dot_index == -1 and len(ids) >= max_len and comma_index != -1:
+            dot_index = comma_index
+            seg_out[dot_index] = dot_id
+
+        if dot_index == -1:
+            if last == -1:
+                last = this_start
+            if i == num_segments - 1:
+                dot_index = len(seg_out) - 1
+        else:
+            last = this_start + dot_index + 1
+
+        if dot_index != -1:
+            punctuations += seg_out[: dot_index + 1]
+
+    # Build output
+    ans = []
+    for j, p in enumerate(punctuations):
+        if j >= len(ids):
+            break
+        t = id2token[ids[j]] if ids[j] < len(id2token) else "<unk>"
+        # Insert space before ASCII tokens
+        if ans and len(ans[-1][0].encode()) == 1 and len(t[0].encode()) == 1:
+            ans.append(" ")
+        ans.append(t)
+        if p != IGNORE_ID and p < len(PUNCT_MARKS) and PUNCT_MARKS[p]:
+            ans.append(PUNCT_MARKS[p])
+
+    return "".join(ans)

python/sherpa_punct_sdk/preprocess.py

@@ -0,0 +1,128 @@
+# Sherpa Onnx Punctuation Preprocessor (CharTokenizer)
+#
+# Model: sherpa-onnx-punct-ct-transformer
+# Tokenizer: character-level for Chinese, word-level for English
+# Vocab: tokens.json (272727 entries)
+# Padding: to 64 tokens
+
+import json
+import os
+from typing import List, Tuple
+import numpy as np
+
+
+class CharTokenizer:
+    """Character/word tokenizer for the sherpa punct CT Transformer model."""
+
+    def __init__(self, tokens_path: str, unk_symbol: str = "<unk>"):
+        if not os.path.exists(tokens_path):
+            raise FileNotFoundError(f"tokens.json not found: {tokens_path}")
+        with open(tokens_path, "r", encoding="utf-8") as f:
+            id2token = json.load(f)
+        self.id2token = id2token
+        self.token2id = {tok: idx for idx, tok in enumerate(id2token)}
+        self.unk_id = self.token2id.get(unk_symbol, 0)
+
+    def tokenize(self, text: str) -> List[int]:
+        """Split text into tokens and return token IDs.
+
+        Chinese characters are segmented individually.
+        English words are kept as whole tokens.
+        """
+        # Split on whitespace
+        word_list = text.split()
+
+        words = []
+        for w in word_list:
+            s = ""
+            for c in w:
+                if len(c.encode()) > 1:
+                    # Multi-byte character (Chinese, Japanese, etc.)
+                    if s == "":
+                        s = c
+                    elif len(s[-1].encode()) > 1:
+                        s += c
+                    else:
+                        words.append(s)
+                        s = c
+                else:
+                    # ASCII character
+                    if s == "":
+                        s = c
+                    elif len(s[-1].encode()) > 1:
+                        words.append(s)
+                        s = c
+                    else:
+                        s += c
+            if s:
+                words.append(s)
+
+        ids = []
+        for w in words:
+            if len(w[0].encode()) > 1:
+                # Chinese phrase: tokenize each character
+                for c in w:
+                    ids.append(self.token2id.get(c, self.unk_id))
+            else:
+                ids.append(self.token2id.get(w, self.unk_id))
+        return ids
+
+    def tokenize_full(self, text: str) -> List[int]:
+        """Tokenize full text without truncation or padding."""
+        return self.tokenize(text)
+
+    def encode(
+        self, text: str, pad_length: int = 64
+    ) -> Tuple[np.ndarray, int]:
+        """Tokenize and pad to fixed length. Truncates if > pad_length.
+
+        Returns:
+            input_array: (1, pad_length) int32 numpy array
+            original_length: actual token count before padding
+        """
+        ids = self.tokenize(text)
+        original_len = len(ids)
+
+        # Truncate or pad to pad_length
+        if len(ids) > pad_length:
+            ids = ids[:pad_length]
+            original_len = pad_length
+
+        padded = np.zeros((1, pad_length), dtype=np.int32)
+        padded[0, : len(ids)] = ids
+
+        return padded, min(original_len, pad_length)
+
+    def encode_long(
+        self, text: str, window_size: int = 64
+    ) -> Tuple[List[np.ndarray], List[int], List[int]]:
+        """Tokenize long text into sliding windows for batched inference.
+
+        Splits full token sequence into windows of window_size.
+        Each window is padded to window_size if shorter.
+
+        Returns:
+            windows: list of (1, window_size) int32 arrays
+            window_token_ids: list of token ID lists per window
+            window_lens: original token lengths per window (before padding)
+        """
+        ids = self.tokenize(text)
+        if not ids:
+            return [], [], []
+
+        windows = []
+        window_token_ids = []
+        window_lens = []
+
+        for start in range(0, len(ids), window_size):
+            chunk = ids[start:start + window_size]
+            chunk_len = len(chunk)
+
+            padded = np.zeros((1, window_size), dtype=np.int32)
+            padded[0, :chunk_len] = chunk
+
+            windows.append(padded)
+            window_token_ids.append(chunk)
+            window_lens.append(chunk_len)
+
+        return windows, window_token_ids, window_lens