sequence_BERTCRF1.py

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import matplotlib.pyplot as plt
import numpy as np
import torch
from TorchCRF import CRF
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from sklearn.metrics import precision_score, recall_score, f1_score
from torch.utils.data import DataLoader,Dataset
import collections
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 只使用 GPU 0



from transformers import BertModel,BertTokenizer,BertConfig
import argparse



class ChemDataset(Dataset):
    def __init__(self, data_dir, tokenizer, max_seq_length, mode):
        self.examples = self._load_examples(data_dir, mode)
        self.label_map = self._get_label_map(data_dir)
        self.tokenizer = tokenizer
        self.max_seq_length = max_seq_length
        self.mode = mode
    def _load_examples(self, data_dir, mode):
        with open(os.path.join(data_dir, mode, "seq.in"), 'r', encoding='utf-8') as seq_in_f:
            with open(os.path.join(data_dir, mode, "seq.out"), 'r', encoding='utf-8') as seq_out_f:
                text_list = [line.strip() for line in seq_in_f.readlines()]
                label_list = [line.strip() for line in seq_out_f.readlines()]

        assert len(text_list) == len(label_list)
        return list(zip(text_list, label_list))

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

    def get_labels(self,data_dir):
        label_set = set()
        for f_type in ["train", "test"]:
            seq_out_dir = os.path.join(os.path.join(data_dir, f_type), "seq.out")
            with open(seq_out_dir) as data_f:
                seq_sentence_list = [seq.split() for seq in data_f.readlines()]
                seq_word_list = [word for seq in seq_sentence_list for word in seq]
                label_set = label_set | set(seq_word_list)
        label_list = list(label_set)
        label_list.sort()
        return ["[Padding]", "[CLS]", "[SEP]"] + label_list

    def _get_label_map(self,data_dir):
        label_set = set()
        for f_type in ["train", "test"]:
            seq_out_dir = os.path.join(os.path.join(data_dir, f_type), "seq.out")
            with open(seq_out_dir) as data_f:
                seq_sentence_list = [seq.split() for seq in data_f.readlines()]
                seq_word_list = [word for seq in seq_sentence_list for word in seq]
                label_set = label_set | set(seq_word_list)
        label_list = list(label_set)
        label_list.sort()
        label_list = ["[Padding]", "[CLS]", "[SEP]"] + label_list
        return {label: i for i, label in enumerate(label_list)}

    def __getitem__(self, idx):
        text, label = self.examples[idx]  # 直接解包 tuple
        # textlist = text.split()  # 分词
        textlist = text
        labellist = label.split()  # 标签分词
        label_map={}
        tokens = []
        labels = []
        label_ids = []
        for i, word in enumerate(labellist):
            tokenized_word = self.tokenizer.tokenize(textlist[i])
            tokens.extend(tokenized_word)
            label = labellist[i]
            for j in range(len(tokenized_word)):
                if j == 0:
                    labels.append(label)
                else:
                    labels.append("[##WordPiece]")

        # 截断
        if len(tokens) >= self.max_seq_length - 2:
            tokens = tokens[:self.max_seq_length - 2]
            labels = labels[:self.max_seq_length - 2]

        # 添加特殊 token
        ntokens = ["[CLS]"] + tokens + ["[SEP]"]
        segment_ids = [0] * len(ntokens)
        label_map = self._get_label_map(data_dir=data_dir)
        # print(label_map)
        label_ids.append(label_map["[CLS]"])
        for i, token in enumerate(tokens):
            label_ids.append(label_map[labels[i]])
        label_ids.append(label_map["[SEP]"])
        # label_ids = [self.label_map.get(label, 0) for label in ntokens]

        input_ids = self.tokenizer.convert_tokens_to_ids(ntokens)
        input_mask = [1] * len(input_ids)

        # 填充
        while len(input_ids) < self.max_seq_length:
            input_ids.append(0)
            input_mask.append(0)
            segment_ids.append(0)
            label_ids.append(0)

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(input_mask, dtype=torch.long),
            "segment_ids": torch.tensor(segment_ids, dtype=torch.long),
            "label_ids": torch.tensor(label_ids, dtype=torch.long),
        }
