import torch import glob import os from transformers import BertTokenizerFast as BertTokenizer, BertForSequenceClassification LABEL_COLUMNS = ["Assertive Tone", "Conversational Tone", "Emotional Tone", "Informative Tone", "None"] tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") model = BertForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=5) id2label = {i:label for i,label in enumerate(LABEL_COLUMNS)} label2id = {label:i for i,label in enumerate(LABEL_COLUMNS)} for ckpt in glob.glob('checkpoints/*.ckpt'): base_name = os.path.basename(ckpt) # 去除文件后缀 model_name = os.path.splitext(base_name)[0] params = torch.load(ckpt, map_location="cpu")['state_dict'] msg = model.load_state_dict(params, strict=True) path = f'models/{model_name}' os.makedirs(path, exist_ok=True) torch.save(model.state_dict(), f'{path}/pytorch_model.bin') config = model.config config.architectures = ['BertForSequenceClassification'] config.label2id = label2id config.id2label = id2label model.config.to_json_file(f'{path}/config.json') tokenizer.save_vocabulary(path)