#!/usr/bin/env python
# coding: utf-8

# In[1]:


import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
from math import exp
from sys import exit
rknn = RKNN(verbose=True)

# 设置当前文件夹为脚本所在文件夹
os.chdir(os.path.dirname(os.path.abspath(__file__)))

# In[3]:


ONNX_MODEL="model.onnx"
RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
DATASET="dataset.txt"
QUANTIZE=False
detailed_performance_log = True


# In[4]:
seq_len = 256

# pre-process config
print('--> Config model')
rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3)
print('done')

# Load ONNX model
print('--> Loading model')
ret = rknn.load_onnx(model=ONNX_MODEL,
                     inputs=["input_ids", "attention_mask", "token_type_ids"],
                     input_size_list=[[1, seq_len], [1, seq_len], [1, seq_len]],
                     outputs=["/encoder/layer.21/output/LayerNorm/Add_1_output_0"])
if ret != 0:
    print('Load model failed!')
    exit(ret)
print('done')

# Build model
print('--> Building model')
ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
if ret != 0:
    print('Build model failed!')
    exit(ret)
print('done')

#export
print('--> Export RKNN model')
ret = rknn.export_rknn(RKNN_MODEL)
if ret != 0:
    print('Export RKNN model failed!')
    exit(ret)
print('done')