def compute_confidence_structure_loss(pred_logits, labels, label_map, mask, conf_threshold=0.8):
    """
    基于预测置信度的结构一致性loss（仅在预测自信时才进行结构惩罚）
    """
    batch_size, seq_len, num_labels = pred_logits.size()
    penalty = 0.0
    count = 0

    pred_probs = torch.softmax(pred_logits, dim=-1)
    pred_labels = pred_probs.argmax(dim=-1)

    for b in range(batch_size):
        for t in range(1, seq_len):
            if mask[b, t] == 0:
                continue

            curr_id = pred_labels[b, t].item()
            prev_id = pred_labels[b, t - 1].item()
            curr_label = label_map[curr_id]
            prev_label = label_map[prev_id]

            if curr_label.startswith("I-"):
                curr_type = curr_label[2:]
                if not (prev_label.startswith("B-") or prev_label.startswith("I-")) or prev_label[2:] != curr_type:
                    # 判断是否满足置信度门槛
                    if pred_probs[b, t, curr_id] >= conf_threshold:
                        penalty += pred_probs[b, t, curr_id]
                        count += 1

    if count > 0:
        structure_loss = penalty / count
    else:
        structure_loss = torch.tensor(0.0, device=pred_logits.device)
    return structure_loss

def compute_structure_loss(pred_logits, labels, label_map, mask):
    """
    结构一致性loss，用于惩罚 BIO标签结构不合法的情况。
    :param pred_logits: 模型输出的 logits, shape: (batch_size, seq_len, num_labels)
    :param labels: gold label ids, shape: (batch_size, seq_len)
    :param label_map: id → label 名称映射，如 {0:'O', 1:'B-LOC', 2:'I-LOC', ...}
    :param mask: attention_mask, shape: (batch_size, seq_len), 0表示padding位置
    :return: 结构惩罚loss（可加到主loss上）
    """
    batch_size, seq_len, num_labels = pred_logits.size()
    penalty = 0.0
    count = 0

    pred_probs = torch.softmax(pred_logits, dim=-1)  # 概率分布
    pred_labels = pred_probs.argmax(dim=-1)  # shape: (batch_size, seq_len)

    for b in range(batch_size):
        for t in range(1, seq_len):
            if mask[b, t] == 0:
                continue  # 跳过padding

            curr_id = pred_labels[b, t].item()
            prev_id = pred_labels[b, t - 1].item()
            curr_label = label_map[curr_id]
            prev_label = label_map[prev_id]

            if curr_label.startswith("I-"):
                curr_type = curr_label[2:]
                if not (prev_label.startswith("B-") or prev_label.startswith("I-")) or (prev_label[2:] != curr_type):
                    # 结构不一致，增加惩罚（惩罚当前token在softmax分布中对I-类的偏好）
                    # 拉低I-X在这个位置的概率，提升O/B类概率
                    # 强化loss: log(P_bad_label) → -log(1-P)
                    penalty += pred_probs[b, t, curr_id]  # 越自信错 → 惩罚越大
                    count += 1

    if count > 0:
        structure_loss = penalty / count
    else:
        structure_loss = torch.tensor(0.0, device=pred_logits.device)

    return structure_loss

