import time
import torch
import mindietorch
import numpy as np
from transformers import AutoTokenizer
import sys
import os
import random
os.environ['TOKENIZERS_PARALLELISM'] = 'true'

# run this first:
# 1. cp /root/.cache/huggingface/modules/transformers_modules/Alibaba-NLP/new-impl/40ced75c3017eb27626c9d4ea981bde21a2662f4/*.py ./
# 2. modify modeling.py, change "from .configuration import NewConfig" to "from configuration import NewConfig"
from modeling import NewModel

max_length = 512

# 初始化数组
string_array = []

# 生成字符串数组
current_length = 256
while current_length <= max_length:
    string_array.append("啊" * current_length)
    current_length *= 2

# 随机排序
random.shuffle(string_array)

MODEL_PATH = "/llm/Alibaba-NLP/gte-multilingual-base/"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True, local_files_only=True)
#model = AutoModel.from_pretrained(MODEL_PATH, trust_remote_code=True)
model = NewModel.from_pretrained(MODEL_PATH, trust_remote_code=True)
model.to(torch.float32).eval()

inputs = tokenizer(string_array, add_special_tokens=True, padding='max_length', max_length=max_length, truncation=True, return_tensors='pt').to("cpu")

# 调用torch.jit.trace生成traced_model
#input_ids = inputs['input_ids']
#attention_mask = inputs['attention_mask']
input_ids = torch.ones((1,max_length), dtype=torch.int)
attention_mask = torch.ones((1,max_length), dtype=torch.int)
traced_model = torch.jit.trace(model, [input_ids, attention_mask], strict=False)
#traced_model = torch.jit.script(model)
traced_model.to(torch.float32).eval().to("npu")

# 把traced_model转化成npu模型
#torch.jit.save(traced_model, 'gte_multilingual_base_traced.pt')

# 编译成HW模型
#min_shape = (1, 1)
#max_shape = (len(string_array), max_length)  # batch, max length
#inputs = []
#inputs.append(mindietorch.Input(min_shape = min_shape, max_shape = max_shape, dtype=input_ids.dtype))
#inputs.append(mindietorch.Input(min_shape = min_shape, max_shape = max_shape, dtype=attention_mask.dtype))
inputs = [mindietorch.Input((1, max_length),dtype=input_ids.dtype), mindietorch.Input((1,max_length), dtype=attention_mask.dtype)]

#loaded_model = loaded_model.eval().to("npu")
compiled_module = mindietorch.compile(
     traced_model,
     inputs=inputs,
     #precision_policy=mindietorch.PrecisionPolicy.FP16,
     #truncate_long_and_double=True,
     truncate_long_and_double=False,
     require_full_compilation=False,
     allow_tensor_replace_int=False,
     min_block_size=3,
     torch_executed_ops=[],
     optimization_level=0,
     soc_version="Ascend310P3",
)

# ts模型保存
compiled_module.save("CompiledModel.ts")

mindietorch.finalize()

import json
import time
import torch
import mindietorch
import numpy as np
from transformers import AutoModel, AutoTokenizer
import sys
from modeling import NewModel
import os
import random
import torch.nn.functional as F

os.environ['TOKENIZERS_PARALLELISM'] = 'true'
torch.manual_seed(88)

max_length = 512

MODEL_PATH = "/llm/Alibaba-NLP/gte-multilingual-base/"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True, local_files_only=True)

cpu_model = NewModel.from_pretrained(MODEL_PATH)
cpu_model.eval()

compiled_model = torch.jit.load("CompiledModel.ts")
mindietorch.set_device(0)
compiled_model.to(torch.float32).eval().to("npu")
modeloutput = compiled_model(torch.ones((1,max_length), dtype=torch.int).to("npu"), torch.ones((1,max_length), dtype=torch.int).to("npu"))

