拥抱面节约负荷模型（Colab）进行预测

Question

使用HuggingFace训练变压器模型以预测目标变量（例如，电影收视率）。我是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)