| """独立加载本目录中的商品标题分类模型，无需原项目和基座权重。""" |
|
|
| import json |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
| from transformers import AutoConfig, AutoModel |
|
|
|
|
| class BertTitleClassifier(nn.Module): |
| def __init__(self, bert_config, num_classes): |
| super().__init__() |
| self.bert = AutoModel.from_config(bert_config) |
| self.classifier = nn.Linear(bert_config.hidden_size, num_classes) |
|
|
| def forward(self, input_ids, attention_mask=None): |
| outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask) |
| return self.classifier(outputs.last_hidden_state[:, 0, :]) |
|
|
| @classmethod |
| def from_pretrained(cls, model_dir, device="cpu"): |
| model_dir = Path(model_dir) |
| labels = json.loads((model_dir / "labels.json").read_text(encoding="utf-8")) |
| bert_config = AutoConfig.from_pretrained(model_dir, local_files_only=True) |
| model = cls(bert_config, len(labels["id2label"])) |
| state = torch.load(model_dir / "model.pt", map_location="cpu", weights_only=True) |
| model.load_state_dict(state) |
| return model.to(device).eval() |
|
|