text = "阿斯蒂芬啊倒萨方发阿发暗暗大师傅奥奥福德阿凡达噶双方的安抚"
for i in range(8192):
    # 生成一个[1,max_length]长度的随机字符串
    random_length = random.randint(1, max_length)
    random_string = ''.join(random.choice(text) for _ in range(random_length))
    #print(random_string)
    inputs = tokenizer(random_string, add_special_tokens=True, padding='max_length', max_length=max_length, truncation=True, return_tensors='pt').to("cpu")
    input_ids = inputs['input_ids'].to(dtype=torch.int)
    attention_mask = inputs['attention_mask'].to(dtype=torch.int)

    # cpu计算
    modeloutput = cpu_model(input_ids, attention_mask)
    cpu_embedding = modeloutput.last_hidden_state[:, 0][:768].detach().cpu()
    pooler_cpu = modeloutput.get('pooler_output', None)
    hidden_states_cpu = modeloutput.get('hidden_states', None)
    attentions_cpu = modeloutput.get('attentions',None)

    # npu计算第一次
    modeloutput = compiled_model(input_ids.to("npu", dtype=torch.int), attention_mask.to("npu", dtype=torch.int))
    npu_embedding1 = modeloutput['last_hidden_state'][:, 0][:768].detach().to(torch.float32).cpu()
    pooler_output1 = modeloutput.get('pooler_output', None)
    hidden_states1 = modeloutput.get('hidden_states', None)
    attentions1 = modeloutput.get('attentions',None)

    # npu计算第二次
    modeloutput = compiled_model(input_ids.to("npu", dtype=torch.int), attention_mask.to("npu", dtype=torch.int))
    npu_embedding2 = modeloutput['last_hidden_state'][:, 0][:768].detach().to(torch.float32).cpu()
    pooler_output2 = modeloutput.get('pooler_output', None)
    hidden_states2 = modeloutput.get('hidden_states', None)
    attentions2 = modeloutput.get('attentions', None)
    #pooler_output2 = pooler_output1
    #hidden_states2 = hidden_states1
    #attentions2 = attentions1

    cosine_similarity = F.cosine_similarity(npu_embedding1, npu_embedding2).item()
    if cosine_similarity>0.999:
        print("*", end='', flush=True)
        continue
    else:
        print("tensor cos:", cosine_similarity, flush=True)

    if pooler_output1 is None and pooler_output2 is None:
        print("tensor pooler_output is None")
    else:
        pooler_output1 = pooler_output1.detach().cpu()
        pooler_output2 = pooler_output2.detach().cpu()
        print("tensor 0 type:", pooler_cpu.dtype, ",", pooler_output1.dtype)
        #print("[0]cpu:", pooler_cpu.shape, ":", pooler_cpu[0:8])
        #print("[0]1:", pooler_output1.shape, ":", pooler_output1[0:8])
        #print("[0]2:", pooler_output2.shape, ":", pooler_output2[0:8])
        if not torch.equal(pooler_output1, pooler_cpu):
            #print("[0]cpu:", pooler_cpu.shape, ":", pooler_cpu.dtype, ":", pooler_cpu)
            print("[0]1:", pooler_output1.dtype, ":", pooler_output1)
            print("[0]2:", pooler_output2.dtype, ":", pooler_output2)
            print("tensor not match")

    if hidden_states1 is None and hidden_states2 is None:
        print("hidden_states is None")
    else:
        hidden_states1 = hidden_states1.detach().cpu()
        hidden_states2 = hidden_states2.detach().cpu()
        print("tensor 1 type:", hidden_states_cpu.dtype, ":", hidden_states1.dtype)
        #print("[1]cpu:", hidden_states_cpu.shape, ":", hidden_states_cpu[0:8])
        #print("[1]1:", hidden_states1.shape, ":", hidden_states1[0:8])
        #print("[1]2:", hidden_states2.shape, ":", hidden_states2[0:8])
        if not torch.equal(hidden_states1, hidden_states_cpu):
            #print("[1]cpu:", hidden_states_cpu.shape, ":", hidden_states_cpu.dtype, ":", hidden_states_cpu)
            print("[1]1:", hidden_states1.dtype, ":", hidden_states1)
            print("[1]2:", hidden_states2.dtype, ":", hidden_states2)
            print("tensor not match")

    if attentions1 is None and attentions2 is None:
        print("attentions is None")
    else:
        attentions1 = attentions1.detach().cpu()
        attentions2 = attentions2.detach().cpu()
        #print("[2]cpu:", attentions_cpu.shape, ":", attentions_cpu)
        #print("[2]1:", attentions1.shape, ":", attentions1)
        #print("[2]2:", attentions2.shape, ":", attentions2)
        print("tensor 2 type:", attentions_cpu, ":", attentions1.dtype)
        if not torch.equal(attentions1, attentions_cpu):
            #print("[2]cpu:", attentions_cpu.shape, ":", attentions_cpu.dtype, ":", attentions_cpu)
            print("[2]1:", attentions1.dtype, ":", attentions1)
            print("[2]2:", attentions2.dtype, ":", attentions2[..., -1])
            print("tensor not match")

    #if torch.equal(npu_embedding1, npu_embedding2):
    #    print(".", end="", flush=True)
    #else:
    #    print(npu_embedding1[0])
    #    print(npu_embedding2[0], flush=True)

mindietorch.finalize()