Upload 10 files

.gitattributes

@@ -33,3 +33,6 @@ 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
+ea50ffd6-c6fe-11ef-8ff3-1c860b30973e filter=lfs diff=lfs merge=lfs -text
+RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.emb filter=lfs diff=lfs merge=lfs -text
+RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.rknn filter=lfs diff=lfs merge=lfs -text

RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.emb

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

RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.onnx

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

RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.rknn

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

convert_rknn.py

@@ -0,0 +1,54 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+import datetime
+from rknn.api import RKNN
+from sys import exit
+
+
+ONNX_MODEL = "RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.onnx"
+RKNN_MODEL = ONNX_MODEL.replace(".onnx", ".rknn")
+DATASET = ""
+QUANTIZE = False
+detailed_performance_log = True
+
+timedate_iso = datetime.datetime.now().isoformat()
+
+rknn = RKNN(verbose=True)
+rknn.config(
+    # mean_values=[x * 255 for x in [0.485, 0.456, 0.406]],
+    # std_values=[x * 255 for x in [0.229, 0.224, 0.225]],
+    quantized_dtype="w8a8",
+    quantized_algorithm="normal",
+    quantized_method="channel",
+    quantized_hybrid_level=0,
+    target_platform="rk3588",
+    quant_img_RGB2BGR=False,
+    float_dtype="float16",
+    optimization_level=3,
+    custom_string=f"converted at {timedate_iso}",
+    remove_weight=False,
+    compress_weight=False,
+    inputs_yuv_fmt=None,
+    single_core_mode=False,
+    dynamic_input=None,
+    model_pruning=False,
+    op_target=None,
+    quantize_weight=False,
+    remove_reshape=False,
+    sparse_infer=False,
+    enable_flash_attention=False,
+    # 隐藏的参数
+    # disable_rules=[],
+    # sram_prefer=False,
+    # nbuf_prefer=False,
+    # check_data=[],
+)
+
+ret = rknn.load_onnx(model=ONNX_MODEL)
+ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+ret = rknn.export_rknn(RKNN_MODEL)
+
+# ret = rknn.init_runtime(target='rk3588',device_id='cbb956772bf5dac9',core_mask=RKNN.NPU_CORE_0,perf_debug=detailed_performance_log)
+# rknn.eval_perf()
+# ret = rknn.accuracy_analysis(inputs=['../embeddings.npy','../state.npy','../scale_ratio.npy'], target='rk3588', device_id=device_id)

ea50ffd6-c6fe-11ef-8ff3-1c860b30973e

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

export_onnx.py

