text:
# coding:utf-8

import os
import json
import time
import pickle
import random
import warnings
import numpy as np
from torch.optim import Optimizer
from collections import defaultdict

import torch
import torch.nn as nn
from torch.optim import AdamW
from torch.backends import cudnn
from torch.optim.lr_scheduler import LambdaLR
from torch.optim.swa_utils import AveragedModel, SWALR

from transformers import (
    BertTokenizer,
    BertTokenizerFast
)
from modeling.modeling_nezha.modeling import NeZhaPreTrainedModel, NeZhaModel

class NeZhaSequenceClassification(NeZhaPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = 3
        self.bert = NeZhaModel(config)
        self.classifier = nn.Linear(config.hidden_size, self.num_labels)
        self.init_weights()

    def forward(
            self,
            input_ids=None,
            attention_mask=None,
            token_type_ids=None,
            position_ids=None,
            head_mask=None,
            inputs_embeds=None,
            labels=None
    ):

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )

        pooled_output = outputs[1]
        logits = self.classifier(pooled_output)

        outputs = (logits,) + outputs[2:]

        if labels is not None:
            # loss_fct = LabelSmoothingLoss(smoothing=0.01)
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            outputs = (loss,) + outputs

        return outputs

def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    return seed

def batch_loader(config, src, tgt, seg, mask):
    ins_num = src.size()[0]
    batch_size = config['batch_size']
    for i in range(ins_num // batch_size):
        src_batch = src[i * batch_size: (i + 1) * batch_size, :]
        tgt_batch = tgt[i * batch_size: (i + 1) * batch_size]
        seg_batch = seg[i * batch_size: (i + 1) * batch_size, :]
        mask_batch = mask[i * batch_size: (i + 1) * batch_size, :]
        yield src_batch, tgt_batch, seg_batch, mask_batch
    if ins_num > ins_num // batch_size * batch_size:
        src_batch = src[ins_num // batch_size * batch_size:, :]
        tgt_batch = tgt[ins_num // batch_size * batch_size:]
        seg_batch = seg[ins_num // batch_size * batch_size:, :]
        mask_batch = mask[ins_num // batch_size * batch_size:, :]
        yield src_batch, tgt_batch, seg_batch, mask_batch

def read_dataset(config, tokenizer):
    start = time.time()
    dataset = []
    seq_length = config['max_seq_len']

    with open(config['data_path'], 'r', encoding='utf-8') as f:
        for line_id, line in enumerate(f):
            sent_a, sent_b, tgt = line.strip().split('\t')
            src_a = tokenizer.convert_tokens_to_ids(['[CLS]'] + tokenizer.tokenize(sent_a) + ['[SEP]'])
            src_b = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(sent_b) + ['[SEP]'])
            src = src_a + src_b
            seg = [0] * len(src_a) + [1] * len(src_b)
            mask = [1] * len(src)
            tgt = int(tgt)
            if len(src) > seq_length:
                src = src[: seq_length]
                seg = seg[: seq_length]
                mask = mask[: seq_length]
            while len(src) < seq_length:
                src.append(0)
                seg.append(0)
                mask.append(0)
            dataset.append((src, tgt, seg, mask))

        data_cache_path = config['normal_data_cache_path']
        if not os.path.exists(os.path.dirname(data_cache_path)):
            os.makedirs(os.path.dirname(data_cache_path))
        with open(data_cache_path, 'wb') as f:
            pickle.dump(dataset, f)

    print("\n>> loading sentences from {},Time cost:{:.2f}".
          format(config['data_path'], ((time.time() - start) / 60.00)))

    return dataset

class LabelSmoothingLoss(nn.Module):
    """
    NLL loss with label smoothing.
    """

    def __init__(self, smoothing=0.01):
        super(LabelSmoothingLoss, self).__init__()
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing

    def forward(self, x, target):
        log_probs = torch.nn.functional.log_softmax(x, dim=-1)

        nll_loss = -log_probs.gather(dim=-1, index=target.unsqueeze(1))
        nll_loss = nll_loss.squeeze(1)
        smooth_loss = -log_probs.mean(dim=-1)
        loss = self.confidence * nll_loss + self.smoothing * smooth_loss
        return loss.mean()

class FGM:
    def __init__(self, config, model):
        self.model = model
        self.backup = {}
        self.emb_name = config['emb_name']
        self.epsilon = config['epsilon']

    def attack(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and self.emb_name in name:
                self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = self.epsilon * param.grad / norm
                    param.data.add_(r_at)

    def restore(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and self.emb_name in name:
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}

class PGD:
    def __init__(self, config, model):
        self.model = model
        self.emb_backup = {}
        self.grad_backup = {}
        self.epsilon = config['epsilon']
        self.emb_name = config['emb_name']
        self.alpha = config['alpha']

    def attack(self, is_first_attack=False):
        for name, param in self.model.named_parameters():
            if param.requires_grad and self.emb_name in name:
                if is_first_attack:
                    self.emb_backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = self.alpha * param.grad / norm
                    param.data.add_(r_at)
                    param.data = self.project(name, param.data, self.epsilon)

    def restore(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and self.emb_name in name:
                assert name in self.emb_backup
                param.data = self.emb_backup[name]
        self.emb_backup = {}

    def project(self, param_name, param_data, epsilon):
        r = param_data - self.emb_backup[param_name]
        if torch.norm(r) > epsilon:
            r = epsilon * r / torch.norm(r)
        return self.emb_backup[param_name] + r

    def backup_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None:
                self.grad_backup[name] = param.grad.clone()

    def restore_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None:
                param.grad = self.grad_backup[name]

class WarmupLinearSchedule(LambdaLR):
    def __init__(self, optimizer, warmup_steps, t_total, last_epoch=-1):
        self.warmup_steps = warmup_steps
        self.t_total = t_total
        super(WarmupLinearSchedule, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)

    def lr_lambda(self, step):
        if step < self.warmup_steps:
            return float(step) / float(max(1, self.warmup_steps))
        return max(0.0, float(self.t_total - step) / float(max(1.0, self.t_total - self.warmup_steps)))

def block_shuffle(config, train_set):
    bs = config['batch_size'] * 100
    num_block = int(len(train_set) / bs)
    slice_ = num_block * bs

    train_set_to_shuffle = train_set[:slice_]
    train_set_left = train_set[slice_:]

    sorted_train_set = sorted(train_set_to_shuffle, key=lambda i: len(i[0]))
    shuffled_train_set = []

    tmp = []
    for i in range(len(sorted_train_set)):
        tmp.append(sorted_train_set[i])
        if (i+1) % bs == 0:
            random.shuffle(tmp)
            shuffled_train_set.extend(tmp)
            tmp = []

    random.shuffle(train_set_left)
    shuffled_train_set.extend(train_set_left)

    return shuffled_train_set

def build_model_and_tokenizer(config):
    tokenizer_path = config['model_path'] + '/vocab.txt'
    if config['tokenizer_fast']:
        tokenizer = BertTokenizerFast.from_pretrained(tokenizer_path)
    else:
        tokenizer = BertTokenizer.from_pretrained(tokenizer_path)
    if config['use_model'] == 'nezha':
        model = NeZhaSequenceClassification.from_pretrained(config['model_path'])
    return tokenizer, model

def build_optimizer(config, model, train_steps):
    param_optimizer = list(model.named_parameters())
    no_decay = ['bias', 'LayerNorm.weight']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)],
         'weight_decay_rate': config['weight_decay']},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)],
         'weight_decay_rate': 0.0}
    ]
    optimizer = AdamW(optimizer_grouped_parameters, lr=config['learning_rate'], eps=1e-8)
    scheduler = WarmupLinearSchedule(optimizer, warmup_steps=train_steps * config['warmup_ratio'],
                                     t_total=train_steps)

    return optimizer, scheduler

