Search code examples
pythonnlphuggingface-transformerspre-trained-model

How to download a HuggingFace model 'transformers.trainer.Trainer'?

In 1 code., I have uploaded hugging face 'transformers.trainer.Trainer' based model using save_pretrained() function In 2nd code, I want to download this uploaded model and use it to make predictions. I need help in this step - How to download the uploaded model & then make a prediction?

Steps to create model:

from transformers import AutoModelForQuestionAnswering, TrainingArguments, Trainer
model = AutoModelForQuestionAnswering.from_pretrained('xlm-roberta-large)
trainer = Trainer(
model,
args,
train_dataset=tokenized_train_ds,
eval_dataset=tokenized_val_ds,
data_collator=data_collator,
tokenizer=tokenizer,)

#Arguments used above not mentioned here - model, args, tokenized_train_ds, tokenized_val_ds, data_collator, tokenizer
#Below step train the pre-trained model
trainer.train()

I then uploaded this 'trainer' model using the below command:-

trainer.save_model('./trainer_sm')

In a different code, I now want to download this model & use it for making predictions, Can someone advise how to do this? I tried the below command to upload it:-

model_sm=AutoModelForQuestionAnswering.from_pretrained("./trainer_sm")

And used it to make predictions by this line of code:-

model_sm.predict(test_features)
AttributeError: 'XLMRobertaForQuestionAnswering' object has no attribute 'predict'

I also used 'use_auth_token=True' as an argument for from_pretrained, but that also didn't work.

Also, type(trainer) is 'transformers.trainer.Trainer' , while type(model_sm) is transformers.models.xlm_roberta.modeling_xlm_roberta.XLMRobertaForQuestionAnswering

Solution

  • What you have saved is the model which the trainer was going to tune and you should be aware that predicting, training, evaluation and etc, are the utilities of transformers.trainer.Trainer object, not transformers.models.xlm_roberta.modeling_xlm_roberta.XLMRobertaForQuestionAnswering. Based on what was mentioned the easiest way to keep things going is creating another instance of the trainer.

    model_sm=AutoModelForQuestionAnswering.from_pretrained("./trainer_sm")
reloaded_trainer = Trainer(
                    model = model_sm,
                    tokenizer = tokenizer,
                    # other arguments if you have changed the defaults 
                    )
reloaded_trainer.predict(test_dataset)