写一个使用bert-base-chinese 进行实体抽取的代码示例

以下是一个使用bert-base-chinese进行实体抽取的代码示例：

import torch
import transformers

tokenizer = transformers.BertTokenizer.from_pretrained('bert-base-chinese')
model = transformers.BertForTokenClassification.from_pretrained('bert-base-chinese', num_labels=5)

text = '我在北京的清华大学上学。'
encoded_input = tokenizer(text, return_tensors='pt')
output = model(**encoded_input)[0]
predicted_labels = torch.argmax(output, dim=2)[0]

entities = []
current_entity = ''
current_entity_type = ''
for i, token in enumerate(tokenizer.tokenize(text)):
    if token.startswith('##'):
        current_entity += token[2:]
    else:
        if current_entity:
            entities.append((current_entity, current_entity_type))
            current_entity = ''
            current_entity_type = ''
        current_entity += token
        current_entity_type = predicted_labels[i].item()

if current_entity:
    entities.append((current_entity, current_entity_type))

print(entities)

在这个示例中，我们首先加载了预训练的BertTokenizer和BertForTokenClassification模型。然后，我们使用tokenizer对输入文本进行编码，并将其传递给模型进行预测。最后，我们通过遍历预测的标签来识别实体，并将它们存储在一个列表中。在这个示例中，我们假设模型可以预测5种实体类型。在实际应用中，您可能需要根据您的数据集和任务进行自定义。