def main():
    config = {
        'use_model': 'nezha',
        'normal_data_cache_path': 'user_data/processed/nezha/all_data.pkl',  # 96423
        'data_path': 'data/train.txt',
        'output_path': 'output_model',
        'model_path': 'pretrain_code/new_model',  # your pretrain model path
        'shuffle_way': '',
        'use_swa': True,
        'tokenizer_fast': False,
        'batch_size': 32,
        'num_epochs': 3,
        'max_seq_len': 128,
        'learning_rate': 2e-5,
        'alpha': 0.3,
        'epsilon': 1.0,
        'adv_k': 3,
        'emb_name': 'word_embeddings.',
        'adv': 'pgd',
        'warmup_ratio': 0.1,
        'weight_decay': 0.01,
        'device': 'cuda',
        'logging_step': 1000,
        'seed': 2021
    }

    warnings.filterwarnings('ignore')
    localtime_start = time.asctime(time.localtime(time.time()))
    print(">> program start at:{}".format(localtime_start))
    print("\n>> loading model from :{}".format(config['model_path']))

    tokenizer, model = build_model_and_tokenizer(config)
    if not os.path.exists(config['normal_data_cache_path']):
        train_set = read_dataset(config, tokenizer)
    else:
        with open(config['normal_data_cache_path'], 'rb') as f:
            train_set = pickle.load(f)

    seed_everything(config['seed'])

    if config['shuffle_way'] == 'block_shuffle':
        train_set = block_shuffle(config, train_set)
    else:
        random.shuffle(train_set)

    train_num = len(train_set)
    train_steps = int(train_num * config['num_epochs'] / config['batch_size']) + 1

    optimizer, scheduler = build_optimizer(config, model, train_steps)
    model.to(config['device'])

    src = torch.LongTensor([example[0] for example in train_set])
    tgt = torch.LongTensor([example[1] for example in train_set])
    seg = torch.LongTensor([example[2] for example in train_set])
    mask = torch.LongTensor([example[3] for example in train_set])

    cudnn.benchmark = True

    total_loss, cur_avg_loss = 0.0, 0.0
    global_steps = 0

    if config['adv'] == '':
        print('\n>> start normal training ...')
    elif config['adv'] == 'fgm':
        print('\n>> start fgm training ...')
    elif config['adv'] == 'pgd':
        print('\n>> start pgd training ...')

    start = time.time()
    for epoch in range(1, config['num_epochs'] + 1):
        model.train()
        for i, (src_batch, tgt_batch, seg_batch, mask_batch) \
                in enumerate(batch_loader(config, src, tgt, seg, mask)):

            src_batch = src_batch.to(config['device'])
            tgt_batch = tgt_batch.to(config['device'])
            seg_batch = seg_batch.to(config['device'])
            mask_batch = mask_batch.to(config['device'])

            output = model(input_ids=src_batch, labels=tgt_batch,
                           token_type_ids=seg_batch, attention_mask=mask_batch)
            loss = output[0]
            optimizer.zero_grad()
            loss.backward()

            total_loss += loss.item()
            cur_avg_loss += loss.item()

            if config['adv'] == 'fgm':
                fgm = FGM(config, model)
                fgm.attack()
                adv_loss = model(input_ids=src_batch, labels=tgt_batch,
                                 token_type_ids=seg_batch, attention_mask=mask_batch)[0]
                adv_loss.backward()
                fgm.restore()

            if config['adv'] == 'pgd':
                pgd = PGD(config, model)
                K = config['adv_k']
                pgd.backup_grad()
                for t in range(K):
                    pgd.attack(is_first_attack=(t == 0))
                    if t != K - 1:
                        model.zero_grad()
                    else:
                        pgd.restore_grad()
                    adv_loss = model(input_ids=src_batch, labels=tgt_batch,
                                     token_type_ids=seg_batch, attention_mask=mask_batch)[0]
                    adv_loss.backward()
                pgd.restore()
            optimizer.step()

            scheduler.step()
            model.zero_grad()

            if (i + 1) % config['logging_step'] == 0:
                print("\n>> epoch - {}, epoch steps - {}, global steps - {}, "
                      "epoch avg loss - {:.4f}, global avg loss - {:.4f}, time cost - {:.2f} min".format
                      (epoch, i + 1, global_steps + 1, cur_avg_loss / config['logging_step'],
                       total_loss / (global_steps + 1),
                       (time.time() - start) / 60.00))
                cur_avg_loss = 0.0
            global_steps += 1

        model_save_path = os.path.join(config['output_path'], f'checkpoint-{global_steps}')
        model_to_save = model.module if hasattr(model, 'module') else model
        print('\n>> model saved ... ...')
        model_to_save.save_pretrained(model_save_path)

        conf = json.dumps(config)
        out_conf_path = os.path.join(config['output_path'], f'checkpoint-{global_steps}' +
                                     '/train_config.json')
        with open(out_conf_path, 'w', encoding='utf-8') as f:
            f.write(conf)

    localtime_end = time.asctime(time.localtime(time.time()))
    print("\n>> program end at:{}".format(localtime_end))

