import torch


def hf_represent(batch_ids, model, layer=-1):
    """
    Encode a batch of sequences of ids using a Hugging Face
    Transformer-based model `model`. The model's `forward` method is
    `output_hidden_states=True`, and we get the hidden states from
    `layer`.


    Parameters
    ----------
    batch_ids : iterable, shape (n_examples, n_tokens)
        Sequences of indices into the model vocabulary.

    model : Hugging Face transformer model

    later : int
        The layer to return. This will get all the hidden states at
        this layer. `layer=0` gives the embedding, and `layer=-1`
        gives the final output states.

    Returns
    -------
    Tensor of shape `(n_examples, n_tokens, n_dimensions)`
       where `n_dimensions` is the dimensionality of the
       Transformer model

    """
    with torch.no_grad():
        reps = model(batch_ids, output_hidden_states=True)
        return reps.hidden_states[layer]