pubmedul2_small / export_checkpoint.py
Siddharth63's picture
new additions
98c533d
import argparse
from transformers import T5ForConditionalGeneration, TFT5ForConditionalGeneration
def main(args):
pt_model = T5ForConditionalGeneration.from_pretrained(args.model_dir, from_flax=True)
pt_model.save_pretrained(args.model_dir)
tf_model = TFT5ForConditionalGeneration.from_pretrained(args.model_dir, from_pt=True)
tf_model.save_pretrained(args.model_dir)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--model_dir', type=str, default='.')