if __name__ == '__main__':
    main()


text:
# ccks2021_track3_baseline
 a baseline to practice
路径可能会有问题，自己改改

torch==1.7.1

pyhton==3.7.1

transformers==4.7.0

cuda==11.0

this is a baseline, you can fix params to improve your score

#
some useful tricks:

1.https://github.com/alphadl/lookahead.pytorch

2.https://github.com/timgaripov/swa

3.https://github.com/lonePatient/multi-sample_dropout_pytorch

4.https://github.com/lonePatient/albert_pytorch/blob/master/prepare_lm_data_ngram.py

#
pretrain model download:
https://github.com/lonePatient/NeZha_Chinese_PyTorch

#
address: https://tianchi.aliyun.com/competition/entrance/531901/score

#
steps:

1.run pretrain_code/run_pretrain.py

2.run run_classify.py

3.run run_predictor.py

4.upload your submit file

repo_language: Python

mime_type: text/plain


text:
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
  <component name="ProjectModuleManager">
    <modules>
      <module fileurl="file://$PROJECT_DIR$/.idea/baseline.iml" filepath="$PROJECT_DIR$/.idea/baseline.iml" />
    </modules>
  </component>
</project>


text:
<?xml version="1.0" encoding="UTF-8"?>
<project version="4">
  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.7 (torch_pure) (4)" project-jdk-type="Python SDK" />
</project>


text:
<?xml version="1.0" encoding="UTF-8"?>
<module type="PYTHON_MODULE" version="4">
  <component name="NewModuleRootManager">
    <content url="file://$MODULE_DIR$" />
    <orderEntry type="jdk" jdkName="Python 3.7 (torch_pure) (4)" jdkType="Python SDK" />
    <orderEntry type="sourceFolder" forTests="false" />
  </component>
  <component name="PyDocumentationSettings">
    <option name="format" value="PLAIN" />
    <option name="myDocStringFormat" value="Plain" />
  </component>
  <component name="TestRunnerService">
    <option name="PROJECT_TEST_RUNNER" value="pytest" />
  </component>
</module>


text:
<component name="InspectionProjectProfileManager">
  <settings>
    <option name="USE_PROJECT_PROFILE" value="false" />
    <version value="1.0" />
  </settings>
</component>


