saymrwulf/transformers
1
0
Fork 0
mirror of https://github.com/saymrwulf/transformers.git synced 2026-05-14 20:58:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix Funnel Transformer conversion script (#9683)

Browse source
This commit is contained in:
Sylvain Gugger 2021-01-20 09:50:20 -05:00 committed by GitHub
parent 8940c7662d
commit a7b62fece5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
src/transformers/models/funnel/convert_funnel_original_tf_checkpoint_to_pytorch.py
View file

@ -19,18 +19,18 @@ import argparse
import torch
from transformers import FunnelConfig, FunnelForPreTraining, load_tf_weights_in_funnel
from transformers import FunnelBaseModel, FunnelConfig, FunnelModel, load_tf_weights_in_funnel
from transformers.utils import logging
logging.set_verbosity_info()
def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path):
def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path, base_model):
# Initialise PyTorch model
config = FunnelConfig.from_json_file(config_file)
print("Building PyTorch model from configuration: {}".format(str(config)))
model = FunnelForPreTraining(config)
model = FunnelBaseModel(config) if base_model else FunnelModel(config)
# Load weights from tf checkpoint
load_tf_weights_in_funnel(model, config, tf_checkpoint_path)
@ -57,5 +57,10 @@ if __name__ == "__main__":
parser.add_argument(
"--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
)
parser.add_argument(
"--base_model", action="store_true", help="Whether you want just the base model (no decoder) or not."
)
args = parser.parse_args()
convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path)
convert_tf_checkpoint_to_pytorch(
args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.base_model
)