import torch 
import torch_npu

from transformers import (
    Trainer,
    TrainingArguments,
    EvalPrediction,
    AutoModelForSequenceClassification,
    AutoTokenizer
)

from datasets import Dataset


import evaluate
import numpy as np

import pandas as pd


import os 
MODEL_PATH = os.environ.get('MODEL_PATH', '/data/ts/ascend/model')
DATASET_PATH = os.environ.get('DATASET_PATH', '/data/ts/ascend/datasets')
BUS_MODEL_PATH = os.environ.get('BUS_MODEL_PATH', '/data/ts/ascend/bus_models')



def train(body):
    # model_name = body['model_name']
    # pretrain_model = body['pretrain_model']
    # dataset = body['dataset']
    
    # train_path = body['train_path']
    # token_name = body['token_name']
    # batch_size = body['batch_size']
    # num_train_epochs = body['num_train_epochs']
    # basic_mode_type = body['basic_mode_type']
    # choose_dataset_size = body['choose_dataset_size']/100
    model_name = 'test'#body['model_name']
    pretrain_model = 'acge_text_embedding'#body['pretrain_model']
    dataset = 'ray-dataset/test.csv'
    
    batch_size = 4
    num_train_epochs = 4
    basic_mode_type = '1'
    choose_dataset_size = 10/100

    # 得到语料名
    dataset_name =dataset.split('/')[-1]
    dataset_path = os.path.join(DATASET_PATH, dataset_name)
    # os.system(f"mc mirror --overwrite --remove oss/ray-dataset/{dataset_name} {dataset_path}")
    # 标签处理
    df = pd.read_csv(dataset_path)
    df,num_labels,label_to_id,id_to_label = convert_label(df)
    # df.to_csv(dataset_path,index=False)
    # 暂时只支持csv
    raw_datasets = Dataset.from_pandas(df)
    # raw_train, raw_val = raw_datasets.split([1-choose_dataset_size, choose_dataset_size])

    #加载模型
    model_info = os.path.join(MODEL_PATH, pretrain_model) 
    os.makedirs(MODEL_PATH, exist_ok=True)
    # os.system(f"mc mirror --overwrite --remove oss/ray-model/{pretrain_model} {model_info}")
    tokenizer = AutoTokenizer.from_pretrained(model_info)
    model = AutoModelForSequenceClassification.from_pretrained(model_info,num_labels=num_labels, id2label=id_to_label, label2id=label_to_id, ignore_mismatched_sizes=True).to('npu:3')

    # token化
    def preprocess_function(examples):
        result = tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
        result["label"] = [int(l) for l in examples["label"]]  # 假设每个标签是一个列表，取第一个元素
        # result["label"] = examples["label"]
        nput_ids = torch.tensor(result.input_ids, device='npu:3')
        token_type_ids = torch.tensor(result.token_type_ids, device='npu:3')
        attention_mask = torch.tensor(result.attention_mask, device='npu:3')
        label = torch.tensor(result.label, device='npu:3')
        result2 = {}
        result2['label'] = label
        result2['input_ids'] = input_ids
        result2['token_type_ids'] = token_type_ids
        result2['attention_mask'] = attention_mask
        
        return result2
    
    dataset_train = raw_datasets.map(preprocess_function, batched=True)
    dataset_val = raw_datasets.map(preprocess_function, batched=True)
    dataset_train = dataset_train.remove_columns(['text'])
     
   

    output_info = os.path.join(BUS_MODEL_PATH, model_name) 
    training_args = TrainingArguments(
        output_dir=output_info,
        num_train_epochs=num_train_epochs,
        learning_rate=2e-5,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        evaluation_strategy="steps",
        eval_steps=100,
        save_strategy="epoch",
        logging_strategy="steps",
        logging_steps=50
    )

  
    trainer = Trainer(
        model=model,
        train_dataset=dataset_train,
        args=training_args,
    )

    print('开始训练')
    
    result = trainer.train()

    print('训练完成')

    bus_model_info = os.path.join(MODEL_PATH, model_name)   
    checkpoint_path = result.checkpoint.path
    os.system(
        f"cp -r {checkpoint_path}/checkpoint/ {bus_model_info}")
    os.system(
        f"mc cp -r {checkpoint_path}/checkpoint/ oss/ray-npu/{model_name}")

    return result


