def get_framework(test_class):
    """Infer the framework from the test class `test_class`."""

    if "ModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
        return "pt"
    elif "TFModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
        return "tf"
    elif "FlaxModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
        return "flax"
    else:
        return None
Infer the framework from the test class `test_class`.
dpkd/transformers/utils/add_pipeline_model_mapping_to_test.py