import jsonfrom tqdm import tqdmfrom transformers import BertForTokenClassificationfrom transformers import BertTokenizerFast
import torchfrom torch import nn
import numpy as np

BERT_NAME = 'bert-base-chinese'

label2id = { 'O': 0, 'B-LOC': 1, 'I-LOC': 2, 'B-ORG': 3, 'I-ORG': 4, 'B-PER': 5, 'I-PER': 6}
tokenizer = BertTokenizerFast.from_pretrained(BERT_NAME)max_seq_len = 256
pad_id = tokenizer.convert_tokens_to_ids('[PAD]')cls_id = tokenizer.convert_tokens_to_ids('[CLS]')sep_id = tokenizer.convert_tokens_to_ids('[SEP]')unk_id = tokenizer.convert_tokens_to_ids('[UNK]')unword_label_id = len(label2id)
def trans2TheDataset(data, max_seq_len):    data2 = []
    for k, v  in data:        k = [tokenizer.convert_tokens_to_ids(i) for i in k]        v = [label2id[i] for i in v]                        input_ids = [cls_id] + k        labels = [unword_label_id] + v                if len(input_ids) > max_seq_len-1:            input_ids = input_ids[:max_seq_len-1]            labels = labels[:max_seq_len-1]                input_ids.append(sep_id)        labels.append(unword_label_id)                klen = len(input_ids)                if klen < max_seq_len:            pad_len = max_seq_len-klen                    input_ids = input_ids + [pad_id] * pad_len            labels = labels + [unword_label_id] * pad_len                        token_type_ids = [0] * max_seq_len            attention_mask = [1] * klen + [0] * pad_len                data2.append((np.asarray(input_ids, dtype='int64'),                       np.asarray(token_type_ids, dtype='int64'),                      np.asarray(attention_mask, dtype='int64'),                      np.asarray(labels, dtype='int64')))    return TheDataset(data2)#train_data拆分为train和dev，其中dev占20%，train占80%train_data_count = len(train_data)dev_data_count = int(train_data_count * 0.2)train_data_count -= dev_data_countprint(train_data_count, dev_data_count)dev_data = train_data[:dev_data_count]train_data = train_data[dev_data_count:]print(len(train_data), len(dev_data))
train_dataset = trans2TheDataset(train_data, max_seq_len)dev_dataset = trans2TheDataset(dev_data, max_seq_len)test_dataset = trans2TheDataset(test_data, max_seq_len)

batch_size = 16num_workers = 1train_dataloader = torch.utils.data.DataLoader(train_dataset, num_workers=num_workers, batch_size=batch_size, shuffle=True)dev_dataloader = torch.utils.data.DataLoader(dev_dataset, num_workers=num_workers, batch_size=batch_size, shuffle=False)test_dataloader = torch.utils.data.DataLoader(test_dataset, num_workers=num_workers, batch_size=batch_size, shuffle=False)

class BertModel(nn.Module):
    def __init__(self, num_labels):        super(BertModel, self).__init__()        self.bert = BertForTokenClassification.from_pretrained(BERT_NAME, num_labels=num_labels)
    def forward(self, input_ids, attention_mask,token_type_ids, labels):        output = self.bert(input_ids=input_ids, attention_mask=attention_mask,                           token_type_ids=token_type_ids, labels=labels,                           output_hidden_states=False,                            output_attentions=False,                           return_dict=False)        return output

model = BertModel(len(label2id)+1)

model = model.cuda()

decay_layer = [n for n, p in model.named_parameters() if 'layer.11' in n ]optimizer_grouped_parameters = [    {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in decay_layer)], 'weight_decay': 0.01},    {'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in decay_layer)], 'weight_decay': 0.0}]

lr = 0.00005optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=lr)

for ep in range(10):
    total_acc_train = 0    total_loss_train = 0    train_count = 0
    model.train()
    for input_ids, token_type_ids, attention_mask, labels  in tqdm(train_dataloader):
        labels = labels.to(device)        attention_mask = attention_mask.to(device)        input_ids = input_ids.to(device)        token_type_ids = token_type_ids.to(device)
        optimizer.zero_grad()        loss, logits = model(input_ids, attention_mask,token_type_ids, labels)                for i in range(logits.shape[0]):            logits_clean = logits[i][labels[i] != unword_label_id]
            preds = logits_clean.argmax(dim=1)
            label_clean = labels[i][labels[i] != unword_label_id]
            acc = (preds == label_clean).float().mean()                    total_acc_train += acc.item()            total_loss_train += loss.item()            train_count += 1        loss.backward()        optimizer.step()        print(f'Ep {ep} Train Acc: {total_acc_train/train_count:.3f} Loss: {total_loss_train/train_count:.3f}')            model.eval()        total_acc_val = 0    total_loss_val = 0    dev_count = 0
    for input_ids, token_type_ids, attention_mask, labels in tqdm(dev_dataloader):
        labels = labels.to(device)        attention_mask = attention_mask.to(device)        input_ids = input_ids.to(device)        token_type_ids = token_type_ids.to(device)                loss, logits = model(input_ids, attention_mask,token_type_ids, labels)
        for i in range(logits.shape[0]):
            logits_clean = logits[i][labels[i] != unword_label_id]
            preds = logits_clean.argmax(dim=1)
            label_clean = labels[i][labels[i] != unword_label_id]
            acc = (preds == label_clean).float().mean()                    total_acc_val += acc.item()            total_loss_val += loss.item()            dev_count += 1                print(f'Ep {ep} Eval Acc: {total_acc_val/dev_count:.3f} Loss: {total_loss_val/dev_count:.3f}')

model.eval()  total_acc = 0total_loss = 0test_count = 0
all_inputs = []all_logits = []all_labels = []all_masks = []
with torch.no_grad():    for input_ids, token_type_ids, attention_mask, labels in tqdm(test_dataloader):
        all_inputs.append(input_ids)        all_labels.append(labels)        all_masks.append(attention_mask)
        labels = labels.to(device)        attention_mask = attention_mask.to(device)        input_ids = input_ids.to(device)        token_type_ids = token_type_ids.to(device)
        loss, logits = model(input_ids, attention_mask,token_type_ids, labels)
        #all_inputs.append(input_ids.to('cpu'))        logits_cpu = logits.to('cpu').detach().numpy()        all_logits.append(logits_cpu)        #all_labels.append(labels.to('cpu'))
        for i in range(logits.shape[0]):
            logits_clean = logits[i][labels[i] != unword_label_id]
            preds = logits_clean.argmax(dim=1)
            label_clean = labels[i][labels[i] != unword_label_id]
            acc = (preds == label_clean).float().mean()
            total_acc += acc.item()            total_loss += loss.item()            test_count += 1
print(f'Test Acc: {total_acc/test_count:.3f} Loss: {total_loss/test_count:.3f}')