Create convert_scan_to_unrolled.py

convert_scan_to_unrolled.py

@@ -0,0 +1,39 @@
+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)
+    # get the key of a scanned module
+    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)
+
+    # append parameters for non-scanned modules (i.e. all modules that do not contain the key 'layers')
+    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("./")