python人工智能NLP BERT 命名实体识别 : 用BERT实现一个 NLP 命名实体识别 任务，代码方案分享1(图文详解)

import torch
from transformers import BertTokenizer, BertModel

# 加载预训练的 BERT 模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 定义 NER 标签
labels = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]

# 将文本转换为 BERT 输入
def convert_to_bert_input(text):
    tokens = tokenizer.tokenize(text)
    input_ids = tokenizer.convert_tokens_to_ids(tokens)
    attention_mask = [1] * len(input_ids)
    return torch.tensor([input_ids]), torch.tensor([attention_mask])

# 训练模型
def train_model(train_data, epochs):
    optimizer = torch.optim.Adam(model.parameters())
    for epoch in range(epochs):
        for batch in train_data:
            input_ids, attention_mask, labels = batch
            outputs = model(input_ids, attention_mask=attention_mask)
            loss = torch.nn.CrossEntropyLoss()(outputs.logits, labels)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

# 评估模型
def evaluate_model(test_data):
    model.eval()
    with torch.no_grad():
        for batch in test_data:
            input_ids, attention_mask, labels = batch
            outputs = model(input_ids, attention_mask=attention_mask)
            predictions = torch.argmax(outputs.logits, dim=-1)
            accuracy = (predictions == labels).sum() / len(labels)
            print(f"Accuracy: {accuracy}")

# 预测实体
def predict_entities(text):
    input_ids, attention_mask = convert_to_bert_input(text)
    outputs = model(input_ids, attention_mask=attention_mask)
    predictions = torch.argmax(outputs.logits, dim=-1)
    entities = []
    for i, prediction in enumerate(predictions):
        if prediction != 0:
            entity = (i, labels[prediction])
            entities.append(entity)
    return entities