NanoSakura-0.3B-Onnx

Model Card for NanoSakura-0.3B (Ja -> Zh)

This model is a highly efficient, custom-built Hybrid Sequence-to-Sequence (Seq2Seq) model designed for Japanese to Chinese translation. It combines the deep understanding capabilities of a Transformer Encoder with the blazing-fast, memory-efficient generation of a Mamba2 (State Space Model) Decoder.

Model Details

Model Description

Traditional Seq2Seq models (like T5 or BART) rely entirely on Transformers. While powerful, the self-attention mechanism in the decoder leads to an $O(N^2)$ computational bottleneck and high KV-cache memory usage during text generation.

This model solves that by introducing a Hybrid Architecture:

1. Encoder (Transformer): Uses Self-Attention + RoPE + SwiGLU to fully capture the global context of the source Japanese text in parallel.
2. Decoder (Mamba2 + Cross-Attention): Replaces self-attention with Mamba2's State Space Model (SSM). This allows for $O(1)$ state-updating generation (no growing KV cache), while retaining Cross-Attention to accurately "look back" at the encoder's features to prevent hallucination.

With only ~287 Million parameters, it achieves excellent translation quality while maintaining an extremely low hardware footprint, making it ideal for edge deployment or high-throughput API services.

• Developed by: telecomadm1145
• Model type: Hybrid Transformer-Mamba2 Seq2Seq
• Language(s) (NLP): Japanese (ja), Chinese (zh)
• License: MIT
• Parameters: ~287M

How to Get Started with the Model

C#:

https://github.com/telecomadm1145/NanoSakura

py:

import os
import json
import argparse
import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer
from huggingface_hub import snapshot_download

class NanoSakuraONNX:
    def __init__(self, repo_id: str = "telecomadm1145/NanoSakura-0.3B-Onnx", onnx_folder: str = "i8"):
        self.repo_id = repo_id
        
        print(f"Downloading/Loading ONNX models from HF Hub: {repo_id} (folder: {onnx_folder})...")
        local_repo_dir = snapshot_download(
            repo_id=repo_id,
            allow_patterns=[f"{onnx_folder}/*", "config.json"]
        )
        
        self.onnx_dir = os.path.join(local_repo_dir, onnx_folder)
        
        config_path = os.path.join(self.onnx_dir, "config.json")
        if not os.path.exists(config_path):
            config_path = os.path.join(local_repo_dir, "config.json")
            
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"找不到 config.json 配置！")
            
        with open(config_path, 'r', encoding='utf-8') as f:
            self.cfg = json.load(f)
            
        print(f"Loading tokenizer from {repo_id}...")
        self.tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
        
        self.bos_token_id = self.cfg.get("bos_token_id", self.tokenizer.bos_token_id)
        self.eos_token_id = self.cfg.get("eos_token_id", self.tokenizer.eos_token_id)
        self.eos_token = self.tokenizer.eos_token if self.tokenizer.eos_token else "<eos>"

        print(f"Loading ONNX sessions from {self.onnx_dir}...")
        providers = ["CPUExecutionProvider"]
        
        self.encoder = ort.InferenceSession(os.path.join(self.onnx_dir, "encoder.onnx"), providers=providers)
        self.cross_kv = ort.InferenceSession(os.path.join(self.onnx_dir, "cross_kv.onnx"), providers=providers)
        self.decoder_step = ort.InferenceSession(os.path.join(self.onnx_dir, "decoder_step.onnx"), providers=providers)
        print("Models loaded successfully!")

    def generate(self, text: str, max_new_tokens: int = 128) -> str:
        if not text.endswith(self.eos_token):
            text = text + self.eos_token

        enc = self.tokenizer(text, return_tensors="np", padding=False, truncation=True)
        input_ids = enc["input_ids"].astype(np.int64)
        attention_mask = enc["attention_mask"].astype(np.int64)
        
        B, S = input_ids.shape

        enc_hidden = self.encoder.run(None, {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        })[0]

        all_k, all_v = self.cross_kv.run(None, {
            "encoder_hidden_states": enc_hidden.astype(np.float32),
        })

        L  = self.cfg["dec_layers"]
        K  = self.cfg["conv_kernel"]
        cd = self.cfg["conv_dim"]
        Hh = self.cfg["n_ssm_heads"]
        N  = self.cfg["d_state"]
        P  = self.cfg["ssm_head_dim"]
        
        conv_states = np.zeros((L, B, K-1, cd), dtype=np.float32)
        ssm_states  = np.zeros((L, B, Hh, N, P), dtype=np.float32)
        
        cross_mask = attention_mask.astype(bool).reshape(B, 1, 1, S)
        
        cur_token_id = np.full((B, 1), self.bos_token_id, dtype=np.int64)
        generated_tokens = []
        
        for step in range(max_new_tokens):
            outs = self.decoder_step.run(None, {
                "dec_token_id": cur_token_id,
                "cross_mask": cross_mask,
                "all_k_enc": all_k,
                "all_v_enc": all_v,
                "all_conv_states": conv_states,
                "all_ssm_states": ssm_states,
            })
            
            logits, conv_states, ssm_states = outs
            
            cur_token_id = np.argmax(logits[:, -1, :], axis=-1)[:, None]
            token_val = cur_token_id[0, 0]
            
            if token_val == self.eos_token_id:
                break
                
            generated_tokens.append(token_val)

        out_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
        return out_text


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="NanoSakura ONNX 推理")
    parser.add_argument("--repo_id", type=str, default="telecomadm1145/NanoSakura-0.3B-Onnx", help="HuggingFace 仓库名")
    parser.add_argument("--onnx_folder", type=str, default="i8", help="仓库中存放ONNX模型的文件夹，默认i8")
    parser.add_argument("--text", type=str, default="私の○○を見てください！", help="输入文本(日文)")
    parser.add_argument("--max_new_tokens", type=int, default=128, help="最大生成长度")
    
    args = parser.parse_args()
    print(f"{args.text}")
    model = NanoSakuraONNX(repo_id=args.repo_id, onnx_folder=args.onnx_folder)
    
    output_text = model.generate(args.text, max_new_tokens=args.max_new_tokens)
    print(f"{output_text}")