@@ -0,0 +1,120 @@
+from rwkv_src.rwkv_model import RWKV_RNN, make_chunks
+import types
+import os
+import torch
+import numpy as np
+import argparse
+import json
+import copy
+from pathlib import Path
+import onnx
+from onnx import shape_inference
+
+parser = argparse.ArgumentParser(description='Convert model')
+parser.add_argument('model', type=Path, help='Path to RWKV pth file')
+parser.add_argument('--chunks', type=int, default=1, help='Number of chunks')
+parser.add_argument('--ext_embedding', action='store_true', default=False, help='Use external embedding')
+parser.add_argument('--prefill_model', action='store_true', help='Convert model for sequential prefill')
+parser.add_argument('--wkv_customop', action='store_true', help='Use custom op for wkv')
+parser_args = parser.parse_args()
+
+seq_length = 32 if parser_args.prefill_model else 1
+
+model_args = types.SimpleNamespace()
+model_args.USE_CUDA = False
+model_args.fp16 = False
+model_args.wkv_customop = parser_args.wkv_customop
+model_args.USE_EMBEDDING = False if parser_args.ext_embedding else True
+
+model_args.MODEL_NAME = str(parser_args.model)
+
+if 'ABC' in model_args.MODEL_NAME or 'MIDI' in model_args.MODEL_NAME or 'x070' in model_args.MODEL_NAME:
+    model_args.RESCALE_LAYER = 0
+else:
+    model_args.RESCALE_LAYER = 6
+
+model = make_chunks(parser_args.chunks, model_args) if parser_args.chunks > 1 else RWKV_RNN(model_args)
+
+if parser_args.prefill_model:
+    model_args.MODEL_NAME = model_args.MODEL_NAME + "_prefill"
+
+os.path.exists("onnx") or os.mkdir("onnx")
+
+if type(model) == list:
+    args = model[0].args
+    if not args.USE_EMBEDDING:
+        model[0].emb_weight.cpu().numpy().astype(np.float32).tofile("onnx/" + args.MODEL_NAME.split("/")[-1] + f"_chunk1of{len(model)}.emb")
+    args = model[0].args
+    fp16 = args.fp16
+    states = []
+    for i in range(args.n_layer):
+        states.append(torch.zeros(1, args.n_embd, dtype=torch.float16 if fp16 else torch.float32))
+        states.append(torch.zeros(args.n_head, args.head_size, args.head_size, dtype=torch.float16 if fp16 else torch.float32))
+        states.append(torch.zeros(1, args.n_embd, dtype=torch.float16 if fp16 else torch.float32))
+    if model[0].device is not torch.device('cpu'):
+        states = [i.to(model[0].device) for i in states]
+
+    for i in range(len(model)):
+        dirname = "onnx/" + args.MODEL_NAME.split("/")[-1] + f"_chunk{i+1}of{len(model)}"
+        os.path.exists(dirname) or os.mkdir(dirname)
+        if i == 0 and args.USE_EMBEDDING:
+            in0 = torch.LongTensor([[1]*seq_length])
+        else:
+            in0 = torch.zeros(1, seq_length, args.n_embd, dtype=torch.float16 if fp16 else torch.float32)
+
+        if model[0].device is not torch.device('cpu'):
+            in0 = in0.to(model[0].device)
+        inputs = {'in0': in0, 'state': [states[j] for j in range(3*model[i].layer_begin, 3*model[i].layer_end)]}
+        input_names = ['in'] + [f'state{j}_in' for j in range(3*model[i].layer_begin, 3*model[i].layer_end)]
+        output_names = ['out'] + [f'state{j}_out' for j in range(3*model[i].layer_begin, 3*model[i].layer_end)]
+
+        if args.wkv_customop:
+            from torch.onnx.symbolic_helper import _get_tensor_sizes
+            from torch.onnx import register_custom_op_symbolic
+            op_name = "rwkv::wkv_chunk" if parser_args.prefill_model else "rwkv::wkv"
+            def onnx_custom_wkv(g, k, v, r, state2, time_first, time_decay):
+                out1, out2 = g.op(op_name, k, v, r, state2, time_first, time_decay, outputs=2)
+                return out1.setType(k.type().with_dtype(torch.float32).with_sizes([seq_length, _get_tensor_sizes(k)[0], 1, args.head_size])),\
+                 out2.setType(k.type().with_dtype(torch.float32).with_sizes([1, _get_tensor_sizes(k)[0], args.head_size, args.head_size]))
+            register_custom_op_symbolic(op_name, onnx_custom_wkv, 9)
+
+        torch.onnx.export(model[i], inputs, dirname + "/" + args.MODEL_NAME.split("/")[-1] + f"_chunk{i+1}of{len(model)}.onnx", input_names=input_names, output_names=output_names, opset_version=17)
+        shape_inference.infer_shapes_path(dirname + "/" + args.MODEL_NAME.split("/")[-1] + f"_chunk{i+1}of{len(model)}.onnx")
+        onnx_model = onnx.load(dirname + "/" + args.MODEL_NAME.split("/")[-1] + f"_chunk{i+1}of{len(model)}.onnx")
+
+        # To make model compatible with other frameworks
+        for initializer in onnx_model.graph.initializer:
+            shape = list(initializer.dims)
+            value_info = onnx.helper.make_tensor_value_info(initializer.name, initializer.data_type, shape)
+            onnx_model.graph.value_info.append(value_info)
+        onnx.save_model(onnx_model, dirname + "/" + args.MODEL_NAME.split("/")[-1] + f"_chunk{i+1}of{len(model)}.onnx", save_as_external_data=True, all_tensors_to_one_file=True)
+        print(f"onnx model chunk{i} saved to {dirname}" + "/" + args.MODEL_NAME.split("/")[-1] + f"_chunk{i+1}of{len(model)}.onnx")
+
+else:
+    args = model.args
+    if not args.USE_EMBEDDING:
+        model.emb_weight.cpu().numpy().astype(np.float32).tofile("onnx/" + args.MODEL_NAME.split("/")[-1] + ".emb")
+    args = model.args
+    fp16 = args.fp16
+    in0 = torch.LongTensor([[1]*seq_length]) if args.USE_EMBEDDING else torch.zeros(1, seq_length, args.n_embd, dtype=torch.float16 if fp16 else torch.float32)
+    states = []
+    for i in range(model.args.n_layer):
+        states.append(torch.zeros(1, model.args.n_embd, dtype=torch.float16 if fp16 else torch.float32))
+        states.append(torch.zeros(model.args.n_head, model.args.head_size, model.args.head_size, dtype=torch.float16 if fp16 else torch.float32))
+        states.append(torch.zeros(1, model.args.n_embd, dtype=torch.float16 if fp16 else torch.float32))
+    if model.device is not torch.device('cpu'):
+        states = [tensor.to(model.device) for tensor in states]
+    inputs = {'in0': in0, 'state': states}
+    input_names = ['in'] + [f'state{i}_in' for i in range(3*model.args.n_layer)]
+    output_names = ['logits'] + [f'state{i}_out' for i in range(3*model.args.n_layer)]
+    torch.onnx.export(model, inputs, "onnx/" + args.MODEL_NAME.split("/")[-1] + ".onnx", input_names=input_names, output_names=output_names, opset_version=17)
+    shape_inference.infer_shapes_path("onnx/" + args.MODEL_NAME.split("/")[-1] + ".onnx")
+    onnx_model = onnx.load("onnx/" + args.MODEL_NAME.split("/")[-1] + ".onnx")
+
+    # To make model compatible with other frameworks
+    for initializer in onnx_model.graph.initializer:
+        shape = list(initializer.dims)
+        value_info = onnx.helper.make_tensor_value_info(initializer.name, initializer.data_type, shape)
+        onnx_model.graph.value_info.append(value_info)
+    onnx.save_model(onnx_model, "onnx/" + args.MODEL_NAME.split("/")[-1] + ".onnx", save_as_external_data=True, all_tensors_to_one_file=True)
+    print(f"onnx model saved to onnx/" + args.MODEL_NAME.split("/")[-1] + ".onnx")

