| """ | |
| 模型定义：基于 BertForSequenceClassification 的二分类模型 | |
| 使用 HuggingFace 原生的 save_pretrained / from_pretrained 实现可靠保存/加载 | |
| """ | |
| from transformers import BertForSequenceClassification | |
| def create_model(model_name: str = "bert-base-chinese", num_labels: int = 2): | |
| """创建分类模型""" | |
| model = BertForSequenceClassification.from_pretrained( | |
| model_name, | |
| num_labels=num_labels, | |
| ) | |
| return model | |
| def load_model(model_path: str): | |
| """加载已训练的模型""" | |
| model = BertForSequenceClassification.from_pretrained(model_path) | |
| return model | |