import contextlib
import transformers


@contextlib.contextmanager
def output_hidden_states_context(encoder):
    assert isinstance(
        encoder,
        (
            transformers.BertModel,
            transformers.RobertaModel,
            transformers.AlbertModel,
            transformers.XLMRobertaModel,
        ),
    )
    if hasattr(encoder.encoder, "output_hidden_states"):
        # Transformers < v2
        modified_obj = encoder.encoder
    elif hasattr(encoder.encoder.config, "output_hidden_states"):
        # Transformers >= v3
        modified_obj = encoder.encoder.config
    else:
        raise RuntimeError(f"Failed to convert model {type(encoder)} to output hidden states")

    old_value = modified_obj.output_hidden_states
    modified_obj.output_hidden_states = True
    yield
    modified_obj.output_hidden_states = old_value