inference.py

@@ -0,0 +1,210 @@
+# import onnxruntime as ort  # Uncomment this line to use onnxruntime
+import ztu_somemodelruntime_rknnlite2 as ort  # Uncomment this line to use rknnlite2 
+import numpy as np
+from pathlib import Path
+from rwkv_tokenizer import RWKV_TOKENIZER
+import time
+
+class RWKVModel:
+    def __init__(self, model_path: str, tokenizer_path: str = None, use_external_embedding: bool = False):
+        # 加载ONNX模型
+        session_options = ort.SessionOptions()
+        # session_options.core_mask = 7 # 00000111 使用0,1,2三个核心
+        self.session = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'], session_options=session_options)
+        
+        # 打印模型输入信息
+        print("\nModel inputs:")
+        for inp in self.session.get_inputs():
+            print(f"{inp.name}: shape={inp.shape}, type={inp.type}")
+        
+        # 获取模型信息
+        self.n_layer = len([x for x in self.session.get_inputs() if 'state' in x.name]) // 3
+        self.n_embd = self.session.get_inputs()[0].shape[-1] if not use_external_embedding else None
+        
+        # 从模型中获取状态向量的维度
+        self.state_shapes = {}
+        for inp in self.session.get_inputs():
+            if 'state' in inp.name:
+                self.state_shapes[inp.name] = inp.shape
+        
+        print("\nNumber of layers:", self.n_layer)
+            
+        # 加载tokenizer
+        if tokenizer_path:
+            self.tokenizer = RWKV_TOKENIZER(tokenizer_path)
+        else:
+            self.tokenizer = None
+            
+        # 加载外部embedding(如果需要)
+        self.use_external_embedding = use_external_embedding
+        if use_external_embedding:
+            emb_path = Path(model_path).parent / (Path(model_path).stem + '.emb')
+            self.embedding = np.fromfile(emb_path, dtype=np.float32)
+            # 重新组织embedding数组的形状
+            vocab_size = len(self.embedding) // 768  # 假设embedding维度是768
+            self.embedding = self.embedding.reshape(vocab_size, 768)
+            self.n_embd = 768
+            print(f"\nEmbedding shape: {self.embedding.shape}")
+            
+        # 初始化状态
+        self.reset_state()
+        
+    def reset_state(self):
+        """重置所有状态为0"""
+        self.states = []
+        for i in range(self.n_layer * 3):
+            state_name = f'state{i}_in'
+            state_shape = self.state_shapes[state_name]
+            self.states.append(np.zeros(state_shape, dtype=np.float32))
+    
+    def _prepare_inputs(self, token_id):
+        """准备模型输入"""
+        inputs = {}
+        
+        # 准备主输入
+        if self.use_external_embedding:
+            # 使用外部embedding
+            embedding = self.embedding[token_id].reshape(1, 1, self.n_embd)
+            inputs['in'] = embedding.astype(np.float32)
+        else:
+            # 使用token id
+            inputs['in'] = np.array([[token_id]], dtype=np.int64)
+            
+        # 添加状态
+        for i in range(len(self.states)):
+            inputs[f'state{i}_in'] = self.states[i]
+            
+        # 打印输入shape
+        if token_id == 0:  # 只打印第一个token的信息
+            print("\nPrepared input shapes:")
+            for k, v in inputs.items():
+                print(f"{k}: shape={v.shape}, type={v.dtype}")
+            
+        return inputs
+        
+    def forward(self, token_id):
+        """单步推理"""
+        # 准备输入
+        inputs = self._prepare_inputs(token_id)
+        
+        # 运行推理
+        outputs = self.session.run(None, inputs)
+        
+        # 打印输出信息(仅第一次)
+        if token_id == 0:
+            print("\nModel outputs:")
+            for i, out in enumerate(outputs):
+                print(f"Output {i}: shape={out.shape}, type={out.dtype}")
+        
+        # 更新状态
+        for i in range(len(self.states)):
+            new_state = outputs[i + 1]  # 第一个输出是logits
+            # 确保维度匹配
+            if new_state.shape != self.states[i].shape:
+                if token_id == 0:
+                    print(f"\nState shape mismatch for state{i}_in:")
+                    print(f"Expected: {self.states[i].shape}")
+                    print(f"Got: {new_state.shape}")
+                # 处理维度
+                if len(self.states[i].shape) == 2:  # (1, 768)
+                    new_state = new_state.squeeze(1)  # (1, 1, 768) -> (1, 768)
+                elif len(self.states[i].shape) == 3:  # (12, 64, 64)
+                    new_state = new_state.squeeze(0)  # (1, 12, 64, 64) -> (12, 64, 64)
+            self.states[i] = new_state
+            
+        return outputs[0]  # 返回logits
+        
+    def generate(self, prompt: str, max_length: int = 100, temperature: float = 1.0, stop_tokens: set = None):
+        """生成文本"""
+        if not self.tokenizer:
+            raise ValueError("需要提供tokenizer才能进行文本生成")
+            
+        # 编码prompt
+        tokens = self.tokenizer.encode(prompt)
+        generated = list(tokens)
+        
+        # 重置状态
+        self.reset_state()
+        
+        # 处理prompt
+        print("\nProcessing prompt...", end='', flush=True)
+        t_start = time.time()
+        for token in tokens:
+            logits = self.forward(token)
+        t_prompt = time.time() - t_start
+        print(f" Done. ({len(tokens)} tokens, {t_prompt:.2f}s, {len(tokens)/t_prompt:.2f} tokens/s)")
+        
+        # 生成新token
+        print("\nGenerating:", end='', flush=True)
+        t_start = time.time()
+        generated_tokens = 0
+        
+        for i in range(max_length):
+            # 获取logits并应用temperature
+            t_token_start = time.time()
+            logits = self.forward(generated[-1])
+            
+            # 打印第一次生成的logits信息
+            if i == 0:
+                print(f"\nLogits shape: {logits.shape}")
+                
+            # 确保logits是1维的
+            logits = logits.reshape(-1)  # 展平成1维
+            
+            if temperature > 0:
+                # 应用temperature并计算概率
+                logits = logits / temperature
+                # 减去最大值以避免exp溢出
+                logits = logits - np.max(logits)
+                probs = np.exp(logits)
+                probs = probs / np.sum(probs)
+                next_token = np.random.choice(len(probs), p=probs)
+            else:
+                next_token = np.argmax(logits)
+                
+            generated.append(next_token)
+            generated_tokens += 1
+            
+            # 检查是否生成了停止标记
+            if stop_tokens and next_token in stop_tokens:
+                break
+            
+            # 实时输出新生成的token
+            new_text = self.tokenizer.decode([next_token])
+            print(new_text, end='', flush=True)
+            
+
+                
+        t_generate = time.time() - t_start
+        print(f"\n\nGeneration finished: {generated_tokens} tokens generated in {t_generate:.2f}s ({generated_tokens/t_generate:.2f} tokens/s)")
+                
+        return self.tokenizer.decode(generated)
+
+def main():
+    import time
+    
+    # 使用示例
+    print("Loading model...")
+    t_start = time.time()
+    model = RWKVModel(
+        model_path='RWKV-x070-World-0.1B-v2.8-20241210-ctx4096.onnx',
+        tokenizer_path='rwkv_vocab_v20230424.txt',
+        use_external_embedding=True
+    )
+    print(f"Model loaded in {time.time() - t_start:.2f}s")
+    
+    prompt = "Here is a example of Quick Sort algorithm implemented in C++:\n```cpp"
+    print(f"\nPrompt: {prompt}")
+    
+    generated_text = model.generate(
+        prompt=prompt,
+        max_length=1024,
+        temperature=0.7,
+        stop_tokens={0, 1, 2, 3}  # 特殊token作为停止标记
+    )
+    
+    print("\nFull text:")
+    print(generated_text)
+
+if __name__ == '__main__':
+    main() 