class BertForSequenceLabeling(nn.Module):
    def __init__(self, pretrained_model, num_labels, config, labels):
        super(BertForSequenceLabeling, self).__init__()
        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.bert = BertModel.from_pretrained(pretrained_model, config=config, ignore_mismatched_sizes=True)  # 加载 BERT
        self.dropout = nn.Dropout(p=0.3)  # 训练时的 dropout
        self.hidden_size = self.bert.config.hidden_size  # BERT 的隐藏层大小
        self.classifier = nn.Linear(self.hidden_size, num_labels)  # 分类器
        self.crf = CRF(num_tags=num_labels, batch_first=True )
        self.labels = labels

    def initialize_crf_constraints(self):
        for i, label in enumerate(self.labels):
            if label.startswith("I-"):
                # 禁止 O -> I
                self.crf.transitions.data[self.labels.index("O"), i] = -10000.0

                # 禁止不同类型 I-X 连续
                type_i = label[2:]
                for j, label_prev in enumerate(self.labels):
                    if label_prev.startswith("I-") or label_prev.startswith("B-"):
                        type_j = label_prev[2:]
                        if type_i != type_j:
                            self.crf.transitions.data[j, i] = -10000.0

    def forward(self, input_ids, input_mask,segment_ids, labels):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=input_mask,
            token_type_ids=segment_ids
        )
        sequence_output = outputs.last_hidden_state  # 获取序列输出 (batch_size, seq_len, hidden_size)
        hidden_size = sequence_output.shape[-1]
        sequence_output = self.dropout(sequence_output)  # Dropout
        # sequence_output = sequence_output.view[-1, hidden_size]
        emissions = self.classifier(sequence_output)  # (batch_size, seq_len, num_labels)


        if labels is not None:
            # CRF loss: 使用 mask 指定有效位置
            loss = -self.crf(emissions, labels, mask=input_mask.bool())

            structure_loss = compute_confidence_structure_loss(
                pred_logits=emissions,
                labels=labels,
                label_map={i: label for i, label in enumerate(self.labels)},
                mask=input_mask
            )

            λ = 0.1
            total_loss = loss + λ * structure_loss  #   不用total_loss，效果较差

            predictions = self.crf.decode(emissions, mask=input_mask.bool())

            return loss, predictions
        else:
            # 推理时解码
            predictions = self.crf.decode(emissions, mask=input_mask.bool())
            return None, predictions


def evaluate(model, test_loader, train_dataset, data_dir):
    print("******Evaluating******")
    model.eval()
    total_loss = 0
    all_preds = []
    all_labels = []


    with torch.no_grad():
        for batch in test_loader:
            input_ids = batch["input_ids"].cuda()
            segment_ids = batch["segment_ids"].cuda()
            attention_mask = batch["attention_mask"].cuda()
            labels = batch["label_ids"].cuda()

            loss, predictions = model(input_ids, attention_mask, segment_ids, labels)
            loss = loss.mean()
            total_loss += loss.item()

            for i in range(len(predictions)):
                pred_seq = predictions[i]
                true_seq = labels[i].cpu().tolist()
                mask_seq = attention_mask[i].cpu().tolist()

                true_seq = [t for t, m in zip(true_seq, mask_seq) if m == 1]
                all_preds.extend(pred_seq)
                all_labels.extend(true_seq)

    all_preds = torch.tensor(all_preds)
    all_labels = torch.tensor(all_labels)

    metrics = compute_metrics(all_preds, all_labels, ignore_index=0)
    avg_loss = total_loss / len(test_loader)
    print(f"Precision F1 Score(Macro): {metrics['f1_macro']:.4f}")
    return avg_loss, metrics['f1_macro']

