File size: 2,586 Bytes
b0132ae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """ERINE + CRF 模型"""
import torch
import torch.nn as nn
from torchcrf import CRF
from transformers import AutoModel, AutoTokenizer
from config import ERNIE_LOCAL, BIO_LABELS, CHECKPOINT, CHECKPOINT_FROZEN, CHECKPOINT_FC2
class ErnieCRF(nn.Module):
"""ERNIE 3.0 encoder + Linear + CRF"""
def __init__(self, model_path, num_labels):
super().__init__()
self.ernie = AutoModel.from_pretrained(model_path)
self.fc = nn.Linear(self.ernie.config.hidden_size, num_labels)
self.crf = CRF(num_labels, batch_first=True)
def forward(self, input_ids, attention_mask, labels=None):
mask = attention_mask.bool()
hidden = self.ernie(input_ids, attention_mask=attention_mask).last_hidden_state
emissions = self.fc(hidden)
if labels is not None:
return -self.crf(emissions, labels, mask=mask, reduction="mean")
return self.crf.decode(emissions, mask=mask)
class ErnieCRF2(nn.Module):
"""ERNIE 3.0 encoder + 双层FC(hidden→3*hidden→3) + CRF"""
def __init__(self, model_path, num_labels, hidden_factor=3):
super().__init__()
self.ernie = AutoModel.from_pretrained(model_path)
self.hidden_size = self.ernie.config.hidden_size
mid_size = self.hidden_size * hidden_factor
self.fc1 = nn.Linear(self.hidden_size, mid_size)
self.fc2 = nn.Linear(mid_size, num_labels)
self.relu = nn.ReLU()
self.dropout = nn.Dropout(0.1)
self.crf = CRF(num_labels, batch_first=True)
def forward(self, input_ids, attention_mask, labels=None):
mask = attention_mask.bool()
hidden = self.ernie(input_ids, attention_mask=attention_mask).last_hidden_state
x = self.dropout(self.relu(self.fc1(hidden)))
emissions = self.fc2(x)
if labels is not None:
return -self.crf(emissions, labels, mask=mask, reduction="mean")
return self.crf.decode(emissions, mask=mask)
def load_model(device="cuda", frozen=False, fc2=False):
"""加载训练好的模型和 tokenizer"""
tokenizer = AutoTokenizer.from_pretrained(ERNIE_LOCAL)
if fc2:
model = ErnieCRF2(ERNIE_LOCAL, len(BIO_LABELS), hidden_factor=2).to(device)
ckpt = CHECKPOINT_FC2
else:
model = ErnieCRF(ERNIE_LOCAL, len(BIO_LABELS)).to(device)
ckpt = CHECKPOINT_FROZEN if frozen else CHECKPOINT
model.load_state_dict(torch.load(ckpt, map_location=device, weights_only=True))
model.eval()
return model, tokenizer
|