rwkv_tokenizer.py

@@ -0,0 +1,89 @@
+from typing import List,Set,Dict
+
+class ABCTokenizer():
+    def __init__(self):
+        self.pad_token_id = 0
+        self.bos_token_id = 2
+        self.eos_token_id = 3
+    def encode(self, text):
+        ids = [ord(c) for c in text]
+        return ids
+    def decode(self, ids):
+        txt = ''.join(chr(idx) if idx > self.eos_token_id else '' for idx in ids if idx != self.eos_token_id)
+        return txt
+
+class RWKV_TOKENIZER():
+    table: List[List[List[bytes]]]
+    good: List[Set[int]]
+    wlen: List[int]
+    def __init__(self, file_name):
+        self.idx2token = {}
+        sorted = [] # must be already sorted
+        lines = open(file_name, "r", encoding="utf-8").readlines()
+        for l in lines:
+            idx = int(l[:l.index(' ')])
+            x = eval(l[l.index(' '):l.rindex(' ')])
+            x = x.encode("utf-8") if isinstance(x, str) else x
+            assert isinstance(x, bytes)
+            assert len(x) == int(l[l.rindex(' '):])
+            sorted += [x]
+            self.idx2token[idx] = x
+
+        self.token2idx = {}
+        for k, v in self.idx2token.items():
+            self.token2idx[v] = int(k)
+
+        # precompute some tables for fast matching
+        self.table = [[[] for j in range(256)] for i in range(256)]
+        self.good = [set() for i in range(256)]
+        self.wlen = [0 for i in range(256)]
+
+        for i in reversed(range(len(sorted))): # reverse order - match longer tokens first
+            s = sorted[i]
+            if len(s) >= 2:
+                s0 = int(s[0])
+                s1 = int(s[1])
+                self.table[s0][s1] += [s]
+                self.wlen[s0] = max(self.wlen[s0], len(s))
+                self.good[s0].add(s1)
+
+    def encodeBytes(self, src: bytes) -> List[int]:
+        src_len: int = len(src)
+        tokens: List[int] = []
+        i: int = 0
+        while i < src_len:
+            s: bytes = src[i : i + 1]
+
+            if i < src_len - 1:
+                s1: int = int(src[i + 1])
+                s0: int = int(src[i])
+                if s1 in self.good[s0]:
+                    sss: bytes = src[i : i + self.wlen[s0]]
+                    try:
+                        s = next(filter(sss.startswith, self.table[s0][s1]))
+                    except:
+                        pass
+            tokens.append(self.token2idx[s])
+            i += len(s)
+
+        return tokens
+
+    def decodeBytes(self, tokens):
+        return b''.join(map(lambda i: self.idx2token[i], tokens))
+
+    def encode(self, src: str):
+        return self.encodeBytes(src.encode("utf-8"))
+
+    def decode(self, tokens):
+        return self.decodeBytes(tokens).decode('utf-8')
+
+    def printTokens(self, tokens):
+        for i in tokens:
+            s = self.idx2token[i]
+            try:
+                s = s.decode('utf-8')
+            except:
+                pass
+            print(f'{repr(s)}{i}', end=' ')
+            # print(repr(s), i)
+        print()

