初始化项目,由ModelHub XC社区提供模型
Model: Norod78/distilgpt2-base-pretrained-he Source: Original Platform
This commit is contained in:
24
converters/convert2flax.py
Normal file
24
converters/convert2flax.py
Normal file
@@ -0,0 +1,24 @@
|
||||
import argparse
|
||||
import logging
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import os
|
||||
from transformers import AutoConfig, FlaxAutoModelForCausalLM
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO,
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
model_path = "./distilgpt2-base-pretrained-he"
|
||||
save_directory = "./tmp/flax/"
|
||||
|
||||
config_path = os.path.join(model_path, 'config.json')
|
||||
|
||||
# Loading from a PyTorch checkpoint file instead of a TensorFlow model (slower)
|
||||
config = AutoConfig.from_pretrained(config_path)
|
||||
model = FlaxAutoModelForCausalLM.from_pretrained(model_path, from_pt=True, config=config)
|
||||
model.save_pretrained(save_directory)
|
||||
Reference in New Issue
Block a user