text:
<component name="InspectionProjectProfileManager">
  <profile version="1.0">
    <option name="myName" value="Project Default" />
    <inspection_tool class="PyPackageRequirementsInspection" enabled="true" level="WARNING" enabled_by_default="true">
      <option name="ignoredPackages">
        <value>
          <list size="282">
            <item index="0" class="java.lang.String" itemvalue="web.py" />
            <item index="1" class="java.lang.String" itemvalue="pybloom" />
            <item index="2" class="java.lang.String" itemvalue="tensorflow" />
            <item index="3" class="java.lang.String" itemvalue="absl-py" />
            <item index="4" class="java.lang.String" itemvalue="protobuf" />
            <item index="5" class="java.lang.String" itemvalue="scipy" />
            <item index="6" class="java.lang.String" itemvalue="six" />
            <item index="7" class="java.lang.String" itemvalue="Werkzeug" />
            <item index="8" class="java.lang.String" itemvalue="opencv-python" />
            <item index="9" class="java.lang.String" itemvalue="h5py" />
            <item index="10" class="java.lang.String" itemvalue="python-dateutil" />
            <item index="11" class="java.lang.String" itemvalue="astor" />
            <item index="12" class="java.lang.String" itemvalue="kiwisolver" />
            <item index="13" class="java.lang.String" itemvalue="gast" />
            <item index="14" class="java.lang.String" itemvalue="numpy" />
            <item index="15" class="java.lang.String" itemvalue="Keras-Applications" />
            <item index="16" class="java.lang.String" itemvalue="Keras-Preprocessing" />
            <item index="17" class="java.lang.String" itemvalue="termcolor" />
            <item index="18" class="java.lang.String" itemvalue="tensorboard" />
            <item index="19" class="java.lang.String" itemvalue="matplotlib" />
            <item index="20" class="java.lang.String" itemvalue="tensorboardX" />
            <item index="21" class="java.lang.String" itemvalue="grpcio" />
            <item index="22" class="java.lang.String" itemvalue="pytz" />
            <item index="23" class="java.lang.String" itemvalue="pyparsing" />
            <item index="24" class="java.lang.String" itemvalue="Markdown" />
            <item index="25" class="java.lang.String" itemvalue="Pillow" />
            <item index="26" class="java.lang.String" itemvalue="gensim" />
            <item index="27" class="java.lang.String" itemvalue="numba" />
            <item index="28" class="java.lang.String" itemvalue="tensorflow-estimator" />
            <item index="29" class="java.lang.String" itemvalue="greenlet" />
            <item index="30" class="java.lang.String" itemvalue="Babel" />
            <item index="31" class="java.lang.String" itemvalue="scikit-learn" />
            <item index="32" class="java.lang.String" itemvalue="testpath" />
            <item index="33" class="java.lang.String" itemvalue="backports.os" />
            <item index="34" class="java.lang.String" itemvalue="py" />
            <item index="35" class="java.lang.String" itemvalue="patsy" />
            <item index="36" class="java.lang.String" itemvalue="ipython-genutils" />
            <item index="37" class="java.lang.String" itemvalue="mccabe" />
            <item index="38" class="java.lang.String" itemvalue="bleach" />
            <item index="39" class="java.lang.String" itemvalue="HeapDict" />
            <item index="40" class="java.lang.String" itemvalue="lxml" />
            <item index="41" class="java.lang.String" itemvalue="soupsieve" />
            <item index="42" class="java.lang.String" itemvalue="jsonschema" />
            <item index="43" class="java.lang.String" itemvalue="xlrd" />
            <item index="44" class="java.lang.String" itemvalue="anaconda-project" />
            <item index="45" class="java.lang.String" itemvalue="menuinst" />
            <item index="46" class="java.lang.String" itemvalue="boto3" />
            <item index="47" class="java.lang.String" itemvalue="regex" />
            <item index="48" class="java.lang.String" itemvalue="fastcache" />
            <item index="49" class="java.lang.String" itemvalue="imageio" />
            <item index="50" class="java.lang.String" itemvalue="pytest-remotedata" />
            <item index="51" class="java.lang.String" itemvalue="Keras" />
            <item index="52" class="java.lang.String" itemvalue="idna" />
            <item index="53" class="java.lang.String" itemvalue="Bottleneck" />
            <item index="54" class="java.lang.String" itemvalue="rsa" />
            <item index="55" class="java.lang.String" itemvalue="networkx" />
            <item index="56" class="java.lang.String" itemvalue="pycurl" />
            <item index="57" class="java.lang.String" itemvalue="json5" />
            <item index="58" class="java.lang.String" itemvalue="pluggy" />
            <item index="59" class="java.lang.String" itemvalue="cffi" />
            <item index="60" class="java.lang.String" itemvalue="pep8" />
            <item index="61" class="java.lang.String" itemvalue="jdcal" />
            <item index="62" class="java.lang.String" itemvalue="alabaster" />
            <item index="63" class="java.lang.String" itemvalue="jupyter" />
            <item index="64" class="java.lang.String" itemvalue="pyOpenSSL" />
            <item index="65" class="java.lang.String" itemvalue="PyWavelets" />
            <item index="66" class="java.lang.String" itemvalue="smart-open" />
            <item index="67" class="java.lang.String" itemvalue="prompt-toolkit" />
            <item index="68" class="java.lang.String" itemvalue="QtAwesome" />
            <item index="69" class="java.lang.String" itemvalue="bert4keras" />
            <item index="70" class="java.lang.String" itemvalue="glob2" />
            <item index="71" class="java.lang.String" itemvalue="Send2Trash" />
            <item index="72" class="java.lang.String" itemvalue="imagesize" />
            <item index="73" class="java.lang.String" itemvalue="et-xmlfile" />
            <item index="74" class="java.lang.String" itemvalue="pathlib2" />
            <item index="75" class="java.lang.String" itemvalue="spyder" />
            <item index="76" class="java.lang.String" itemvalue="pylint" />
            <item index="77" class="java.lang.String" itemvalue="statsmodels" />
            <item index="78" class="java.lang.String" itemvalue="isort" />
            <item index="79" class="java.lang.String" itemvalue="olefile" />
            <item index="80" class="java.lang.String" itemvalue="keyring" />
            <item index="81" class="java.lang.String" itemvalue="unicodecsv" />
            <item index="82" class="java.lang.String" itemvalue="pytest-astropy" />
            <item index="83" class="java.lang.String" itemvalue="traitlets" />
            <item index="84" class="java.lang.String" itemvalue="pywinpty" />
            <item index="85" class="java.lang.String" itemvalue="joblib" />
            <item index="86" class="java.lang.String" itemvalue="nltk" />
            <item index="87" class="java.lang.String" itemvalue="atomicwrites" />
            <item index="88" class="java.lang.String" itemvalue="partd" />
            <item index="89" class="java.lang.String" itemvalue="Click" />
            <item index="90" class="java.lang.String" itemvalue="hanziconv" />
            <item index="91" class="java.lang.String" itemvalue="fsspec" />
            <item index="92" class="java.lang.String" itemvalue="filelock" />
            <item index="93" class="java.lang.String" itemvalue="numpydoc" />
            <item index="94" class="java.lang.String" itemvalue="pyzmq" />
            <item index="95" class="java.lang.String" itemvalue="sentencepiece" />
            <item index="96" class="java.lang.String" itemvalue="astropy" />
            <item index="97" class="java.lang.String" itemvalue="entrypoints" />
            <item index="98" class="java.lang.String" itemvalue="anaconda-navigator" />
            <item index="99" class="java.lang.String" itemvalue="bkcharts" />
            <item index="100" class="java.lang.String" itemvalue="sphinxcontrib-websupport" />
            <item index="101" class="java.lang.String" itemvalue="google-api-core" />
            <item index="102" class="java.lang.String" itemvalue="beautifulsoup4" />
            <item index="103" class="java.lang.String" itemvalue="path.py" />
            <item index="104" class="java.lang.String" itemvalue="tokenizers" />
            <item index="105" class="java.lang.String" itemvalue="clyent" />
            <item index="106" class="java.lang.String" itemvalue="navigator-updater" />
            <item index="107" class="java.lang.String" itemvalue="libarchive-c" />
            <item index="108" class="java.lang.String" itemvalue="cryptography" />
            <item index="109" class="java.lang.String" itemvalue="pyahocorasick" />
            <item index="110" class="java.lang.String" itemvalue="widgetsnbextension" />
            <item index="111" class="java.lang.String" itemvalue="multipledispatch" />
            <item index="112" class="java.lang.String" itemvalue="numexpr" />
            <item index="113" class="java.lang.String" itemvalue="jupyter-core" />
            <item index="114" class="java.lang.String" itemvalue="rope" />
            <item index="115" class="java.lang.String" itemvalue="wcwidth" />
            <item index="116" class="java.lang.String" itemvalue="llvmlite" />
            <item index="117" class="java.lang.String" itemvalue="rouge" />
            <item index="118" class="java.lang.String" itemvalue="Jinja2" />
            <item index="119" class="java.lang.String" itemvalue="pycrypto" />
            <item index="120" class="java.lang.String" itemvalue="google-resumable-media" />
            <item index="121" class="java.lang.String" itemvalue="conda" />
            <item index="122" class="java.lang.String" itemvalue="sortedcollections" />
            <item index="123" class="java.lang.String" itemvalue="pyflakes" />
            <item index="124" class="java.lang.String" itemvalue="asn1crypto" />
            <item index="125" class="java.lang.String" itemvalue="parso" />
            <item index="126" class="java.lang.String" itemvalue="pytest-doctestplus" />
            <item index="127" class="java.lang.String" itemvalue="ipython" />
            <item index="128" class="java.lang.String" itemvalue="xlwt" />
            <item index="129" class="java.lang.String" itemvalue="packaging" />
            <item index="130" class="java.lang.String" itemvalue="JPype1" />
            <item index="131" class="java.lang.String" itemvalue="chardet" />
            <item index="132" class="java.lang.String" itemvalue="jmespath" />
            <item index="133" class="java.lang.String" itemvalue="s3transfer" />
            <item index="134" class="java.lang.String" itemvalue="sphinxcontrib-jsmath" />
            <item index="135" class="java.lang.String" itemvalue="PyYAML" />
            <item index="136" class="java.lang.String" itemvalue="pickleshare" />
            <item index="137" class="java.lang.String" itemvalue="defusedxml" />
            <item index="138" class="java.lang.String" itemvalue="pycparser" />
            <item index="139" class="java.lang.String" itemvalue="pyasn1-modules" />
            <item index="140" class="java.lang.String" itemvalue="tables" />
            <item index="141" class="java.lang.String" itemvalue="Pygments" />
            <item index="142" class="java.lang.String" itemvalue="sphinxcontrib-qthelp" />
            <item index="143" class="java.lang.String" itemvalue="docutils" />
            <item index="144" class="java.lang.String" itemvalue="gevent" />
            <item index="145" class="java.lang.String" itemvalue="comtypes" />
            <item index="146" class="java.lang.String" itemvalue="pywin32" />
            <item index="147" class="java.lang.String" itemvalue="qtconsole" />
            <item index="148" class="java.lang.String" itemvalue="terminado" />
            <item index="149" class="java.lang.String" itemvalue="transformers" />
            <item index="150" class="java.lang.String" itemvalue="distributed" />
            <item index="151" class="java.lang.String" itemvalue="jupyter-client" />
            <item index="152" class="java.lang.String" itemvalue="ipykernel" />
            <item index="153" class="java.lang.String" itemvalue="nbconvert" />
            <item index="154" class="java.lang.String" itemvalue="attrs" />
            <item index="155" class="java.lang.String" itemvalue="conda-package-handling" />
            <item index="156" class="java.lang.String" itemvalue="psutil" />
            <item index="157" class="java.lang.String" itemvalue="jedi" />
            <item index="158" class="java.lang.String" itemvalue="cytoolz" />
            <item index="159" class="java.lang.String" itemvalue="backports.tempfile" />
            <item index="160" class="java.lang.String" itemvalue="msgpack" />
            <item index="161" class="java.lang.String" itemvalue="decorator" />
            <item index="162" class="java.lang.String" itemvalue="pandocfilters" />
            <item index="163" class="java.lang.String" itemvalue="backports.shutil-get-terminal-size" />
            <item index="164" class="java.lang.String" itemvalue="google-cloud-storage" />
            <item index="165" class="java.lang.String" itemvalue="pycodestyle" />
            <item index="166" class="java.lang.String" itemvalue="pycosat" />
            <item index="167" class="java.lang.String" itemvalue="pyasn1" />
            <item index="168" class="java.lang.String" itemvalue="requests" />
            <item index="169" class="java.lang.String" itemvalue="bitarray" />
            <item index="170" class="java.lang.String" itemvalue="pyrsistent" />
            <item index="171" class="java.lang.String" itemvalue="sphinxcontrib-devhelp" />
            <item index="172" class="java.lang.String" itemvalue="mkl-fft" />
            <item index="173" class="java.lang.String" itemvalue="XlsxWriter" />
            <item index="174" class="java.lang.String" itemvalue="seaborn" />
            <item index="175" class="java.lang.String" itemvalue="zipp" />
            <item index="176" class="java.lang.String" itemvalue="pkginfo" />
            <item index="177" class="java.lang.String" itemvalue="itsdangerous" />
            <item index="178" class="java.lang.String" itemvalue="pyreadline" />
            <item index="179" class="java.lang.String" itemvalue="ipywidgets" />
            <item index="180" class="java.lang.String" itemvalue="tornado" />
            <item index="181" class="java.lang.String" itemvalue="botocore" />
            <item index="182" class="java.lang.String" itemvalue="torch" />
            <item index="183" class="java.lang.String" itemvalue="pyhanlp" />
            <item index="184" class="java.lang.String" itemvalue="singledispatch" />
            <item index="185" class="java.lang.String" itemvalue="sortedcontainers" />
            <item index="186" class="java.lang.String" itemvalue="mistune" />
            <item index="187" class="java.lang.String" itemvalue="pandas" />
            <item index="188" class="java.lang.String" itemvalue="toolz" />
            <item index="189" class="java.lang.String" itemvalue="Sphinx" />
            <item index="190" class="java.lang.String" itemvalue="future" />
            <item index="191" class="java.lang.String" itemvalue="mpmath" />
            <item index="192" class="java.lang.String" itemvalue="jupyter-console" />
            <item index="193" class="java.lang.String" itemvalue="ruamel-yaml" />
            <item index="194" class="java.lang.String" itemvalue="bokeh" />
            <item index="195" class="java.lang.String" itemvalue="cachetools" />
            <item index="196" class="java.lang.String" itemvalue="webencodings" />
            <item index="197" class="java.lang.String" itemvalue="sphinxcontrib-applehelp" />
            <item index="198" class="java.lang.String" itemvalue="html5lib" />
            <item index="199" class="java.lang.String" itemvalue="missingno" />
            <item index="200" class="java.lang.String" itemvalue="googleapis-common-protos" />
            <item index="201" class="java.lang.String" itemvalue="spyder-kernels" />
            <item index="202" class="java.lang.String" itemvalue="QtPy" />
            <item index="203" class="java.lang.String" itemvalue="astroid" />
            <item index="204" class="java.lang.String" itemvalue="MarkupSafe" />
            <item index="205" class="java.lang.String" itemvalue="mkl-random" />
            <item index="206" class="java.lang.String" itemvalue="pytest-arraydiff" />
            <item index="207" class="java.lang.String" itemvalue="locket" />
            <item index="208" class="java.lang.String" itemvalue="snowballstemmer" />
            <item index="209" class="java.lang.String" itemvalue="contextlib2" />
            <item index="210" class="java.lang.String" itemvalue="conda-verify" />
            <item index="211" class="java.lang.String" itemvalue="certifi" />
            <item index="212" class="java.lang.String" itemvalue="sympy" />
            <item index="213" class="java.lang.String" itemvalue="notebook" />
            <item index="214" class="java.lang.String" itemvalue="xlwings" />
            <item index="215" class="java.lang.String" itemvalue="win-unicode-console" />
            <item index="216" class="java.lang.String" itemvalue="sacremoses" />
            <item index="217" class="java.lang.String" itemvalue="pyodbc" />
            <item index="218" class="java.lang.String" itemvalue="boto" />
            <item index="219" class="java.lang.String" itemvalue="wrapt" />
            <item index="220" class="java.lang.String" itemvalue="sphinxcontrib-htmlhelp" />
            <item index="221" class="java.lang.String" itemvalue="pytest-openfiles" />
            <item index="222" class="java.lang.String" itemvalue="anaconda-client" />
            <item index="223" class="java.lang.String" itemvalue="backcall" />
            <item index="224" class="java.lang.String" itemvalue="PySocks" />
            <item index="225" class="java.lang.String" itemvalue="backports.functools-lru-cache" />
            <item index="226" class="java.lang.String" itemvalue="mock" />
            <item index="227" class="java.lang.String" itemvalue="dask" />
            <item index="228" class="java.lang.String" itemvalue="google-cloud-core" />
            <item index="229" class="java.lang.String" itemvalue="scikit-image" />
            <item index="230" class="java.lang.String" itemvalue="jieba" />
            <item index="231" class="java.lang.String" itemvalue="more-itertools" />
            <item index="232" class="java.lang.String" itemvalue="SQLAlchemy" />
            <item index="233" class="java.lang.String" itemvalue="tblib" />
            <item index="234" class="java.lang.String" itemvalue="cloudpickle" />
            <item index="235" class="java.lang.String" itemvalue="importlib-metadata" />
            <item index="236" class="java.lang.String" itemvalue="simplegeneric" />
            <item index="237" class="java.lang.String" itemvalue="backports.weakref" />
            <item index="238" class="java.lang.String" itemvalue="zict" />
            <item index="239" class="java.lang.String" itemvalue="urllib3" />
            <item index="240" class="java.lang.String" itemvalue="jupyterlab" />
            <item index="241" class="java.lang.String" itemvalue="Cython" />
            <item index="242" class="java.lang.String" itemvalue="Flask" />
            <item index="243" class="java.lang.String" itemvalue="nose" />
            <item index="244" class="java.lang.String" itemvalue="pytest" />
            <item index="245" class="java.lang.String" itemvalue="sphinxcontrib-serializinghtml" />
            <item index="246" class="java.lang.String" itemvalue="jupyterlab-server" />
            <item index="247" class="java.lang.String" itemvalue="conda-build" />
            <item index="248" class="java.lang.String" itemvalue="nbformat" />
            <item index="249" class="java.lang.String" itemvalue="prometheus-client" />
            <item index="250" class="java.lang.String" itemvalue="tqdm" />
            <item index="251" class="java.lang.String" itemvalue="win-inet-pton" />
            <item index="252" class="java.lang.String" itemvalue="lazy-object-proxy" />
            <item index="253" class="java.lang.String" itemvalue="colorama" />
            <item index="254" class="java.lang.String" itemvalue="ply" />
            <item index="255" class="java.lang.String" itemvalue="google-auth" />
            <item index="256" class="java.lang.String" itemvalue="openpyxl" />
            <item index="257" class="java.lang.String" itemvalue="pymysql" />
            <item index="258" class="java.lang.String" itemvalue="tensorflow_gpu" />
            <item index="259" class="java.lang.String" itemvalue="scikit_learn" />
            <item index="260" class="java.lang.String" itemvalue="pytorch" />
            <item index="261" class="java.lang.String" itemvalue="autopep8" />
            <item index="262" class="java.lang.String" itemvalue="google-pasta" />
            <item index="263" class="java.lang.String" itemvalue="pkuseg" />
            <item index="264" class="java.lang.String" itemvalue="google-auth-oauthlib" />
            <item index="265" class="java.lang.String" itemvalue="typed-ast" />
            <item index="266" class="java.lang.String" itemvalue="opt-einsum" />
            <item index="267" class="java.lang.String" itemvalue="toml" />
            <item index="268" class="java.lang.String" itemvalue="tensorflow-gpu-estimator" />
            <item index="269" class="java.lang.String" itemvalue="streamz" />
            <item index="270" class="java.lang.String" itemvalue="astunparse" />
            <item index="271" class="java.lang.String" itemvalue="snownlp" />
            <item index="272" class="java.lang.String" itemvalue="tensorflow-gpu" />
            <item index="273" class="java.lang.String" itemvalue="argparse" />
            <item index="274" class="java.lang.String" itemvalue="wikipedia" />
            <item index="275" class="java.lang.String" itemvalue="tb-nightly" />
            <item index="276" class="java.lang.String" itemvalue="keras" />
            <item index="277" class="java.lang.String" itemvalue="thulac" />
            <item index="278" class="java.lang.String" itemvalue="pytorch-pretrained-bert" />
            <item index="279" class="java.lang.String" itemvalue="apex" />
            <item index="280" class="java.lang.String" itemvalue="cxxfilt" />
            <item index="281" class="java.lang.String" itemvalue="simpletransformers" />
          </list>
        </value>
      </option>
    </inspection_tool>
    <inspection_tool class="PyPep8NamingInspection" enabled="true" level="WEAK WARNING" enabled_by_default="true">
      <option name="ignoredErrors">
        <list>
          <option value="N803" />
        </list>
      </option>
    </inspection_tool>
    <inspection_tool class="PyUnresolvedReferencesInspection" enabled="true" level="WARNING" enabled_by_default="true">
      <option name="ignoredIdentifiers">
        <list>
          <option value="demo.paquURL.crawler" />
          <option value="tuple.scatter3D" />
          <option value="tuple.view_init" />
          <option value="tuple.set_xlabel" />
          <option value="tuple.set_ylabel" />
          <option value="tuple.set_zlabel" />
          <option value="data.IMDB.load_data" />
          <option value="transformers.convert_bert_original_tf_checkpoint_to_pytorch" />
          <option value="modeling.modeling_nezha.configuration.NeZhaConfig.__getitem__" />
        </list>
      </option>
    </inspection_tool>
  </profile>