def compute_metrics(preds, labels, ignore_index=0):
    """
    Compute precision, recall, and F1-score (micro/macro) using PyTorch & Sklearn.

    Args:
        preds (torch.Tensor): Model predictions, shape (batch_size, seq_len).
        labels (torch.Tensor): Ground truth labels, shape (batch_size, seq_len).
        num_labels (int): Number of unique labels (including O and Padding).
        ignore_index (int): Label index to ignore in calculation (e.g., `[PAD]`).

    Returns:
        dict: Precision, Recall, F1-score (macro & micro).
    """

    # Flatten the predictions and labels
    preds = preds.view(-1).cpu().numpy()
    labels = labels.view(-1).cpu().numpy()

    # Mask out ignored tokens (e.g., [PAD])
    mask = labels != ignore_index
    preds = preds[mask]
    labels = labels[mask]

    # Compute precision, recall, and F1-score
    precision_macro = precision_score(labels, preds, average="macro", zero_division=0)
    recall_macro = recall_score(labels, preds, average="macro", zero_division=0)
    f1_macro = f1_score(labels, preds, average="macro", zero_division=0)

    precision_micro = precision_score(labels, preds, average="micro", zero_division=0)
    recall_micro = recall_score(labels, preds, average="micro", zero_division=0)
    f1_micro = f1_score(labels, preds, average="micro", zero_division=0)

    return {
        "precision_macro": precision_macro,
        "recall_macro": recall_macro,
        "f1_macro": f1_macro,
        "precision_micro": precision_micro,
        "recall_micro": recall_micro,
        "f1_micro": f1_micro,
    }

