python - HuggingFace 保存加载模型 (Colab) 进行预测

Question

使用 HuggingFace 训练 Transformer 模型以预测目标变量(例如，电影评级)。我是 Python 的新手，这可能是一个简单的问题，但我不知道如何保存经过训练的分类器模型(通过 Colab)然后重新加载以便对新数据进行目标变量预测。例如，我使用 HuggingFace 资源中的示例训练了一个模型来预测 imbd 评级，如下所示。我尝试了多种方法(save_model、save_pretrained)，要么根本无法保存它，要么在加载时不知道调用什么来获得预测。对于涉及保存、加载然后根据测试数据模型创建新预测分数的步骤，我们将不胜感激。

#example mainly from here: https://huggingface.co/transformers/training.html
!pip install transformers
!pip install datasets

from datasets import load_dataset
raw_datasets = load_dataset("imdb")

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

def tokenize_function(examples):
    return tokenizer(examples["text"], max_length = 128, padding="max_length", truncation=True) 

tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)

#choosing small datasets for example#
small_train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(1000))
small_eval_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(500))

### TRAINING classification ###
from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased", num_labels=2)

from transformers import TrainingArguments
from transformers import Trainer

training_args = TrainingArguments("test_trainer", evaluation_strategy="epoch", num_train_epochs=2, weight_decay=.0001, learning_rate=0.00001, per_device_train_batch_size=32) 

trainer = Trainer(model=model, args=training_args, train_dataset=small_train_dataset, eval_dataset=small_eval_dataset)
trainer.train()

y_test_predicted_original = model_loaded.predict(small_eval_dataset)

#### Saving ###
from google.colab import drive
drive.mount('/content/gdrive')
%cd /content/gdrive/My\ Drive/FOLDER

trainer.save_pretrained ("Trained model") #assumed this would save but did not
model.save_pretrained ("Trained model") #did save

### Loading Model and Creating Predicted Scores ###

#perhaps this....#
from transformers import BertConfig, BertModel
conf = BertConfig.from_pretrained("Trained model", num_labels=2)
model_loaded = AutoModelForSequenceClassification.from_pretrained("Trained model", config=conf)

#or...#
model_loaded = AutoModelForSequenceClassification.from_pretrained("Trained model", local_files_only=True)
model_loaded 

#with ultimate goal of getting predicted scores (not sure what to call here)...
y_test_predicted_loaded = model_loaded.predict(small_eval_dataset)

Answer 1

保存模型

trainer.save_model("Trained model")

加载模型

model_loaded = AutoModelForSequenceClassification.from_pretrained("Trained model")

预测

trainer = Trainer(model = model_loaded)
test_results = trainer.predict(test_dataset)