</component>


text:
# coding:utf-8

import os
import pickle

import torch
import random
import warnings
import numpy as np
import pandas as pd
from tqdm import tqdm
from typing import Dict
from collections import defaultdict
from torch.utils.data import Dataset

from transformers import (
    BertTokenizer,
    DataCollatorForLanguageModeling,
    DataCollatorForWholeWordMask,
    PreTrainedTokenizer,
    BertConfig,
    BertForMaskedLM
)
from transformers.utils import logging

from modeling.modeling_nezha.modeling import NeZhaForMaskedLM, NeZhaConfig
from simple_trainer import Trainer
from pretrain_args import TrainingArguments

warnings.filterwarnings('ignore')
logger = logging.get_logger(__name__)

def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    return seed

def read_data(config, train_file_path, test_file_path, tokenizer: BertTokenizer) -> dict:
    train_df = pd.read_csv(train_file_path, header=None, sep='\t')
    test_df = pd.read_csv(test_file_path, header=None, sep='\t')

    pretrain_df = pd.concat([train_df, test_df], axis=0)

    inputs = defaultdict(list)
    for i, row in tqdm(pretrain_df.iterrows(), desc=f'preprocessing pretrain data ... ...', total=len(pretrain_df)):
        sentence_a, sentence_b = row[0], row[1]
        inputs_dict = tokenizer.encode_plus(sentence_a, sentence_b, add_special_tokens=True,
                                            return_token_type_ids=True, return_attention_mask=True)
        inputs['input_ids'].append(inputs_dict['input_ids'])
        inputs['token_type_ids'].append(inputs_dict['token_type_ids'])
        inputs['attention_mask'].append(inputs_dict['attention_mask'])

    data_cache_path = config['data_cache_path']

    if not os.path.exists(os.path.dirname(data_cache_path)):
        os.makedirs(os.path.dirname(data_cache_path))
    with open(data_cache_path, 'wb') as f:
        pickle.dump(inputs, f)

    return inputs