if __name__=="__main__":
    config = BertConfig.from_pretrained("./pretrained_model/bert_base_chinese")
    config.max_position_embeddings = 512
    data_name = "ChemNER+"
    data_dir = f"data/{data_name}"

    tokenizer = BertTokenizer.from_pretrained("./pretrained_model/bert_base_chinese",do_lower_case=True)

    train_dataset = ChemDataset(data_dir=data_dir, tokenizer=tokenizer, max_seq_length=128, mode="train")
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

    test_dataset = ChemDataset(data_dir=data_dir, tokenizer=tokenizer, max_seq_length=128, mode="test")
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

    labels = train_dataset.get_labels(data_dir=data_dir)

    pretrained_model = "./pretrained_model/bert_base_chinese"
    # 初始化模型
    num_labels = len(labels)
    model = BertForSequenceLabeling(pretrained_model, num_labels, config, labels=labels).cuda()
    # optimizer = optim.AdamW(model.parameters(), lr=2e-5)
    optimizer = optim.AdamW([
        {'params': model.bert.parameters(), 'lr': 2e-5},  # 预训练层，较小lr
        {'params': model.classifier.parameters(), 'lr': 2e-5},  # 分类层，较大lr
        {'params': model.crf.parameters(), 'lr': 1e-3, 'weight_decay': 0.05},  # CRF transition层，更大lr
    ])
    model.initialize_crf_constraints()  # 初始化转移矩阵

    # args.do_train=False

    train_loss = []
    eval_loss = []
    eval_f1 = []
    num_epochs = 30
    do_train = False
    do_eval = False
    # 训练循环
    if do_train==True:
        print("******Training******")

        for epoch in range(num_epochs):
            model.train()
            total_loss = 0
            for batch in train_loader:
                # print("Batch keys:", batch.keys())
                input_ids = batch["input_ids"].cuda()
                segment_ids = batch["segment_ids"].cuda()
                attention_mask = batch["attention_mask"].cuda()
                labels = batch["label_ids"].cuda()
                # print("Labels min:", labels,labels.min().item(), "max:", labels.max().item())
                optimizer.zero_grad()
                loss, predictions = model(input_ids, attention_mask, segment_ids,labels)
                loss = loss.mean()
                loss.backward()
                optimizer.step()
                print(f"Batch_loss    Loss: {loss}")
                total_loss += loss.item()
            epoch_loss = total_loss / len(train_loader)
            print(f"Epoch {epoch+1}, Loss: {epoch_loss}")
            epoch_eval_loss, epoch_eval_f1 = evaluate(model, test_loader, train_dataset, data_dir)
            train_loss.append(epoch_loss)
            eval_loss.append(epoch_eval_loss)
            eval_f1.append(epoch_eval_f1)

        output_dir = os.path.join(os.getcwd(), "output_model", f"{data_name}_bertcrf_epoch{num_epochs}")

        # 确保目录存在
        os.makedirs(output_dir, exist_ok=True)
        output_model_file = os.path.join(os.getcwd(), "output_model", f"{data_name}_bertcrf_epoch{num_epochs}", "bert_ner.pth")
        torch.save(model.state_dict(), output_model_file)
        print(f"Model saved to {output_model_file}!")

    # args.do_eval=True

    if do_eval==True:
        print("******Evaluating******")
        model_path = f"./output_model/{data_name}_bertcrf_epoch{num_epochs}/bert_ner.pth"

        model.load_state_dict(torch.load(model_path))
        model.eval()
        print(f"Load model {model_path} successfully!")
        total_loss = 0
        correct = 0
        total = 0
        all_preds = []
        all_labels = []

        label_list = train_dataset.get_labels(data_dir=data_dir)
        id2label = {i: label for i, label in enumerate(label_list)}
        all_pred_labels = []

        with torch.no_grad():
            for batch in test_loader:
                input_ids = batch["input_ids"].cuda()
                segment_ids = batch["segment_ids"].cuda()
                attention_mask = batch["attention_mask"].cuda()
                labels = batch["label_ids"].cuda()

                loss, predictions = model(input_ids, attention_mask, segment_ids, labels)
                loss = loss.mean()
                total_loss += loss.item()

                # 注意 predictions 是 list of list
                for i in range(len(predictions)):
                    pred_seq = predictions[i]  # 长度 = 有效token数
                    true_seq = labels[i].cpu().tolist()
                    mask_seq = attention_mask[i].cpu().tolist()

                    true_seq = [t for t, m in zip(true_seq, mask_seq) if m == 1]
                    all_preds.extend(pred_seq)
                    all_labels.extend(true_seq)

                    # 保存标签名
                    pred_label_names = [id2label[int(p)] for p in pred_seq]
                    all_pred_labels.append(pred_label_names)
        all_preds=np.array(all_preds)
        all_labels=np.array(all_labels)
        # print("all_preds shape:", all_preds.shape)
        # print("all_labels shape:", all_labels.shape)
        all_preds = torch.tensor(all_preds)
        all_labels = torch.tensor(all_labels)
        # 合并所有 batch


        # 计算评估指标
        metrics = compute_metrics(all_preds, all_labels, ignore_index=0)

        avg_loss = total_loss / len(test_loader)
        # accuracy = correct / total
        print(f"Precision Loss: {avg_loss:.4f}")
        print(f"Precision (Macro): {metrics['precision_macro']:.4f}, Recall (Macro): {metrics['recall_macro']:.4f}, F1 (Macro): {metrics['f1_macro']:.4f}")
        print(f"Precision (Micro): {metrics['precision_micro']:.4f}, Recall (Micro): {metrics['recall_micro']:.4f}, F1 (Micro): {metrics['f1_micro']:.4f}")


        output_file = os.path.join(os.getcwd(), "output_model", f"{data_name}_bertcrf_epoch{num_epochs}", "test_predictions.txt")
        with open(output_file, "w", encoding="utf-8") as f_out:
            for pred_seq in all_pred_labels:
                # 将一个样本的所有预测标签拼接为用空格分隔的一行
                line_str = " ".join(pred_seq)
                f_out.write(line_str + "\n")

        print(f"Predictions have been saved to {output_file}!")

    plt.figure()
    plt.plot(range(1, num_epochs + 1), train_loss, label='Train Loss', marker='x')

    # 绘制测试损失（注意只在某些 epoch 上有评估）
    plt.plot(range(1, num_epochs + 1), eval_loss, label='Eval Loss', marker='o')

    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title('Training & Evaluation Loss')
    plt.legend()
    plt.grid(True)
    plt.savefig('result/loss_curve_3.png')  # 保存图像
    plt.show()

    plt.figure()
    plt.plot(range(1, num_epochs + 1), eval_f1, label='F1 Macro')
    plt.xlabel('epoch')
    plt.ylabel('F1 Score')
    plt.title('F1 Score on test')
    plt.legend()
    plt.grid(True)
    plt.savefig('result/f1_curve_3.png')
    plt.show()