帮我写一段bert-bilstm-crf-ner模型用于中文命名实体识别的代码

以下是一个简单的bert-bilstm-crf-ner模型用于中文命名实体识别的代码：

import torch
import torch.nn as nn
from transformers import BertModel

class BertBilstmCrfNer(nn.Module):
    def __init__(self, bert_path, num_tags):
        super(BertBilstmCrfNer, self).__init__()
        self.bert = BertModel.from_pretrained(bert_path)
        self.dropout = nn.Dropout(.5)
        self.bilstm = nn.LSTM(input_size=768, hidden_size=256, num_layers=2, batch_first=True, bidirectional=True)
        self.fc = nn.Linear(512, num_tags)
        self.crf = nn.CRF(num_tags)

    def forward(self, input_ids, attention_mask):
        bert_output = self.bert(input_ids=input_ids, attention_mask=attention_mask)[]
        bert_output = self.dropout(bert_output)
        lstm_output, _ = self.bilstm(bert_output)
        lstm_output = self.dropout(lstm_output)
        logits = self.fc(lstm_output)
        mask = attention_mask.bool()
        tags = self.crf.decode(logits, mask=mask)
        return tags

这个模型使用了BERT作为输入特征提取器，然后通过一个双向LSTM进行特征提取，最后使用一个全连接层将特征映射到标签空间，并使用CRF进行标签解码。