ztu_somemodelruntime_rknnlite2.py

@@ -0,0 +1,509 @@
+# 模块级常量和函数
+from rknnlite.api import RKNNLite
+import numpy as np
+import os
+import warnings
+import logging
+from typing import List, Dict, Union, Optional
+
+# 配置日志
+logger = logging.getLogger("somemodelruntime_rknnlite2")
+logger.setLevel(logging.ERROR)  # 默认只输出错误信息
+if not logger.handlers:
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
+    logger.addHandler(handler)
+
+# ONNX Runtime日志级别到Python logging级别的映射
+_LOGGING_LEVEL_MAP = {
+    0: logging.DEBUG,    # Verbose
+    1: logging.INFO,     # Info
+    2: logging.WARNING,  # Warning
+    3: logging.ERROR,    # Error
+    4: logging.CRITICAL  # Fatal
+}
+
+def set_default_logger_severity(level: int) -> None:
+    """
+    Sets the default logging severity. 0:Verbose, 1:Info, 2:Warning, 3:Error, 4:Fatal
+    
+    Args:
+        level: 日志级别(0-4)
+    """
+    if level not in _LOGGING_LEVEL_MAP:
+        raise ValueError(f"无效的日志级别: {level}, 应该是0-4之间的整数")
+    logger.setLevel(_LOGGING_LEVEL_MAP[level])
+
+def set_default_logger_verbosity(level: int) -> None:
+    """
+    Sets the default logging verbosity level. To activate the verbose log, 
+    you need to set the default logging severity to 0:Verbose level.
+    
+    Args:
+        level: 日志级别(0-4)
+    """
+    set_default_logger_severity(level)
+
+# NPU核心模式常量
+NPU_CORE_AUTO = 0      # 自动选择
+NPU_CORE_0 = 1        # 使用核心0
+NPU_CORE_1 = 2        # 使用核心1
+NPU_CORE_2 = 4        # 使用核心2
+NPU_CORE_0_1 = 3      # 使用核心0和1
+NPU_CORE_0_1_2 = 7    # 使用所有核心
+NPU_CORE_ALL = 0xffff # 使用所有核心
+
+# RKNN tensor type到numpy dtype的映射
+RKNN_DTYPE_MAP = {
+    0: np.float32,  # RKNN_TENSOR_FLOAT32
+    1: np.float16,  # RKNN_TENSOR_FLOAT16
+    2: np.int8,     # RKNN_TENSOR_INT8
+    3: np.uint8,    # RKNN_TENSOR_UINT8
+    4: np.int16,    # RKNN_TENSOR_INT16
+    5: np.uint16,   # RKNN_TENSOR_UINT16
+    6: np.int32,    # RKNN_TENSOR_INT32
+    7: np.uint32,   # RKNN_TENSOR_UINT32
+    8: np.int64,    # RKNN_TENSOR_INT64
+    9: bool,        # RKNN_TENSOR_BOOL
+    10: np.int8,    # RKNN_TENSOR_INT4 (用int8表示)
+}
+
+def get_available_providers() -> List[str]:
+    """
+    获取可用的设备提供者列表(为保持接口兼容性的占位函数)
+    
+    Returns:
+        list: 可用的设备提供者列表,总是返回["CPUExecutionProvider"]
+    """
+    return ["CPUExecutionProvider"]
+
+def get_version_info() -> Dict[str, str]:
+    """
+    获取版本信息
+    
+    Returns:
+        dict: 包含API和驱动版本信息的字典
+    """
+    runtime = RKNNLite()
+    version = runtime.get_sdk_version()
+    return {
+        "api_version": version.split('\n')[2].split(': ')[1].split(' ')[0],
+        "driver_version": version.split('\n')[3].split(': ')[1]
+    }
+
+class IOTensor:
+    """输入/输出张量的信息封装类"""
+    def __init__(self, name, shape, type=None):
+        self.name = name.decode() if isinstance(name, bytes) else name
+        self.shape = shape
+        self.type = type
+
+    def __str__(self):
+        return f"IOTensor(name='{self.name}', shape={self.shape}, type={self.type})"
+
+class SessionOptions:
+    """会话选项类"""
+    def __init__(self):
+        self.async_mode = False  # 是否使用异步模式
+        self.core_mask = 0      # NPU核心选择
+        self.perf_debug = False # 是否启用性能分析
+
+class InferenceSession:
+    """
+    RKNNLite运行时封装类,API风格类似ONNX Runtime
+    """
+    
+    def __init__(self, model_path: str, verbose: bool = False, session_options: Optional[SessionOptions] = None, **kwargs):
+        """
+        初始化运行时并加载模型
+        
+        Args:
+            model_path: 模型文件路径(.rknn或.onnx)
+            verbose: 是否打印详细日志
+            session_options: 会话选项
+            **kwargs: 其他初始化参数
+        """
+        # 只在verbose=True时开启详细日志
+        if verbose:
+            set_default_logger_severity(0)  # Verbose
+        
+        self.model_path = self._process_model_path(model_path)
+        self.runtime = RKNNLite(verbose=verbose)
+        
+        # 加载模型
+        logger.debug(f"正在加载模型: {self.model_path}")
+        ret = self.runtime.load_rknn(self.model_path)
+        if ret != 0:
+            logger.error(f"加载RKNN模型失败: {self.model_path}")
+            raise RuntimeError(f'加载RKNN模型失败: {self.model_path}')
+        logger.debug("模型加载成功")
+            
+        # 应用会话选项
+        options = session_options or SessionOptions()
+        
+        # 初始化运行时
+        logger.debug("正在初始化运行时环境")
+        ret = self.runtime.init_runtime(
+            async_mode=options.async_mode,
+            core_mask=options.core_mask
+        )
+        if ret != 0:
+            logger.error("初始化运行时环境失败")
+            raise RuntimeError('初始化运行时环境失败')
+        logger.debug("运行时环境初始化成功")
+
+        # 获取输入输出信息
+        self._init_io_info()
+        
+        # 保存选项
+        self.options = options
+
+    def get_performance_info(self) -> Dict[str, float]:
+        """
+        获取性能信息
+        
+        Returns:
+            dict: 包含性能信息的字典
+        """
+        if not self.options.perf_debug:
+            raise RuntimeError("性能分析未启用,请在SessionOptions中设置perf_debug=True")
+            
+        perf = self.runtime.rknn_runtime.get_run_perf()
+        return {
+            "run_duration": perf.run_duration / 1000.0  # 转换为毫秒
+        }
+
+    def set_core_mask(self, core_mask: int) -> None:
+        """
+        设置NPU核心使用模式
+        
+        Args:
+            core_mask: NPU核心掩码,使用NPU_CORE_*常量
+        """
+        ret = self.runtime.rknn_runtime.set_core_mask(core_mask)
+        if ret != 0:
+            raise RuntimeError("设置NPU核心模式失败")
+
+    def _process_model_path(self, model_path):
+        """处理模型路径,支持.onnx和.rknn文件"""
+        if not os.path.exists(model_path):
+            logger.error(f"模型文件不存在: {model_path}")
+            raise FileNotFoundError(f"模型文件不存在: {model_path}")
+            
+        # 如果是ONNX文件
+        if model_path.lower().endswith('.onnx'):
+            logger.warning(
+                "检测到ONNX模型文件。注意：SomeModelRuntime不会自动转换ONNX到RKNN。"
+                "请先使用RKNN Toolkit转换模型。"
+                "现在尝试加载同名的.rknn文件。"
+            )
+            # 构造RKNN文件路径
+            rknn_path = os.path.splitext(model_path)[0] + '.rknn'
+            if not os.path.exists(rknn_path):
+                logger.error(f"RKNN模型文件不存在: {rknn_path}")
+                raise FileNotFoundError(
+                    f"RKNN模型文件不存在: {rknn_path}\n"
+                    "请先使用RKNN Toolkit将ONNX模型转换为RKNN格式。"
+                )
+            return rknn_path
+            
+        return model_path
+        
+    def _convert_nhwc_to_nchw(self, shape):
+        """将NHWC格式的shape转换为NCHW格式"""
+        if len(shape) == 4:
+            # NHWC -> NCHW
+            n, h, w, c = shape
+            return [n, c, h, w]
+        return shape
+        
+    def _init_io_info(self):
+        """初始化模型的输入输出信息"""
+        runtime = self.runtime.rknn_runtime
+        
+        # 获取输入输出数量
+        n_input, n_output = runtime.get_in_out_num()
+        
+        # 获取输入信息
+        self.input_tensors = []
+        for i in range(n_input):
+            attr = runtime.get_tensor_attr(i)
+            shape = [attr.dims[j] for j in range(attr.n_dims)]
+            # 对四维输入进行NHWC到NCHW的转换
+            shape = self._convert_nhwc_to_nchw(shape)
+            # 获取dtype
+            dtype = RKNN_DTYPE_MAP.get(attr.type, None)
+            tensor = IOTensor(attr.name, shape, dtype)
+            self.input_tensors.append(tensor)
+            
+        # 获取输出信息
+        self.output_tensors = []
+        for i in range(n_output):
+            attr = runtime.get_tensor_attr(i, is_output=True)
+            shape = runtime.get_output_shape(i)
+            # 获取dtype
+            dtype = RKNN_DTYPE_MAP.get(attr.type, None)
+            tensor = IOTensor(attr.name, shape, dtype)
+            self.output_tensors.append(tensor)
+        
+    def get_inputs(self):
+        """
+        获取模型输入信息
+        
+        Returns:
+            list: 包含输入信息的列表
+        """
+        return self.input_tensors
+        
+    def get_outputs(self):
+        """
+        获取模型输出信息
+        
+        Returns:
+            list: 包含输出信息的列表
+        """
+        return self.output_tensors
+        
+    def run(self, output_names=None, input_feed=None, data_format="nchw", **kwargs):
+        """
+        执行模型推理
+        
+        Args:
+            output_names: 输出节点名称列表,指定需要返回哪些输出
+            input_feed: 输入数据字典或列表
+            data_format: 输入数据格式,"nchw"或"nhwc"
+            **kwargs: 其他运行时参数
+            
+        Returns:
+            list: 模型输出结果列表,如果指定了output_names则只返回指定的输出
+        """
+        if input_feed is None:
+            logger.error("input_feed不能为None")
+            raise ValueError("input_feed不能为None")
+            
+        # 准备输入数据
+        if isinstance(input_feed, dict):
+            # 如果是字典,按照模型输入顺序排列
+            inputs = []
+            input_map = {tensor.name: i for i, tensor in enumerate(self.input_tensors)}
+            for tensor in self.input_tensors:
+                if tensor.name not in input_feed:
+                    raise ValueError(f"缺少输入: {tensor.name}")
+                inputs.append(input_feed[tensor.name])
+        elif isinstance(input_feed, (list, tuple)):
+            # 如果是列表,确保长度匹配
+            if len(input_feed) != len(self.input_tensors):
+                raise ValueError(f"输入数量不匹配: 期望{len(self.input_tensors)}, 实际{len(input_feed)}")
+            inputs = list(input_feed)
+        else:
+            logger.error("input_feed必须是字典或列表类型")
+            raise ValueError("input_feed必须是字典或列表类型")
+            
+        # 执行推理
+        try:
+            logger.debug("开始执行推理")
+            all_outputs = self.runtime.inference(inputs=inputs, data_format=data_format)
+            
+            # 如果没有指定output_names,返回所有输出
+            if output_names is None:
+                return all_outputs
+                
+            # 获取指定的输出
+            output_map = {tensor.name: i for i, tensor in enumerate(self.output_tensors)}
+            selected_outputs = []
+            for name in output_names:
+                if name not in output_map:
+                    raise ValueError(f"未找到输出节点: {name}")
+                selected_outputs.append(all_outputs[output_map[name]])
+                    
+            return selected_outputs
+            
+        except Exception as e:
+            logger.error(f"推理执行失败: {str(e)}")
+            raise RuntimeError(f"推理执行失败: {str(e)}")
+        
+    def close(self):
+        """
+        关闭会话,释放资源
+        """
+        if self.runtime is not None:
+            logger.info("正在释放运行时资源")
+            self.runtime.release()
+            self.runtime = None
+            
+    def __enter__(self):
+        return self
+        
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+    def end_profiling(self) -> Optional[str]:
+        """
+        结束性能分析的存根方法
+        
+        Returns:
+            Optional[str]: None
+        """
+        warnings.warn("end_profiling()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return None
+        
+    def get_profiling_start_time_ns(self) -> int:
+        """
+        获取性能分析开始时间的存根方法
+        
+        Returns:
+            int: 0
+        """
+        warnings.warn("get_profiling_start_time_ns()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return 0
+        
+    def get_modelmeta(self) -> Dict[str, str]:
+        """
+        获取模型元数据的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_modelmeta()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+        
+    def get_session_options(self) -> SessionOptions:
+        """
+        获取会话选项
+        
+        Returns:
+            SessionOptions: 当前会话选项
+        """
+        return self.options
+        
+    def get_providers(self) -> List[str]:
+        """
+        获取当前使用的providers的存根方法
+        
+        Returns:
+            List[str]: ["CPUExecutionProvider"]
+        """
+        warnings.warn("get_providers()是存根方法,始终返回CPUExecutionProvider", RuntimeWarning, stacklevel=2)
+        return ["CPUExecutionProvider"]
+        
+    def get_provider_options(self) -> Dict[str, Dict[str, str]]:
+        """
+        获取provider选项的存根方法
+        
+        Returns:
+            Dict[str, Dict[str, str]]: 空字典
+        """
+        warnings.warn("get_provider_options()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {} 
+
+    def get_session_config(self) -> Dict[str, str]:
+        """
+        获取会话配置的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_session_config()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def get_session_state(self) -> Dict[str, str]:
+        """
+        获取会话状态的存根方法
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_session_state()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def set_session_config(self, config: Dict[str, str]) -> None:
+        """
+        设置会话配置的存根方法
+        
+        Args:
+            config: 会话配置字典
+        """
+        warnings.warn("set_session_config()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_memory_info(self) -> Dict[str, int]:
+        """
+        获取内存使用信息的存根方法
+        
+        Returns:
+            Dict[str, int]: 空字典
+        """
+        warnings.warn("get_memory_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def set_memory_pattern(self, enable: bool) -> None:
+        """
+        设置内存模式的存根方法
+        
+        Args:
+            enable: 是否启用内存模式
+        """
+        warnings.warn("set_memory_pattern()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def disable_memory_pattern(self) -> None:
+        """
+        禁用内存模式的存根方法
+        """
+        warnings.warn("disable_memory_pattern()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_optimization_level(self) -> int:
+        """
+        获取优化级别的存根方法
+        
+        Returns:
+            int: 0
+        """
+        warnings.warn("get_optimization_level()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return 0
+
+    def set_optimization_level(self, level: int) -> None:
+        """
+        设置优化级别的存根方法
+        
+        Args:
+            level: 优化级别
+        """
+        warnings.warn("set_optimization_level()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+
+    def get_model_metadata(self) -> Dict[str, str]:
+        """
+        获取模型元数据的存根方法(与get_modelmeta不同的接口)
+        
+        Returns:
+            Dict[str, str]: 空字典
+        """
+        warnings.warn("get_model_metadata()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return {}
+
+    def get_model_path(self) -> str:
+        """
+        获取模型路径
+        
+        Returns:
+            str: 模型文件路径
+        """
+        return self.model_path
+
+    def get_input_type_info(self) -> List[Dict[str, str]]:
+        """
+        获取输入类型信息的存根方法
+        
+        Returns:
+            List[Dict[str, str]]: 空列表
+        """
+        warnings.warn("get_input_type_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return []
+
+    def get_output_type_info(self) -> List[Dict[str, str]]:
+        """
+        获取输出类型信息的存根方法
+        
+        Returns:
+            List[Dict[str, str]]: 空列表
+        """
+        warnings.warn("get_output_type_info()是存根方法,不提供实际功能", RuntimeWarning, stacklevel=2)
+        return []