def convert_label(df):
    # 创建标签到ID的映射
    label_to_id = {}
    id_to_label = {}
    id_counter = 0
    # 遍历所有标签，创建映射
    for label in df['label'].unique():
        if label not in label_to_id:
            label_to_id[label] = str(id_counter)
            id_to_label[str(id_counter)] = label
            id_counter += 1
    
    # 将DataFrame中的标签替换为ID
    df['label'] = df['label'].map(label_to_id)
    return df,id_counter,label_to_id,id_to_label

if __name__ == "__main__":
    train(None)

开始训练
  3%|█▍                                         | 8/250 [00:01<00:32,  7.37it/s][E compiler_depend.ts:270] call aclnnFlashAttentionScore failed, detail:EZ9999: Inner Error!
EZ9999: 2025-01-20-03:22:58.402.844  Cannot find bin of op FlashAttentionScore, integral key 0/1/|float/ND/float/ND/float/ND/float/ND/float/ND/float/ND/float/ND/.
        TraceBack (most recent call last):
        Cannot find binary for op FlashAttentionScore.
        Kernel Run failed. opType: 100, FlashAttentionScore
        launch failed for FlashAttentionScore, errno:561000.

[ERROR] 2025-01-20-03:22:58 (PID:9022, Device:0, RankID:-1) ERR01005 OPS internal error
Exception raised from operator() at build/CMakeFiles/torch_npu.dir/compiler_depend.ts:459 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x68 (0xffff8de53898 in /home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x6c (0xffff8de0c2a8 in /home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/lib/libc10.so)
frame #2: <unknown function> + 0xd4ebe4 (0xfffda7cb8be4 in /home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so)
frame #3: <unknown function> + 0xe77bd0 (0xfffda7de1bd0 in /home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so)
frame #4: <unknown function> + 0x5ea594 (0xfffda7554594 in /home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so)
frame #5: <unknown function> + 0x5eaa48 (0xfffda7554a48 in /home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so)
frame #6: <unknown function> + 0x5e7f50 (0xfffda7551f50 in /home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch_npu/lib/libtorch_npu.so)
frame #7: <unknown function> + 0x946ec (0xffff8de7a6ec in /home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/lib/libc10.so)
frame #8: <unknown function> + 0x7a80 (0xffffa43a7a80 in /lib64/libpthread.so.0)
frame #9: <unknown function> + 0xe4d0c (0xffffa41dbd0c in /lib64/libc.so.6)


Traceback (most recent call last):
  File "/data/holly/hl-ascend-model/train/train1.py", line 175, in <module>
    train(None)
  File "/data/holly/hl-ascend-model/train/train1.py", line 144, in train
    result = trainer.train()
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/trainer.py", line 2171, in train
    return inner_training_loop(
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/trainer.py", line 2531, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/trainer.py", line 3678, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/trainer.py", line 3734, in compute_loss
    outputs = model(**inputs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py", line 1665, in forward
    outputs = self.bert(
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py", line 1142, in forward
    encoder_outputs = self.encoder(
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py", line 695, in forward
    layer_outputs = layer_module(
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py", line 627, in forward
    layer_output = apply_chunking_to_forward(
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/pytorch_utils.py", line 255, in apply_chunking_to_forward
    return forward_fn(*input_tensors)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py", line 640, in feed_forward_chunk
    layer_output = self.output(intermediate_output, attention_output)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/transformers/models/bert/modeling_bert.py", line 553, in forward
    hidden_states = self.dropout(hidden_states)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/modules/dropout.py", line 58, in forward
    return F.dropout(input, self.p, self.training, self.inplace)
  File "/home/ma-user/anaconda3/envs/MindSpore/lib/python3.9/site-packages/torch/nn/functional.py", line 1266, in dropout
    return _VF.dropout_(input, p, training) if inplace else _VF.dropout(input, p, training)
RuntimeError: The Inner error is reported as above. The process exits for this inner error, and the current working operator name is aclnnFlashAttentionScore.
Since the operator is called asynchronously, the stacktrace may be inaccurate. If you want to get the accurate stacktrace, pleace set the environment variable ASCEND_LAUNCH_BLOCKING=1.
[ERROR] 2025-01-20-03:22:58 (PID:9022, Device:0, RankID:-1) ERR00100 PTA call acl api failed
  3%|█▍                                         | 8/250 [00:03<01:43,  2.33it/s