|
from models.modeling_flax_wav2vec2 import FlaxWav2Vec2ForCTC |
|
from transformers import AutoConfig, AutoFeatureExtractor, AutoTokenizer |
|
from flax.traverse_util import flatten_dict, unflatten_dict |
|
import collections |
|
|
|
model_id = "sanchit-gandhi/flax-wav2vec2-ctc-cv9-baseline-50k" |
|
|
|
config = AutoConfig.from_pretrained(model_id) |
|
config.use_scan = False |
|
|
|
unrolled_model = FlaxWav2Vec2ForCTC.from_pretrained(model_id, config=config) |
|
model = FlaxWav2Vec2ForCTC.from_pretrained(model_id) |
|
|
|
def scanned_to_unrolled(params): |
|
new_params = collections.defaultdict(dict) |
|
|
|
for key, stacked_weights in flatten_dict(params['wav2vec2']['encoder']['layers']['FlaxWav2Vec2EncoderLayers']).items(): |
|
for layer, weights in enumerate(stacked_weights): |
|
new_key = (str(layer),) + key |
|
new_params[new_key] = weights |
|
|
|
new_params = unflatten_dict({('wav2vec2', 'encoder', 'layers') : unflatten_dict(new_params)}) |
|
new_params = flatten_dict(new_params) |
|
|
|
|
|
for k in flatten_dict(params): |
|
if 'layers' not in k or 'adapter' in k: |
|
new_params[k] = flatten_dict(params)[k] |
|
|
|
return unflatten_dict(new_params) |
|
|
|
unrolled_model.params = scanned_to_unrolled(model.params) |
|
|
|
unrolled_model.save_pretrained("./") |
|
|
|
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) |
|
feature_extractor.save_pretrained("./") |
|
tokenizer = AutoTokenizer.from_pretrained("patrickvonplaten/wav2vec2_ctc_cv9_tokenizer") |
|
tokenizer.save_pretrained("./") |