class LineByLineTextDataset(Dataset):
    def __init__(self, tokenizer: PreTrainedTokenizer, train_file_path: str, test_file_path: str, block_size: int):
        assert os.path.isfile(train_file_path), f"Input file path {train_file_path} not found"
        logger.info(f"Creating features from dataset file at {train_file_path}")

        assert os.path.isfile(test_file_path), f"Input file path {test_file_path} not found"
        logger.info(f"Creating features from dataset file at {test_file_path}")

        with open(train_file_path, encoding="utf-8") as f:
            train_lines = [line for line in f.read().splitlines() if (len(line) > 0 and not line.isspace())]

        with open(test_file_path, encoding="utf-8") as f:
            test_lines = [line for line in f.read().splitlines() if (len(line) > 0 and not line.isspace())]

        lines = train_lines + test_lines

        batch_encoding = tokenizer(lines, add_special_tokens=True, truncation=True, max_length=block_size)
        self.examples = batch_encoding["input_ids"]
        self.examples = [{"input_ids": torch.tensor(e, dtype=torch.long)} for e in self.examples]

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i) -> Dict[str, torch.tensor]:
        return self.examples[i]

def main():
    """
    download pretrain model from https://github.com/lonePatient/NeZha_Chinese_PyTorch,
    we only use pretrain model name : nezha-cn-base, nezha-base-wwm
    """
    config = {
        'pretrain_type': 'dynamic_mask',  # dynamic_mask, whole_word_mask
        'data_cache_path': '',
        'train_data_path': '../data/train.txt',
        'test_data_path': '../data/test.txt',
    }

    mlm_probability = 0.15
    num_train_epochs = 150
    seq_length = 100
    batch_size = 32
    learning_rate = 6e-5
    save_steps = 5000
    seed = 2021

    """
    if use bert, replace NeZhaForMaskedLM -> BertForMaskedLM,
                         NeZhaConfig -> BertConfig
                         model_name -> your pretrain_model name

    example:
            >>      model_name = bert-base-uncased
                    model_path = '../pretrain_model/' + model_name + '/pytorch_model.bin'
                    config_path = '../pretrain_model/' + model_name + '/config.json'
                    vocab_file = '../pretrain_model/' + model_name + '/vocab.txt'
                    ... ...
                    model_config = BertConfig.from_pretrained(config_path)
                    BertForMaskedLM.from_pretrained(pretrained_model_name_or_path=model_path,
                                                 config=model_config)
    """

    # put dowm your file path
    if config['pretrain_type'] == 'whole_word_mask':
        model_name = 'nezha-base-wwm'
    else:
        model_name = 'nezha-cn-base'

    config['data_cache_path'] = '../user_data/pretrain/'+config['pretrain_type']+'/data.pkl'

    model_path = '../pretrain_model/' + model_name + '/pytorch_model.bin'
    config_path = '../pretrain_model/' + model_name + '/config.json'

    vocab_file = '../pretrain_model/' + model_name + '/vocab.txt'
    tokenizer = BertTokenizer.from_pretrained(vocab_file)

    model_config = NeZhaConfig.from_pretrained(config_path)

    assert os.path.isfile(model_path), f"Input file path {model_path} not found, " \
                                       f"please download relative pretrain model in huggingface or" \
                                       f"https://github.com/lonePatient/NeZha_Chinese_PyTorch " \
                                       f"model name:nezha-cn-base or nezha-base-wwm"

    if config['pretrain_type'] == 'dynamic_mask':
        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer,
                                                        mlm=True,
                                                        mlm_probability=mlm_probability)
        model = NeZhaForMaskedLM.from_pretrained(pretrained_model_name_or_path=model_path,
                                                 config=model_config)
        model_save_path = 'mlm_model'

    if config['pretrain_type'] == 'whole_word_mask':
        data_collator = DataCollatorForWholeWordMask(tokenizer=tokenizer,
                                                     mlm=True,
                                                     mlm_probability=mlm_probability)
        model = NeZhaForMaskedLM.from_pretrained(pretrained_model_name_or_path=model_path,
                                                 config=model_config)
        model_save_path = 'whole_word_mask_model'

    dataset = LineByLineTextDataset(tokenizer=tokenizer,
                                    train_file_path=config['train_data_path'],
                                    test_file_path=config['test_data_path'],
                                    block_size=seq_length)

    training_args = TrainingArguments(
        output_dir='record',
        num_train_epochs=num_train_epochs,
        learning_rate=learning_rate,
        per_device_train_batch_size=batch_size,
        save_steps=save_steps,
        logging_steps=500,
        save_total_limit=5,
        prediction_loss_only=True,
        seed=seed
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=dataset
    )

    trainer.train()
    trainer.save_model(model_save_path)
    tokenizer.save_pretrained(model_save_path)

if __name__ == '__main__':
    main()