def calculate_perplexity(model: GPT2LMHeadModel, query: torch.Tensor, response: torch.Tensor) -> float:
    """
    **Overview:**
        Calculate the perplexity of the response given a language model, query token ids and response token ids. \
        In essence, the perplexity is the exponential result of cross entropy loss, which can reflect the quality of \
        the generation to some extent.
    **Arguments:**
        - model: The language model to calculate perplexity.
        - query: The token ids for query.
        - response: The token ids for response.
    """
    # Concatenate the query and response.
    total_input = torch.cat([query, response], dim=0)
    # Calculate the logits given the token ids.
    logits = model(total_input, return_dict=True).logits

    # Shift the output logits and input ids to match their dimension.
    # For the i-th shifted logits, it corresponds to the i-th shifted label.
    start = query.shape[0]
    shifted_logits = logits[start:-1, :]
    shifted_labels = total_input[start+1:]

    # Use cross entropy loss to calculate the perplexity.
    loss_fct = torch.nn.CrossEntropyLoss()
    loss = loss_fct(shifted_logits, shifted_labels)
    ppl = torch.exp(loss).item()
    return ppl

    **Overview:**
        Calculate the perplexity of the response given a language model, query token ids and response token ids.         In essence, the perplexity is the exponential result of cross entropy loss, which can reflect the quality of         the generation to some extent.
    **Arguments:**
        - model: The language model to calculate perplexity.
        - query: The token ids for query.
        - response: The token ids for response.
    
chapter8_large/lm_env.py