How to load my own model

[zhihui-shao]

I trained a model using train.py and got the checkpoint folder, how do I load this model for inference?

[syncdoth]

Since the model almost fully supports huggingface transformers API, you can load your model with .from_pretrained(CKPT_DIR).

Example:

from retnet.modeling_retnet import RetNetModelWithLMHead
from transformers import AutoTokenizer

model = RetNetModelWithLMHead.from_pretrained('checkpoints/checkpoint-xxx')
# Trainer API should save tokenizers too in the same directory
tokenizer = AutoTokenizer.from_pretrained('checkpoints/checkpoint-xxx')