import transformers

from transformers import TFWav2Vec2Model
from tests.models.wav2vec2.test_modeling_tf_wav2vec2 import TFWav2Vec2ModelTest

self = TFWav2Vec2ModelTest()
self.setUp()

model_class = TFWav2Vec2Model
allow_missing_keys = False

config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

pt_model_class_name = model_class.__name__[2:]  # Skip the "TF" at the beginning
pt_model_class = getattr(transformers, pt_model_class_name)

tf_model = model_class(config)
pt_model = pt_model_class(config)

tf_inputs_dict = self._prepare_for_class(inputs_dict, model_class)

# Check we can load pt model in tf and vice-versa with model => model functions
try:
    _tf_model = transformers.load_pytorch_model_in_tf2_model(
        tf_model, pt_model, tf_inputs=tf_inputs_dict, allow_missing_keys=allow_missing_keys
    )
except:
    _tf_model = None

try:
    _pt_model = transformers.load_tf2_model_in_pytorch_model(
        pt_model, tf_model, allow_missing_keys=allow_missing_keys
    )
except:
    _pt_model = None

if _tf_model is None:
    print("_tf_model fails")
else:
    print("_tf_model OK")

if _pt_model is None:
    print("_pt_model fails")
else:
    print("_pt_model OK")